Merge pull request #27 from lukevs/add-inspectors-and-classifications

Add inspectors and classifications

mev_inspect/decode.py

@@ -0,0 +1,36 @@
+from typing import Dict, Optional
+
+from hexbytes import HexBytes
+from eth_abi import decode_abi
+
+from mev_inspect.schemas.abi import ABI, ABIFunctionDescription
+from mev_inspect.schemas.call_data import CallData
+
+
+class ABIDecoder:
+    def __init__(self, abi: ABI):
+        self._functions_by_selector: Dict[str, ABIFunctionDescription] = {
+            description.get_selector(): description
+            for description in abi
+            if isinstance(description, ABIFunctionDescription)
+        }
+
+    def decode(self, data: str) -> Optional[CallData]:
+        hex_data = HexBytes(data)
+        selector, params = hex_data[:4], hex_data[4:]
+
+        func = self._functions_by_selector.get(selector)
+
+        if func is None:
+            return None
+
+        names = [input.name for input in func.inputs]
+        types = [input.type for input in func.inputs]
+
+        decoded = decode_abi(types, params)
+
+        return CallData(
+            function_name=func.name,
+            function_signature=func.get_signature(),
+            inputs={name: value for name, value in zip(names, decoded)},
+        )

mev_inspect/inspectors/__init__.py

@@ -0,0 +1 @@
+from .base import Inspector

mev_inspect/inspectors/base.py

@@ -0,0 +1,11 @@
+from abc import ABC, abstractmethod
+from typing import Optional
+
+from mev_inspect.schemas.blocks import NestedTrace
+from mev_inspect.schemas.classifications import Classification
+
+
+class Inspector(ABC):
+    @abstractmethod
+    def inspect(self, nested_trace: NestedTrace) -> Optional[Classification]:
+        pass

mev_inspect/inspectors/uniswap.py

@@ -1,9 +1,14 @@
 import json
+from typing import Optional
 
 from web3 import Web3
 
 from mev_inspect import utils
 from mev_inspect.config import load_config
+from mev_inspect.schemas.blocks import NestedTrace, TraceType
+from mev_inspect.schemas.classifications import Classification
+
+from .base import Inspector
 
 config = load_config()
 
@@ -14,7 +19,7 @@ sushiswap_router_address = config["ADDRESSES"]["SushiswapV2Router"]
 uniswap_pair_abi = json.loads(config["ABI"]["UniswapV2Pair"])
 
 
-class UniswapInspector:
+class UniswapInspector(Inspector):
     def __init__(self, base_provider) -> None:
         self.w3 = Web3(base_provider)
 
@@ -79,18 +84,20 @@ class UniswapInspector:
 
         return result
 
-    def inspect(self, calls):
-        for call in calls:
-            print("\n", call)
-            if (
-                call["type"] == "call"
-                and (
-                    call["action"]["to"] == uniswap_router_address.lower()
-                    or call["action"]["to"] == sushiswap_router_address.lower()
-                )
-                and utils.check_trace_for_signature(
-                    call, self.uniswap_router_trade_signatures
-                )
-            ):
-                # print("WIP, here is where there is a call that matches what we are looking for")
-                1 == 1
+    def inspect(self, nested_trace: NestedTrace) -> Optional[Classification]:
+        trace = nested_trace.trace
+
+        if (
+            trace.type == TraceType.call
+            and (
+                trace.action["to"] == uniswap_router_address.lower()
+                or trace.action["to"] == sushiswap_router_address.lower()
+            )
+            and utils.check_trace_for_signature(
+                trace, self.uniswap_router_trade_signatures
+            )
+        ):
+            # print("WIP, here is where there is a call that matches what we are looking for")
+            1 == 1
+
+        return None

mev_inspect/processor.py

@@ -1,15 +1,43 @@
-from mev_inspect.schemas.utils import to_original_json_dict
+from typing import List
+
+from mev_inspect.inspectors import Inspector
+from mev_inspect.schemas.blocks import Block, NestedTrace, TraceType
+from mev_inspect.schemas.classifications import (
+    Classification,
+    UnknownClassification,
+)
+from mev_inspect.traces import as_nested_traces
 
 
 class Processor:
