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', False, 'import click\n'), (90, 'click.option', 'click.option', False, 'import click\n'), (91, 'click.option', 'click.option', False, 'import click\n'), (20, 're.sub', 're.sub', False, 'import re\n'), (21, 're.sub', 're.sub', False, 'import re\n'), (22, 're.sub', 're.sub', False, 'import re\n'), (96, 'pathlib.Path', 'Path', False, 'from pathlib import Path\n'), (97, 'pathlib.Path', 'Path', False, 'from pathlib import Path\n'), (106, 'os.path.isdir', 'os.path.isdir', False, 'import os\n'), (61, 'os.path.basename', 'os.path.basename', False, 'import os\n'), (62, 'os.access', 'os.access', False, 'import os\n'), (63, 'os.makedirs', 'os.makedirs', False, 'import os\n'), (100, 'os.path.exists', 'os.path.exists', False, 'import os\n'), (101, 'os.mkdir', 'os.mkdir', False, 'import os\n'), (107, 'os.listdir', 'os.listdir', False, 'import os\n'), (112, 'os.path.exists', 'os.path.exists', False, 'import os\n'), (80, '_pickle.dump', 'cPickle.dump', True, 'import _pickle as cPickle\n'), (108, 'os.path.islink', 'os.path.islink', False, 'import os\n'), (108, 'os.path.isdir', 'os.path.isdir', False, 'import os\n')]