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()

[(88, 'click.command', 'click.command', ([], {}), False, 'import click\n'), (89, 'click.option', 'click.option', (['"""-i"""', '"""--input"""', '"""ipath"""'], {'help': '"""input directory / file"""', 'required': '(True)'}), False, 'import click\n'), (90, 'click.option', 'click.option', (['"""-o"""', '"""--output"""', '"""opath"""'], {'default': '"""asm"""', 'help': '"""output directory"""'}), False, 'import click\n'), (91, 'click.option', 'click.option', (['"""-l"""', '"""--len"""', '"""minlen"""'], {'default': '(1)', 'help': '"""ignore assembly code with instructions amount smaller than minlen"""'}), False, 'import click\n'), (20, 're.sub', 're.sub', (['"""0x[0-9a-f]+"""', '"""CONST"""', 'opcode'], {}), False, 'import re\n'), (21, 're.sub', 're.sub', (['"""\\\\*[0-9]"""', '"""*CONST"""', 'opcode'], {}), False, 'import re\n'), (22, 're.sub', 're.sub', (['""" [0-9]"""', '""" CONST"""', 'opcode'], {}), False, 'import re\n'), (96, 'pathlib.Path', 'Path', (['ipath'], {}), False, 'from pathlib import Path\n'), (97, 'pathlib.Path', 'Path', (['opath'], {}), False, 'from pathlib import Path\n'), (106, 'os.path.isdir', 'os.path.isdir', (['ipath'], {}), False, 'import os\n'), (61, 'os.path.basename', 'os.path.basename', (['filename'], {}), False, 'import os\n'), (62, 'os.access', 'os.access', (['results_dir', 'os.F_OK'], {}), False, 'import os\n'), (63, 'os.makedirs', 'os.makedirs', (['results_dir'], {}), False, 'import os\n'), (100, 'os.path.exists', 'os.path.exists', (['opath'], {}), False, 'import os\n'), (101, 'os.mkdir', 'os.mkdir', (['opath'], {}), False, 'import os\n'), (107, 'os.listdir', 'os.listdir', (['ipath'], {}), False, 'import os\n'), (112, 'os.path.exists', 'os.path.exists', (['ipath'], {}), False, 'import os\n'), (80, '_pickle.dump', 'cPickle.dump', (['fv', 'fp'], {}), True, 'import _pickle as cPickle\n'), (108, 'os.path.islink', 'os.path.islink', (['(ipath / f)'], {}), False, 'import os\n'), (108, 'os.path.isdir', 'os.path.isdir', (['(ipath / f)'], {}), False, 'import os\n')]