From cd08dbe16e67fc62f9af2deb9d1f8c6381ae2815 Mon Sep 17 00:00:00 2001
From: NotPeopling2day <32708219+NotPeopling2day@users.noreply.github.com>
Date: Tue, 10 May 2022 00:39:45 +0200
Subject: [PATCH] feat: add calltree creation method (#4)

---
 .pre-commit-config.yaml |   5 --
 evm_trace/__init__.py   |   4 +-
 evm_trace/base.py       | 127 +++++++++++++++++++++++++++++++++++++++-
 pyproject.toml          |   2 -
 setup.py                |   1 +
 5 files changed, 129 insertions(+), 10 deletions(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index ac9e94f..5903073 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -4,11 +4,6 @@ repos:
     hooks:
     -   id: check-yaml
 
--   repo: https://github.com/asottile/seed-isort-config
-    rev: v2.2.0
-    hooks:
-    -   id: seed-isort-config
-
 -   repo: https://github.com/pre-commit/mirrors-isort
     rev: v5.9.3
     hooks:
diff --git a/evm_trace/__init__.py b/evm_trace/__init__.py
index d9c9ee1..4d84409 100644
--- a/evm_trace/__init__.py
+++ b/evm_trace/__init__.py
@@ -1,3 +1,3 @@
-from .base import CallTreeNode, CallType, TraceFrame
+from .base import CallTreeNode, CallType, TraceFrame, get_calltree_from_trace
 
-__all__ = ["CallTreeNode", "CallType", "TraceFrame"]
+__all__ = ["CallTreeNode", "CallType", "get_calltree_from_trace", "TraceFrame"]
diff --git a/evm_trace/base.py b/evm_trace/base.py
index 3882d4a..80fa0d1 100644
--- a/evm_trace/base.py
+++ b/evm_trace/base.py
@@ -1,6 +1,8 @@
+import math
 from enum import Enum
-from typing import Any, Dict, List
+from typing import Any, Dict, Iterator, List
 
+from eth_utils import to_int
 from hexbytes import HexBytes
 from pydantic import BaseModel, Field, ValidationError, validator
 
@@ -43,6 +45,7 @@ class CallTreeNode(BaseModel):
     call_type: CallType
     address: Any
     value: int = 0
+    depth: int = 0
     gas_limit: int
     gas_cost: int  # calculated from call starting and return
     calldata: Any = HexBytes(b"")
@@ -54,3 +57,125 @@ class CallTreeNode(BaseModel):
     @validator("address", "calldata", "returndata", pre=True)
     def validate_hexbytes(cls, v) -> HexBytes:
         return _convert_hexbytes(cls, v)
+
+
+def get_calltree_from_trace(
+    trace: Iterator[TraceFrame], show_internal=False, **root_node_kwargs
+) -> CallTreeNode:
+    """
+    Creates a CallTreeNode from a given transaction trace.
+
+    Args:
+        trace (Iterator[TraceFrame]): Iterator of transaction trace frames.
+        show_internal (bool): Boolean whether to display internal calls. Defaulted to False.
+        root_node_kwargs (dict): Keyword argments passed to the root ``CallTreeNode``.
+
+    Returns:
+        :class:`~evm_trace.base.CallTreeNode: Call tree of transaction trace.
+    """
+
+    return _create_node_from_call(
+        trace=trace,
+        show_internal=show_internal,
+        **root_node_kwargs,
+    )
+
+
+def _extract_memory(offset: HexBytes, size: HexBytes, memory: List[HexBytes]) -> HexBytes:
+    """
+    Extracts memory from the EVM stack.
+
+    Args:
+        offset (HexBytes): Offset byte location in memory.
+        size (HexBytes): Number of bytes to return.
+        memory (List[HexBytes]): Memory stack.
+
+    Returns:
+        HexBytes: Byte value from memory stack.
+    """
+
+    size_int = to_int(size)
+
+    if size_int == 0:
+        return HexBytes(b"")
+
+    offset_int = to_int(offset)
+
+    # Compute the word that contains the first byte
+    start_word = math.floor(offset_int / 32)
+    # Compute the word that contains the last byte
+    stop_word = math.ceil((offset_int + size_int) / 32)
+
+    byte_slice = b""
+
+    for word in memory[start_word:stop_word]:
+        byte_slice += word
+
+    offset_index = offset_int % 32
+    return HexBytes(byte_slice[offset_index:size_int])
+
+
+def _create_node_from_call(
+    trace: Iterator[TraceFrame], show_internal: bool = False, **node_kwargs
+) -> CallTreeNode:
+    """
+    Use specified opcodes to create a branching callnode
+    https://www.evm.codes/
+    """
+
+    if show_internal:
+        raise NotImplementedError()
+
+    node = CallTreeNode(**node_kwargs)
+
+    for frame in trace:
+        if frame.op in ("CALL", "DELEGATECALL", "STATICCALL"):
+            child_node_kwargs = {}
+
+            child_node_kwargs["address"] = frame.stack[-2][-20:]  # address is 20 bytes in EVM
+            child_node_kwargs["depth"] = frame.depth
+            # TODO: Validate gas values
+            child_node_kwargs["gas_limit"] = int(frame.stack[-1].hex(), 16)
+            child_node_kwargs["gas_cost"] = frame.gas_cost
+
+            if frame.op == "CALL":
+                child_node_kwargs["call_type"] = CallType.MUTABLE
+                child_node_kwargs["value"] = int(frame.stack[-3].hex(), 16)
+                child_node_kwargs["calldata"] = _extract_memory(
+                    offset=frame.stack[-4], size=frame.stack[-5], memory=frame.memory
+                )
+            elif frame.op == "DELEGATECALL":
+                child_node_kwargs["call_type"] = CallType.DELEGATE
+                child_node_kwargs["calldata"] = _extract_memory(
+                    offset=frame.stack[-3], size=frame.stack[-4], memory=frame.memory
+                )
+            else:
+                child_node_kwargs["call_type"] = CallType.STATIC
+                child_node_kwargs["calldata"] = _extract_memory(
+                    offset=frame.stack[-3], size=frame.stack[-4], memory=frame.memory
+                )
+
+            child_node = _create_node_from_call(trace=trace, **child_node_kwargs)
+            node.calls.append(child_node)
+
+        # TODO: Handle internal nodes using JUMP and JUMPI
+
+        elif frame.op in ("SELFDESTRUCT", "STOP"):
+            # TODO: Handle the internal value transfer in SELFDESTRUCT
+            # TODO: Handle "execution halted" vs. gas limit reached
+            node.selfdestruct = frame.op == "SELFDESTRUCT"
+            break
+
+        elif frame.op in ("RETURN", "REVERT"):
+            node.returndata = _extract_memory(
+                offset=frame.stack[-1], size=frame.stack[-2], memory=frame.memory
+            )
+            # TODO: Handle "execution halted" vs. gas limit reached
+            node.failed = frame.op == "REVERT"
+            break
+
+        # TODO: Handle invalid opcodes (`node.failed = True`)
+        # NOTE: ignore other opcodes
+
+    # TODO: Handle "execution halted" vs. gas limit reached
+    return node
diff --git a/pyproject.toml b/pyproject.toml
index 443eb0c..f2ce3dc 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -35,7 +35,5 @@ markers = "fuzzing: Run Hypothesis fuzz test suite"
 line_length = 100
 force_grid_wrap = 0
 include_trailing_comma = true
-known_third_party = ["hexbytes", "pydantic", "pytest", "setuptools"]
-known_first_party = ["MODULE_NAME"]
 multi_line_output = 3
 use_parentheses = true
diff --git a/setup.py b/setup.py
index 6b49741..c24c1a5 100644
--- a/setup.py
+++ b/setup.py
@@ -56,6 +56,7 @@
         "importlib-metadata ; python_version<'3.8'",
         "pydantic>=1.9.0,<2.0",
         "hexbytes>=0.2.2,<1.0.0",
+        "eth-utils==1.10.0",
     ],  # NOTE: Add 3rd party libraries here
     python_requires=">=3.7.2,<4",
     extras_require=extras_require,