Manticore script to check whether the invariant sum_of_balances == total_supply is always satisfied in a ERC20 contract

balanceof_invariant_check.py

from manticore.ethereum import *

import itertools

import binascii

import time





start_time = time.time()



################ Script #######################



m = ManticoreEVM()

sol_file = open("PFAToken.sol")

source_code = sol_file.read()





print("[+] Creating user accounts")

user_account = m.create_account(balance=1000000000000000, name='user_account')



# INFO, add more accounts if this is needed to trigger some edge-cases in the

# contract or could help in the analysis (execution may take more time)

#user_account2 = m.create_account(balance=1000000000000000, name='user_account2')



print("[+] Creating contract account")

contract_account = m.solidity_create_contract(source_code, owner=user_account,

                contract_name="PFAToken", name='contract_account')



print("[+] First transaction")

m.transaction(caller=m.make_symbolic_value(name="CALLER1"),

              address=contract_account,

              data=m.make_symbolic_buffer(320,name="DATA1"),

              value=m.make_symbolic_value(name="VALUE1"),

              gas=10000000000)



# WARNING, before uncommenting this be sure to know what you whish for...

#print("[+] Second transaction")

#m.transaction(caller=m.make_symbolic_value(name="CALLER2"),

#              address=contract_account,

#              data=m.make_symbolic_buffer(320,name="DATA2"),

#              value=m.make_symbolic_value(name="VALUE2"),

#              gas=10000000000)



print(f"[+] Generated {m.count_running_states()} running states")



# Get total supply

print("[+] Get total supply")

contract_account.totalSupply()

total_supplies = []                         # List of BitVec containing total supplies

for state in m.running_states:

    r = state.platform.transactions[-1].return_data

    r = r.read_BE(0,32)                     # Converts ArrayProxy to BitVec

    total_supplies.append(r)

    # DEBUG, here we just print a possible totalSupply value for every state

    r = state.solve_one(r)                              

    print(r)



# Get balances

print("[+] Retrieve all balances")

all_balances = []

for acc_name in itertools.chain(m.normal_accounts, m.contract_accounts):

    acc_addr = int(m.get_account(acc_name))

    contract_account.balanceOf(acc_addr)

    balances_acc = []

    for state in m.running_states:

        r = state.platform.transactions[-1].return_data

        r = r.read_BE(0,32)

        balances_acc.append(r)

        # DEBUG, print 10 possible balanceOf values for this account and state

        r = state.solve_n(r,10)                              

        print(r)

    all_balances.append(balances_acc)



# Compute sum of balances

print("[+] Summing up balances")

sum_balances = [ sum(x) for x in zip(*all_balances)]



# Verify invariant

print("[+] Verifying Invariant")

tmp = zip(m.running_states, total_supplies, sum_balances)

for state,totalSupply,sum_balanceOf in tmp:

    if (state.can_be_true(totalSupply != sum_balanceOf)):

        print("~~~~~~~> Invariant violation found")

        m.generate_testcase(state, 'TotalSupply != SumBalances')

        # DEBUG, show up to 10 possible sums of balances

        print("[+] Possible sums of balances")

        r = state.solve_n(sum_balanceOf,10)                              

        print(r)

        

        print("## Unconstrained input ##")

        input_data = state.platform.transactions[1].data

        r = state.solve_one(input_data)

        r = binascii.hexlify(r)

        print(r)



        print("## Constrained input ##")

        state.constrain(totalSupply != sum_balanceOf)

        r = state.solve_one(input_data)

        h = binascii.hexlify(r)

        print(h)



        # Generate and dump testcase

        m.generate_testcase(state, 'CONSTRAINED VIOLATION')



##################### END #########################

print("--- %s seconds ---" % (time.time() - start_time))