import re
import os
import click
import r2pipe
import hashlib
from pathlib import Path
import _pickle as cPickle

def sha3(data):
    return hashlib.sha3_256(data.encode()).hexdigest()

def validEXE(filename):
    magics = [bytes.fromhex('7f454c46')]
    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:
        # add label
        if labels.get(op.get('offset')) is not None:
            output += f'LABEL{labels[op["offset"]]}:\n'
        # add instruction
        if labels.get(op.get('jump')) is not None:
            output += f' {op["type"]} LABEL{labels[op["jump"]]}\n'
        else:
            output += f' {normalize(op["opcode"])}\n'

    return output

def bin2asm(filename, opath, minlen):
    #
    # Create directory where results will be written to.
    #
    results_dir = os.path.join(opath, os.path.basename(filename))
    if not os.access(results_dir, os.F_OK):
        os.makedirs(results_dir)

    r = r2pipe.open(str(filename))
    r.cmd('aaaa')

    count = 0

    fp = open("%s/fv.pcl" % (results_dir), 'wb')
    for fn in r.cmdj('aflj'):
        r.cmd(f's {fn["offset"]}')
        asm = fn2asm(r.cmdj('pdfj'), minlen)
        if asm:
            fv = [
                fn["name"],
                asm
                 ]

            cPickle.dump(fv, fp)
            count += 1

    fp.close()
    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=1, 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()

[((2135, 2150), 'click.command', 'click.command', ([], {}), '()\n', (2148, 2150), False, 'import click\n'), ((2152, 2240), 'click.option', 'click.option', (['"""-i"""', '"""--input"""', '"""ipath"""'], {'help': '"""input directory / file"""', 'required': '(True)'}), "('-i', '--input', 'ipath', help='input directory / file',\n    required=True)\n", (2164, 2240), False, 'import click\n'), ((2238, 2317), 'click.option', 'click.option', (['"""-o"""', '"""--output"""', '"""opath"""'], {'default': '"""asm"""', 'help': '"""output directory"""'}), "('-o', '--output', 'opath', default='asm', help='output directory')\n", (2250, 2317), False, 'import click\n'), ((2319, 2446), 'click.option', 'click.option', (['"""-l"""', '"""--len"""', '"""minlen"""'], {'default': '(1)', 'help': '"""ignore assembly code with instructions amount smaller than minlen"""'}), "('-l', '--len', 'minlen', default=1, help=\n    'ignore assembly code with instructions amount smaller than minlen')\n", (2331, 2446), False, 'import click\n'), ((425, 463), 're.sub', 're.sub', (['"""0x[0-9a-f]+"""', '"""CONST"""', 'opcode'], {}), "('0x[0-9a-f]+', 'CONST', opcode)\n", (431, 463), False, 'import re\n'), ((478, 514), 're.sub', 're.sub', (['"""\\\\*[0-9]"""', '"""*CONST"""', 'opcode'], {}), "('\\\\*[0-9]', '*CONST', opcode)\n", (484, 514), False, 'import re\n'), ((528, 562), 're.sub', 're.sub', (['""" [0-9]"""', '""" CONST"""', 'opcode'], {}), "(' [0-9]', ' CONST', opcode)\n", (534, 562), False, 'import re\n'), ((2555, 2566), 'pathlib.Path', 'Path', (['ipath'], {}), '(ipath)\n', (2559, 2566), False, 'from pathlib import Path\n'), ((2579, 2590), 'pathlib.Path', 'Path', (['opath'], {}), '(opath)\n', (2583, 2590), False, 'from pathlib import Path\n'), ((2731, 2751), 'os.path.isdir', 'os.path.isdir', (['ipath'], {}), '(ipath)\n', (2744, 2751), False, 'import os\n'), ((1577, 1603), 'os.path.basename', 'os.path.basename', (['filename'], {}), '(filename)\n', (1593, 1603), False, 'import os\n'), ((1616, 1647), 'os.access', 'os.access', (['results_dir', 'os.F_OK'], {}), '(results_dir, os.F_OK)\n', (1625, 1647), False, 'import os\n'), ((1657, 1681), 'os.makedirs', 'os.makedirs', (['results_dir'], {}), '(results_dir)\n', (1668, 1681), False, 'import os\n'), ((2633, 2654), 'os.path.exists', 'os.path.exists', (['opath'], {}), '(opath)\n', (2647, 2654), False, 'import os\n'), ((2664, 2679), 'os.mkdir', 'os.mkdir', (['opath'], {}), '(opath)\n', (2672, 2679), False, 'import os\n'), ((2770, 2787), 'os.listdir', 'os.listdir', (['ipath'], {}), '(ipath)\n', (2780, 2787), False, 'import os\n'), ((2976, 2997), 'os.path.exists', 'os.path.exists', (['ipath'], {}), '(ipath)\n', (2990, 2997), False, 'import os\n'), ((2026, 2046), '_pickle.dump', 'cPickle.dump', (['fv', 'fp'], {}), '(fv, fp)\n', (2038, 2046), True, 'import _pickle as cPickle\n'), ((2808, 2833), 'os.path.islink', 'os.path.islink', (['(ipath / f)'], {}), '(ipath / f)\n', (2822, 2833), False, 'import os\n'), ((2842, 2866), 'os.path.isdir', 'os.path.isdir', (['(ipath / f)'], {}), '(ipath / f)\n', (2855, 2866), False, 'import os\n')]