restore.py 2.51 KB
from utils import *
import re
from capstone import *

# 완전히 다른 놈이면 0.000... 이렇게 나오더라
THRESH_HOLD     = 0.80
API_ELF_PATH    = 'API.ino.elf'
API_DB_PATH     = 'sym_func.db'

class RestoreAPI:
    def __init__(self, binPath, functionOffsets):
        funcs = readFunctionInfoFromElf(API_ELF_PATH)
        funcs = self.__refineByDB(funcs)
        apiList = {}
        angrFuncList = readFunctionInfoFromBin(binPath, functionOffsets)
        for func in funcs:
            found = self.__findMostSimilarFunction(func, angrFuncList)
            if found['prob'] >= THRESH_HOLD:
                apiList[found['offset']] = {
                    'type': func['type'],
                    'name': func['name']
                }
                print(found['prob'], hex(found['offset']), func['name'])
        print('[*] found api count:', len(apiList))
        self.apiList = apiList

    def __findMostSimilarFunction(self, src, angrFuncList):
        BoB = {'offset': -1, 'prob': 0.0, 'ops': []}
        for angrFunc in angrFuncList:
            prob = compareFunction(src['ops'], angrFunc['ops'])
            if prob > BoB['prob']:
                BoB['offset'] = angrFunc['offset']
                BoB['prob'] = prob
                BoB['ops'] = angrFunc['ops']
        return BoB

    def __refineByDB(self, funcs):
        dbMap = self.__readDB(API_DB_PATH)
        ret = []
        for func in funcs:
            info = self.__findInfoByName(dbMap, func['name'])
            if info == None:
                continue
            name = func['name']
            name = name[name.find('.')+1:]
            func['name'] = name
            print(f'db: {name}')
            func['type'] = info['type']
            ret.append(func)
        return ret

    def __readDB(self, dbPath):
        ret = []
        currentType = None
        with open(dbPath, 'rt') as f:
            for s in f.readlines():
                s = s.strip(' \n')
                if s == '':
                    continue
                if s[0] == '#':
                    continue
                t = re.findall('^!\((.+)\)$', s)
                if len(t) == 1:
                    currentType = t[0]
                else:
                    # it must be substring of function name!
                    ret.append({'name': s, 'type': currentType})
        return ret

    def __findInfoByName(self, dbMap, name):
        for v in dbMap:
            if v['name'] in name:
                return v
        return None

    def restore(self):
        return self.apiList