pe_analyzer.py

from pwn import u32
import pefile
import os
import subprocess
from analyzer.config import *
from collections import namedtuple
from analyzer import packer1_unpacker_x86
from analyzer import packer1_unpacker_x64

class pe_file():
    def __init__(self,name=None,pe=None,stirng=None):
        self._name = None
        self._pe = None
        self._strings = None

    @property
    def name(self):
        return self._name

    @name.setter
    def name(self,n):
        self._name = n

    @property
    def pe(self):
        return self._pe

    @pe.setter
    def pe(self,p):
        self._pe = p

    @property
    def strings(self):
        return self._strings

    @strings.setter
    def strings(self,s):
        self._strings = s

class pe_analyzer():
    def __init__(self):
        self.files = {}
        self.result = {}

    @staticmethod
    def repair_one_file(path):
        with open(path, "rb") as f:
            data = f.read()
        data = bytearray(data)
        data[0:2] = b"MZ"
        pe_addr = u32(data[0x3c:0x3c+4])
        data[pe_addr:pe_addr+2] = b"PE"
        return data

    @staticmethod
    def packer_scanner(data):
        needCheck = False
        for cha in PACKERS:
            if cha.encode() in data:
                needCheck = True
        if not(needCheck):
            return 0,None

        pe = pefile.PE(data=data)
        for section in pe.sections:
            for cha in PACKERS:
                if cha.encode() in section.Name:
                    # print("PACKER: %d" % PACKERS[cha])
                    return PACKERS[cha],pe
        return 0,None

    @staticmethod
    def unpack(name,pe,packer,rawdata):
        out_path = os.path.join(TMP,name)
        if packer == PACKERS["UPX"]:
            the_last_upx_section_addr = None
            for section in pe.sections:
                if "UPX".encode() in section.Name:
                    the_last_upx_section_addr = section.VirtualAddress
            if the_last_upx_section_addr == None:
                print("section not found: %s" % path)
                assert False
            pe.NT_HEADERS.OPTIONAL_HEADER.DATA_DIRECTORY[1].VirtualAddress = the_last_upx_section_addr
            pe.NT_HEADERS.OPTIONAL_HEADER.DATA_DIRECTORY[1].Size = 0xb4
        elif packer == PACKERS["MPRESS"]:
            pass
        else:
            # elif packer == PACKERS["\x20\x20\x20\x00\x20\x20\x20\x20"]:
            print("PACKER1: " + name)
            # with open(out_path+".bak", "wb") as f:
            #     f.write(rawdata)
            if pe.FILE_HEADER.Machine == 0x8664:
                unpack_data = packer1_unpacker_x64.emu_unpacker(rawdata)
            else:
                unpack_data = packer1_unpacker_x86.emu_unpacker(rawdata)
            # assert unpack_data != None and "unpack_data is None!!!!"
            # print(unpack_data)
            with open(out_path, "wb") as f:
                f.write(unpack_data)
            return out_path
        # else:
            # pass
        pe.write(out_path)
        os.system("%s %s -o %s" % (RET, out_path, out_path))
        return out_path

    def parse_one_file(self,path):
        name = path.split("/")[-1]
        data = self.repair_one_file(path)
        try:
            packer,pe = self.packer_scanner(data)
        except:
            packer,pe = None,None
        if packer:
            tmp = self.unpack(name,pe,packer, data)
            with open(tmp,"rb") as f:
                data = f.read()
        else:
            tmp = os.path.join(TMP,name)
            with open(tmp, "wb") as f:
                f.write(data)
        res_aes = self.check_aes(data)
        if res_aes:
            subprocess.check_output(['rm',tmp])
            return (name,1)
        else:
            s = subprocess.check_output(['strings','-el',tmp])
            s += subprocess.check_output(['strings',tmp])
            subprocess.check_output(['rm',tmp])
            res_str = self.check_strings(s)
            if res_str: return (name,1)
            else: return (name,0)

    @staticmethod
    def check_aes(data):
        if aes_const in data:
            return 1
        return 0

    @staticmethod
    def check_strings(strings):
        for pattern in list(patterns.keys()):
            l = pattern.findall(strings.decode())
            if l:
                return 1
        return 0