mono-debundle

mono_debundle.py

#!/usr/bin/env python
# -*- encoding: utf-8 -*-
# Author: TheCjw<thecjw@qq.com>
# Created on 14:58 2015/5/24

__author__ = "TheCjw"

import argparse
import os
import struct
import zlib

try:
    import pefile
    from tabulate import tabulate
except ImportError as e:
    print e
    exit(0)


class MonoDebundle(object):
    def __init__(self, input_path, output_dir=""):
        self.input_path = input_path
        self.output_dir = output_dir

        self.pe_file = None
        self.pe_mapped_image = None

        self.image_start = 0
        self.image_end = 0

        self.assemblies = []

    def _is_valid(self):
        return self.pe_file and \
               self.pe_mapped_image and \
               self.image_start and \
               self.image_end

    def _is_zlib_header(self, address):
        if self._is_valid() == False:
            return False

        if address < self.image_start or address > self.image_end:
            return False

        if address % 4 != 0:
            return False

        header = self.pe_mapped_image[address - self.image_start:address - self.image_start + 2]
        if len(header) != 2:
            return False

        return header[0] == "\x78" and header[1] == "\x9c"

    def _get_ansi_string_from_va(self, address):
        if self._is_valid() == False:
            return ""

        if address < self.image_start or address > self.image_end:
            return False

        rva = address - self.image_start
        chars = []

        walker = rva
        while True:
            # may crash when meets the end.
            if self.pe_mapped_image[walker:walker + 1] == "\x00":
                break
            chars.append(self.pe_mapped_image[walker:walker + 1])

            walker += 1

        return "".join(chars)

    def get_assemblies_list(self):

        try:
            self.pe_file = pefile.PE(self.input_path, fast_load=True)
            self.pe_mapped_image = self.pe_file.get_memory_mapped_image()
            self.image_start = self.pe_file.OPTIONAL_HEADER.ImageBase
            self.image_end = self.pe_file.OPTIONAL_HEADER.ImageBase + self.pe_file.OPTIONAL_HEADER.SizeOfImage

            bundled_assemblies_descriptor_rva = 0

            for section in self.pe_file.sections:

                # 忽略.text节
                if section.IMAGE_SCN_MEM_EXECUTE:
                    continue

                if section.IMAGE_SCN_CNT_INITIALIZED_DATA:
                    section_data = section.get_data()
                    for i in xrange(0, len(section_data), 4):

                        address = struct.unpack("I", section_data[i:i + 4])[0]
                        if self._is_zlib_header(address) == False:
                            continue

                        next_zlib_address = struct.unpack("I", section_data[i + 0x10:i + 0x10 + 4])[0]
                        if self._is_zlib_header(next_zlib_address) == False:
                            continue

                        bundled_assemblies_descriptor_rva = section.VirtualAddress + i - 4
                        break

                if bundled_assemblies_descriptor_rva != 0:
                    break

            if bundled_assemblies_descriptor_rva == 0:
                return 0

            walker = bundled_assemblies_descriptor_rva
            while True:
                des = struct.unpack_from("IIII", self.pe_mapped_image, walker)
                name_va = des[0]
                data_va = des[1]
                size = des[2]
                compressed_size = des[3]

                if self._is_zlib_header(data_va) == False:
                    break

                # Convert VA to file Offset.
                self.assemblies.append(
                    [self._get_ansi_string_from_va(name_va),
                     self.pe_file.get_offset_from_rva(data_va - self.image_start),
                     size,
                     compressed_size])

                walker += 0x10

        except Exception as e:
            print e

        return len(self.assemblies)

    def extract_assemblies(self):
        self.get_assemblies_list()
        if len(self.assemblies) == 0:
            print "No assembly found."
            return

        if os.path.isabs(self.output_dir) == False:
            self.output_dir = os.path.abspath(self.output_dir)

        if os.path.exists(self.output_dir) == False:
            os.mkdir(self.output_dir)

        # print self.output_dir
        for info in self.assemblies:
            rva = self.pe_file.get_rva_from_offset(info[1])
            compressed_data = self.pe_mapped_image[rva:rva + info[3]]
            decompressed_data = zlib.decompress(compressed_data)

            with open(os.path.join(self.output_dir, info[0]), "wb") as file:
                file.write(decompressed_data)
                file.flush()

    def replace_assembly(self, file_to_replace):
        if os.path.exists(file_to_replace) == False:
            print file_to_replace, "is not exists."
            return

        self.get_assemblies_list()
        if len(self.assemblies) == 0:
            print "No assembly found."
            return

        base_file_name = os.path.basename(file_to_replace)
        for info in self.assemblies:
            if base_file_name.upper() == info[0].upper():

                file_data = open(file_to_replace, "rb").read()
                compressed_data = zlib.compress(file_data, 7)
                if len(compressed_data) > info[3]:
                    print file_to_replace, "is too large."
                    break

                compressed_data += "\x00" * (info[3] - len(compressed_data))
                self.pe_file.set_bytes_at_offset(info[1], compressed_data)

                self.pe_file.write(os.path.join(self.output_dir, "patched.exe"))


def parse_args():
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(help="commands")

    task_list = subparsers.add_parser("search",
                                      help="search bundled files")
    task_list.set_defaults(command="search")
    task_list.add_argument("input_path",
                           action="store",
                           help="Input file path")

    task_extract = subparsers.add_parser("extract",
                                         help="Extract bundle files")
    task_extract.set_defaults(command="extract")
    task_extract.add_argument("input_path",
                              action="store",
                              help="Input file path")
    task_extract.add_argument("output_dir",
                              action="store",
                              help="Output directory")

    task_replace = subparsers.add_parser("replace",
                                         help="Replace specific file")
    task_replace.set_defaults(command="replace")
    task_replace.add_argument("input_path",
                              action="store",
                              help="Input file path")
    task_replace.add_argument("output_dir",
                              action="store",
                              help="Output directory")
    task_replace.add_argument("file_to_replace",
                              action="store",
                              help="File to replace")
    args = parser.parse_args()

    return args


def main():
    args = parse_args()
    if args.command == "search":
        debundle = MonoDebundle(args.input_path)
        debundle.get_assemblies_list()
        table_out = [["Assembly Name", "Offset", "Size", "Compressed Size"]]

        for info in debundle.assemblies:
            table_out.append([info[0], hex(info[1]), hex(info[2]), hex(info[3])])

        print tabulate(table_out)

    elif args.command == "extract":
        MonoDebundle(args.input_path, args.output_dir).extract_assemblies()

    elif args.command == "replace":
        MonoDebundle(args.input_path, args.output_dir).replace_assembly(args.file_to_replace)


if __name__ == "__main__":
    main()