142 lines
3.5 KiB
Python
142 lines
3.5 KiB
Python
|
#!/usr/bin/env python3
|
|||
|
|
import re
|
|||
|
|
import os
|
|||
|
|
import click
|
|||
|
|
import r2pipe
|
|||
|
|
import hashlib
|
|||
|
|
from pathlib import Path
|
|||
|
|
import csv
|
|||
|
|
|
|||
|
|
def sha3(data):
|
|||
|
|
return hashlib.sha3_256(data.encode()).hexdigest()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def validEXE(filename):
|
|||
|
|
magics = [bytes.fromhex('4d5a9000')]
|
|||
|
|
with open(filename, 'rb') as f:
|
|||
|
|
header = f.read(4)
|
|||
|
|
return header in magics
|
|||
|
|
|
|||
|
|
|
|||
|
|
def normalize(opcode):
|
|||
|
|
opcode = opcode.replace(' - ', ' + ')
|
|||
|
|
opcode = re.sub(r'0x[0-9a-f]+', 'CONST', opcode)
|
|||
|
|
opcode = re.sub(r'\*[0-9]', '*CONST', opcode)
|
|||
|
|
opcode = re.sub(r' [0-9]', ' CONST', opcode)
|
|||
|
|
return opcode
|
|||
|
|
|
|||
|
|
|
|||
|
|
def fn2asm(pdf, minlen):
|
|||
|
|
# check
|
|||
|
|
if pdf is None:
|
|||
|
|
return
|
|||
|
|
if len(pdf['ops']) < minlen:
|
|||
|
|
return
|
|||
|
|
if 'invalid' in [op['type'] for op in pdf['ops']]:
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
ops = pdf['ops']
|
|||
|
|
|
|||
|
|
# set label
|
|||
|
|
labels, scope = {}, [op['offset'] for op in ops]
|
|||
|
|
assert (None not in scope)
|
|||
|
|
for i, op in enumerate(ops):
|
|||
|
|
if op.get('jump') in scope:
|
|||
|
|
labels.setdefault(op.get('jump'), i)
|
|||
|
|
|
|||
|
|
# dump output
|
|||
|
|
output = ''
|
|||
|
|
|
|||
|
|
|
|||
|
|
for op in ops:
|
|||
|
|
output+=normalize(op["bytes"])
|
|||
|
|
|
|||
|
|
# output += f'{hex(op["offset"])+":"+normalize(op["opcode"])+":"+normalize(op["bytes"])}\n'
|
|||
|
|
# exit()
|
|||
|
|
|
|||
|
|
return output
|
|||
|
|
|
|||
|
|
|
|||
|
|
def bin2asm(filename, opath, minlen):
|
|||
|
|
# check
|
|||
|
|
# if not validEXE(filename):
|
|||
|
|
# print("fail Exe")
|
|||
|
|
# return 0
|
|||
|
|
|
|||
|
|
# print("666")
|
|||
|
|
r = r2pipe.open(str(filename))
|
|||
|
|
r.cmd('aaaa')
|
|||
|
|
|
|||
|
|
# print(str(filename))
|
|||
|
|
# exit()
|
|||
|
|
count = 0
|
|||
|
|
# header = ['func_name', 'bytes']
|
|||
|
|
csv_data = []
|
|||
|
|
|
|||
|
|
saved_asm_list = []
|
|||
|
|
with open('./asm_func/asm_hex/func_bytes.csv',encoding="utf8") as f:
|
|||
|
|
csv_reader = csv.reader(f)
|
|||
|
|
for line in csv_reader:
|
|||
|
|
# print("good")
|
|||
|
|
# print(line[1])
|
|||
|
|
asm=line[1]
|
|||
|
|
|
|||
|
|
if asm != "bytes":
|
|||
|
|
saved_asm_list.append(line[1])
|
|||
|
|
|
|||
|
|
|
|||
|
|
for fn in r.cmdj('aflj'):
|
|||
|
|
r.cmd(f's {fn["offset"]}')
|
|||
|
|
asm = fn2asm(r.cmdj('pdfj'), minlen)
|
|||
|
|
#如果该asm未重复且不为空，则记录
|
|||
|
|
if asm not in saved_asm_list and asm:
|
|||
|
|
csv_data.append([fn["name"], asm])
|
|||
|
|
count += 1
|
|||
|
|
|
|||
|
|
with open('./asm_func/asm_hex/func_bytes.csv','a+',encoding='utf-8',newline='') as fp:
|
|||
|
|
writer =csv.writer(fp)
|
|||
|
|
# writer.writerow(header)
|
|||
|
|
writer.writerows(csv_data)
|
|||
|
|
|
|||
|
|
print(f'[+] {filename}')
|
|||
|
|
|
|||
|
|
return count
|
|||
|
|
|
|||
|
|
|
|||
|
|
@click.command()
|
|||
|
|
@click.option('-i', '--input', 'ipath', help='input directory / file', required=True)
|
|||
|
|
@click.option('-o', '--output', 'opath', default='asm', help='output directory')
|
|||
|
|
@click.option('-l', '--len', 'minlen', default=10,
|
|||
|
|
help='ignore assembly code with instructions amount smaller than minlen')
|
|||
|
|
def cli(ipath, opath, minlen):
|
|||
|
|
'''
|
|||
|
|
Extract assembly functions from binary executable
|
|||
|
|
'''
|
|||
|
|
ipath = Path(ipath)
|
|||
|
|
opath = Path(opath)
|
|||
|
|
|
|||
|
|
# create output directory
|
|||
|
|
if not os.path.exists(opath):
|
|||
|
|
os.mkdir(opath)
|
|||
|
|
|
|||
|
|
fcount, bcount = 0, 0
|
|||
|
|
|
|||
|
|
# directory
|
|||
|
|
if os.path.isdir(ipath):
|
|||
|
|
for f in os.listdir(ipath):
|
|||
|
|
if not os.path.islink(ipath / f) and not os.path.isdir(ipath / f):
|
|||
|
|
|
|||
|
|
fcount += bin2asm(ipath / f, opath, minlen)
|
|||
|
|
bcount += 1
|
|||
|
|
# file
|
|||
|
|
elif os.path.exists(ipath):
|
|||
|
|
fcount += bin2asm(ipath, opath, minlen)
|
|||
|
|
bcount += 1
|
|||
|
|
else:
|
|||
|
|
print(f'[Error] No such file or directory: {ipath}')
|
|||
|
|
print(f'[+] Total scan binary: {bcount} => Total generated assembly functions: {fcount}')
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == '__main__':
|
|||
|
|
cli()
|