-    def __init__(self, base_provider, inspectors) -> None:
-        self.base_provider = base_provider
-        self.inspectors = inspectors
+    def __init__(self, inspectors: List[Inspector]) -> None:
+        self._inspectors = inspectors
 
-    def get_transaction_evaluations(self, block_data):
-        for transaction_hash in block_data.transaction_hashes:
-            traces = block_data.get_filtered_traces(transaction_hash)
-            traces_json = [to_original_json_dict(trace) for trace in traces]
+    def get_transaction_evaluations(
+        self,
+        block: Block,
+    ) -> List[Classification]:
+        transaction_traces = (
+            trace for trace in block.traces if trace.type != TraceType.reward
+        )
 
-            for inspector in self.inspectors:
-                inspector.inspect(traces_json)
+        return [
+            self._run_inspectors(nested_trace)
+            for nested_trace in as_nested_traces(transaction_traces)
+        ]
+
+    def _run_inspectors(self, nested_trace: NestedTrace) -> Classification:
+        for inspector in self._inspectors:
+            classification = inspector.inspect(nested_trace)
+
+            if classification is not None:
+                return classification
+
+        internal_classifications = [
+            self._run_inspectors(subtrace) for subtrace in nested_trace.subtraces
+        ]
+
+        return UnknownClassification(
+            trace=nested_trace.trace,
+            internal_classifications=internal_classifications,
+        )

mev_inspect/schemas/abi.py

@@ -24,6 +24,7 @@ NON_FUNCTION_DESCRIPTION_TYPES = Union[
 
 
 class ABIDescriptionInput(BaseModel):
+    name: str
     type: str
 
 

mev_inspect/schemas/call_data.py

@@ -0,0 +1,9 @@
+from typing import Any, Dict
+
+from pydantic import BaseModel
+
+
+class CallData(BaseModel):
+    function_name: str
+    function_signature: str
+    inputs: Dict[str, Any]

mev_inspect/schemas/classifications.py

@@ -0,0 +1,14 @@
+from typing import List
+
+from pydantic import BaseModel
+
+from .blocks import Trace
+
+
+class Classification(BaseModel):
+    pass
+
+
+class UnknownClassification(Classification):
+    trace: Trace
+    internal_classifications: List[Classification]

mev_inspect/utils.py

@@ -2,14 +2,16 @@ from typing import List
 
 from hexbytes.main import HexBytes
 
+from mev_inspect.schemas.blocks import Trace
 
-def check_trace_for_signature(trace: dict, signatures: List[str]):
-    if trace["action"]["input"] == None:
+
+def check_trace_for_signature(trace: Trace, signatures: List[str]):
+    if trace.action["input"] == None:
         return False
 
     ## Iterate over all signatures, and if our trace matches any of them set it to True
     for signature in signatures:
-        if HexBytes(trace["action"]["input"]).startswith(signature):
+        if HexBytes(trace.action["input"]).startswith(signature):
             ## Note that we are turning the input into hex bytes here, which seems to be fine
             ## Working with strings was doing weird things
             return True

testing_file.py

@@ -3,7 +3,7 @@ import argparse
 from web3 import Web3
 
 from mev_inspect import block
-from mev_inspect.inspector_uniswap import UniswapInspector
+from mev_inspect.inspectors.uniswap import UniswapInspector
 from mev_inspect.processor import Processor
 
 parser = argparse.ArgumentParser(description="Inspect some blocks.")
@@ -24,11 +24,18 @@ base_provider = Web3.HTTPProvider(args.rpc)
 
 ## Get block data that we need
 block_data = block.create_from_block_number(args.block_number[0], base_provider)
+print(f"Total traces: {len(block_data.traces)}")
+
+total_transactions = len(
+    set(t.transaction_hash for t in block_data.traces if t.transaction_hash is not None)
+)
+print(f"Total transactions: {total_transactions}")
 
 ## Build a Uniswap inspector
 uniswap_inspector = UniswapInspector(base_provider)
 
 ## Create a processor, pass in an ARRAY of inspects
-processor = Processor(base_provider, [uniswap_inspector, uniswap_inspector])
+processor = Processor([uniswap_inspector, uniswap_inspector])
 
-processor.get_transaction_evaluations(block_data)
+classifications = processor.get_transaction_evaluations(block_data)
+print(f"Returned {len(classifications)} classifications")