diff --git a/.gitignore b/.gitignore index 40ea1a2..83df4b6 100644 --- a/.gitignore +++ b/.gitignore @@ -19,3 +19,6 @@ cache # k8s .helm + +# env +.envrc diff --git a/mev_inspect/classifiers/specs/aave.py b/mev_inspect/classifiers/specs/aave.py index c12e785..4def777 100644 --- a/mev_inspect/classifiers/specs/aave.py +++ b/mev_inspect/classifiers/specs/aave.py @@ -1,14 +1,17 @@ from mev_inspect.schemas.classified_traces import ( - Classification, - ClassifierSpec, Protocol, ) +from mev_inspect.schemas.classifiers import ( + ClassifierSpec, + LiquidationClassifier, +) + AAVE_SPEC = ClassifierSpec( abi_name="AaveLendingPool", protocol=Protocol.aave, - classifications={ - "liquidationCall(address,address,address,uint256,bool)": Classification.liquidate, + classifiers={ + "liquidationCall(address,address,address,uint256,bool)": LiquidationClassifier, }, ) diff --git a/mev_inspect/classifiers/specs/balancer.py b/mev_inspect/classifiers/specs/balancer.py index 835d1bf..9b05dfd 100644 --- a/mev_inspect/classifiers/specs/balancer.py +++ b/mev_inspect/classifiers/specs/balancer.py @@ -1,16 +1,19 @@ from mev_inspect.schemas.classified_traces import ( - Classification, - ClassifierSpec, Protocol, ) +from mev_inspect.schemas.classifiers import ( + ClassifierSpec, + SwapClassifier, +) + BALANCER_V1_SPECS = [ ClassifierSpec( abi_name="BPool", protocol=Protocol.balancer_v1, - classifications={ - "swapExactAmountIn(address,uint256,address,uint256,uint256)": Classification.swap, - "swapExactAmountOut(address,uint256,address,uint256,uint256)": Classification.swap, + classifiers={ + "swapExactAmountIn(address,uint256,address,uint256,uint256)": SwapClassifier, + "swapExactAmountOut(address,uint256,address,uint256,uint256)": SwapClassifier, }, ), ] diff --git a/mev_inspect/classifiers/specs/curve.py b/mev_inspect/classifiers/specs/curve.py index 79610f9..a631b8c 100644 --- a/mev_inspect/classifiers/specs/curve.py +++ b/mev_inspect/classifiers/specs/curve.py @@ -1,8 +1,11 @@ from mev_inspect.schemas.classified_traces import ( - ClassifierSpec, Protocol, ) +from mev_inspect.schemas.classifiers import ( + ClassifierSpec, +) + """ Deployment addresses found here https://curve.readthedocs.io/ref-addresses.html diff --git a/mev_inspect/classifiers/specs/erc20.py b/mev_inspect/classifiers/specs/erc20.py index 02cd3ae..aee3179 100644 --- a/mev_inspect/classifiers/specs/erc20.py +++ b/mev_inspect/classifiers/specs/erc20.py @@ -1,15 +1,30 @@ -from mev_inspect.schemas.classified_traces import ( - Classification, +from mev_inspect.schemas.classified_traces import DecodedCallTrace +from mev_inspect.schemas.classifiers import ( ClassifierSpec, + TransferClassifier, ) +from mev_inspect.schemas.transfers import ERC20Transfer + + +class ERC20TransferClassifier(TransferClassifier): + @staticmethod + def get_transfer(trace: DecodedCallTrace) -> ERC20Transfer: + return ERC20Transfer( + block_number=trace.block_number, + transaction_hash=trace.transaction_hash, + trace_address=trace.trace_address, + amount=trace.inputs["amount"], + to_address=trace.inputs["recipient"], + from_address=trace.inputs.get("sender", trace.from_address), + token_address=trace.to_address, + ) ERC20_SPEC = ClassifierSpec( abi_name="ERC20", - classifications={ - "transferFrom(address,address,uint256)": Classification.transfer, - "transfer(address,uint256)": Classification.transfer, - "burn(address)": Classification.burn, + classifiers={ + "transferFrom(address,address,uint256)": ERC20TransferClassifier, + "transfer(address,uint256)": ERC20TransferClassifier, }, ) diff --git a/mev_inspect/classifiers/specs/uniswap.py b/mev_inspect/classifiers/specs/uniswap.py index e6c9f5a..ad6d146 100644 --- a/mev_inspect/classifiers/specs/uniswap.py +++ b/mev_inspect/classifiers/specs/uniswap.py @@ -1,8 +1,10 @@ from mev_inspect.schemas.classified_traces import ( - Classification, - ClassifierSpec, Protocol, ) +from mev_inspect.schemas.classifiers import ( + ClassifierSpec, + SwapClassifier, +) UNISWAP_V3_CONTRACT_SPECS = [ @@ -66,8 +68,8 @@ UNISWAP_V3_CONTRACT_SPECS = [ UNISWAP_V3_GENERAL_SPECS = [ ClassifierSpec( abi_name="UniswapV3Pool", - classifications={ - "swap(address,bool,int256,uint160,bytes)": Classification.swap, + classifiers={ + "swap(address,bool,int256,uint160,bytes)": SwapClassifier, }, ), ClassifierSpec( @@ -97,8 +99,8 @@ UNISWAPPY_V2_CONTRACT_SPECS = [ UNISWAPPY_V2_PAIR_SPEC = ClassifierSpec( abi_name="UniswapV2Pair", - classifications={ - "swap(uint256,uint256,address,bytes)": Classification.swap, + classifiers={ + "swap(uint256,uint256,address,bytes)": SwapClassifier, }, ) diff --git a/mev_inspect/classifiers/specs/weth.py b/mev_inspect/classifiers/specs/weth.py index df85790..e84f2f3 100644 --- a/mev_inspect/classifiers/specs/weth.py +++ b/mev_inspect/classifiers/specs/weth.py @@ -1,16 +1,35 @@ from mev_inspect.schemas.classified_traces import ( - Classification, - ClassifierSpec, Protocol, ) +from mev_inspect.schemas.classifiers import ( + ClassifierSpec, + DecodedCallTrace, + TransferClassifier, +) +from mev_inspect.schemas.transfers import ERC20Transfer + + +class WethTransferClassifier(TransferClassifier): + @staticmethod + def get_transfer(trace: DecodedCallTrace) -> ERC20Transfer: + return ERC20Transfer( + block_number=trace.block_number, + transaction_hash=trace.transaction_hash, + trace_address=trace.trace_address, + amount=trace.inputs["wad"], + to_address=trace.inputs["dst"], + from_address=trace.from_address, + token_address=trace.to_address, + ) + WETH_SPEC = ClassifierSpec( abi_name="WETH9", protocol=Protocol.weth, valid_contract_addresses=["0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2"], - classifications={ - "transferFrom(address,address,uint256)": Classification.transfer, - "transfer(address,uint256)": Classification.transfer, + classifiers={ + "transferFrom(address,address,uint256)": WethTransferClassifier, + "transfer(address,uint256)": WethTransferClassifier, }, ) diff --git a/mev_inspect/classifiers/specs/zero_ex.py b/mev_inspect/classifiers/specs/zero_ex.py index 4f47a03..f558b1d 100644 --- a/mev_inspect/classifiers/specs/zero_ex.py +++ b/mev_inspect/classifiers/specs/zero_ex.py @@ -1,7 +1,9 @@ from mev_inspect.schemas.classified_traces import ( - ClassifierSpec, Protocol, ) +from mev_inspect.schemas.classifiers import ( + ClassifierSpec, +) ZEROX_CONTRACT_SPECS = [ diff --git a/mev_inspect/classifiers/trace.py b/mev_inspect/classifiers/trace.py index 9968eb6..ba7f140 100644 --- a/mev_inspect/classifiers/trace.py +++ b/mev_inspect/classifiers/trace.py @@ -67,8 +67,11 @@ class TraceClassifier: if call_data is not None: signature = call_data.function_signature - classification = spec.classifications.get( - signature, Classification.unknown + classifier = spec.classifiers.get(signature) + classification = ( + Classification.unknown + if classifier is None + else classifier.get_classification() ) return DecodedCallTrace( diff --git a/mev_inspect/inspect_block.py b/mev_inspect/inspect_block.py index 597e2db..e9af60b 100644 --- a/mev_inspect/inspect_block.py +++ b/mev_inspect/inspect_block.py @@ -64,6 +64,8 @@ def inspect_block( write_classified_traces(db_session, classified_traces) transfers = get_transfers(classified_traces) + logger.info(f"Found {len(transfers)} transfers") + if should_write_transfers: delete_transfers_for_block(db_session, block_number) write_transfers(db_session, transfers) diff --git a/mev_inspect/schemas/classified_traces.py b/mev_inspect/schemas/classified_traces.py index d50cb8b..87336b6 100644 --- a/mev_inspect/schemas/classified_traces.py +++ b/mev_inspect/schemas/classified_traces.py @@ -1,15 +1,12 @@ from enum import Enum from typing import Any, Dict, List, Optional -from pydantic import BaseModel - from .blocks import Trace class Classification(Enum): unknown = "unknown" swap = "swap" - burn = "burn" transfer = "transfer" liquidate = "liquidate" @@ -62,12 +59,5 @@ class DecodedCallTrace(CallTrace): protocol: Optional[Protocol] gas: Optional[int] gas_used: Optional[int] - function_name: Optional[str] - function_signature: Optional[str] - - -class ClassifierSpec(BaseModel): - abi_name: str - protocol: Optional[Protocol] = None - valid_contract_addresses: Optional[List[str]] = None - classifications: Dict[str, Classification] = {} + function_name: str + function_signature: str diff --git a/mev_inspect/schemas/classifiers.py b/mev_inspect/schemas/classifiers.py new file mode 100644 index 0000000..5cee8d0 --- /dev/null +++ b/mev_inspect/schemas/classifiers.py @@ -0,0 +1,44 @@ +from abc import ABC, abstractmethod +from typing import Dict, List, Optional, Type + +from pydantic import BaseModel + +from .classified_traces import Classification, DecodedCallTrace, Protocol +from .transfers import ERC20Transfer + + +class Classifier(ABC): + @staticmethod + @abstractmethod + def get_classification() -> Classification: + raise NotImplementedError() + + +class TransferClassifier(Classifier): + @staticmethod + def get_classification() -> Classification: + return Classification.transfer + + @staticmethod + @abstractmethod + def get_transfer(trace: DecodedCallTrace) -> ERC20Transfer: + raise NotImplementedError() + + +class SwapClassifier(Classifier): + @staticmethod + def get_classification() -> Classification: + return Classification.swap + + +class LiquidationClassifier(Classifier): + @staticmethod + def get_classification() -> Classification: + return Classification.liquidate + + +class ClassifierSpec(BaseModel): + abi_name: str + protocol: Optional[Protocol] = None + valid_contract_addresses: Optional[List[str]] = None + classifiers: Dict[str, Type[Classifier]] = {} diff --git a/mev_inspect/transfers.py b/mev_inspect/transfers.py index fbe3a48..90959cd 100644 --- a/mev_inspect/transfers.py +++ b/mev_inspect/transfers.py @@ -1,6 +1,13 @@ -from typing import Dict, List, Optional, Sequence +from typing import Dict, List, Optional, Sequence, Tuple -from mev_inspect.schemas.classified_traces import Classification, ClassifiedTrace +from mev_inspect.classifiers.specs import ALL_CLASSIFIER_SPECS +from mev_inspect.schemas.classifiers import ClassifierSpec, TransferClassifier +from mev_inspect.schemas.classified_traces import ( + Classification, + ClassifiedTrace, + DecodedCallTrace, + Protocol, +) from mev_inspect.schemas.transfers import ERC20Transfer, EthTransfer, TransferGeneric from mev_inspect.traces import is_child_trace_address, get_child_traces @@ -18,9 +25,21 @@ def get_eth_transfers(traces: List[ClassifiedTrace]) -> List[EthTransfer]: def get_transfers(traces: List[ClassifiedTrace]) -> List[ERC20Transfer]: transfers = [] + specs_by_abi_name_and_protocol: Dict[ + Tuple[str, Optional[Protocol]], ClassifierSpec + ] = {(spec.abi_name, spec.protocol): spec for spec in ALL_CLASSIFIER_SPECS} + for trace in traces: - if trace.classification == Classification.transfer: - transfers.append(ERC20Transfer.from_trace(trace)) + if not isinstance(trace, DecodedCallTrace): + continue + + abi_name_and_protocol = (trace.abi_name, trace.protocol) + spec = specs_by_abi_name_and_protocol.get(abi_name_and_protocol) + + if spec is not None: + classifier = spec.classifiers.get(trace.function_signature) + if classifier is not None and issubclass(classifier, TransferClassifier): + transfers.append(classifier.get_transfer(trace)) return transfers diff --git a/tests/helpers.py b/tests/helpers.py index 540a327..21f3264 100644 --- a/tests/helpers.py +++ b/tests/helpers.py @@ -56,6 +56,8 @@ def make_swap_trace( classification=Classification.swap, from_address=from_address, to_address=pool_address, + function_name="swap", + function_signature="swap()", inputs={recipient_input_key: recipient_address}, abi_name=abi_name, block_hash=str(block_number),