From 8201bfd75f97cfcf717e3ea82205ec85928933fe Mon Sep 17 00:00:00 2001
From: Nikhil Parasaram <nikhilparasaram@gmail.com>
Date: Thu, 29 Aug 2024 09:41:18 +0100
Subject: [PATCH] Fix issues with statespace on solc 0.8.26 (#1863)

* Fix issues with statespace on solc 0.8.26

* Sort the import
---
 mythril/analysis/callgraph.py | 46 +++++++++++++++++------------------
 mythril/laser/ethereum/cfg.py | 38 ++++++++++++++++-------------
 2 files changed, 43 insertions(+), 41 deletions(-)

diff --git a/mythril/analysis/callgraph.py b/mythril/analysis/callgraph.py
index dcbd0216d..35c041442 100644
--- a/mythril/analysis/callgraph.py
+++ b/mythril/analysis/callgraph.py
@@ -2,6 +2,7 @@
 graphs."""
 
 import re
+from typing import Dict, List
 
 from jinja2 import Environment, PackageLoader, select_autoescape
 from z3 import Z3Exception
@@ -125,56 +126,52 @@
 }
 
 
-def extract_nodes(statespace):
+def extract_nodes(statespace) -> List[Dict]:
     """
+    Extract nodes from the given statespace and create a list of node dictionaries
+    with visual attributes for graph representation.
 
-    :param statespace:
-    :param color_map:
-    :return:
+    :param statespace: The statespace object containing nodes and states information.
+    :return: A list of dictionaries representing each node with its attributes.
     """
     nodes = []
     color_map = {}
-    for node_key in statespace.nodes:
-        node = statespace.nodes[node_key]
+    for node_key, node in statespace.nodes.items():
         instructions = [state.get_current_instruction() for state in node.states]
         code_split = []
+
         for instruction in instructions:
-            if instruction["opcode"].startswith("PUSH"):
-                code_line = "%d %s %s" % (
-                    instruction["address"],
-                    instruction["opcode"],
-                    instruction["argument"],
-                )
+            address = instruction["address"]
+            opcode = instruction["opcode"]
+            if opcode.startswith("PUSH"):
+                code_line = f"{address} {opcode} {instruction.get('argument', '')}"
             elif (
-                instruction["opcode"].startswith("JUMPDEST")
+                opcode.startswith("JUMPDEST")
                 and NodeFlags.FUNC_ENTRY in node.flags
-                and instruction["address"] == node.start_addr
+                and address == node.start_addr
             ):
                 code_line = node.function_name
             else:
-                code_line = "%d %s" % (instruction["address"], instruction["opcode"])
+                code_line = f"{address} {opcode}"
 
-            code_line = re.sub(
-                "([0-9a-f]{8})[0-9a-f]+", lambda m: m.group(1) + "(...)", code_line
-            )
+            code_line = re.sub(r"([0-9a-f]{8})[0-9a-f]+", r"\1(...)", code_line)
             code_split.append(code_line)
 
         truncated_code = (
             "\n".join(code_split)
-            if (len(code_split) < 7)
+            if len(code_split) < 7
             else "\n".join(code_split[:6]) + "\n(click to expand +)"
         )
 
-        if node.get_cfg_dict()["contract_name"] not in color_map.keys():
+        contract_name = node.get_cfg_dict()["contract_name"]
+        if contract_name not in color_map:
             color = default_colors[len(color_map) % len(default_colors)]
-            color_map[node.get_cfg_dict()["contract_name"]] = color
+            color_map[contract_name] = color
 
         nodes.append(
             {
                 "id": str(node_key),
-                "color": color_map.get(
-                    node.get_cfg_dict()["contract_name"], default_colors[0]
-                ),
+                "color": color_map.get(contract_name, default_colors[0]),
                 "size": 150,
                 "fullLabel": "\n".join(code_split),
                 "label": truncated_code,
@@ -182,6 +179,7 @@ def extract_nodes(statespace):
                 "isExpanded": False,
             }
         )
+
     return nodes
 
 
diff --git a/mythril/laser/ethereum/cfg.py b/mythril/laser/ethereum/cfg.py
index 2b32a39c5..65aa5e8e1 100644
--- a/mythril/laser/ethereum/cfg.py
+++ b/mythril/laser/ethereum/cfg.py
@@ -59,25 +59,29 @@ def __init__(
 
     def get_cfg_dict(self) -> Dict:
         """
+        Generate a configuration dictionary for the current state of the contract.
 
-        :return:
+        :return: A dictionary containing the contract's configuration details.
         """
-        code = ""
-        for state in self.states:
-            instruction = state.get_current_instruction()
-
-            code += str(instruction["address"]) + " " + instruction["opcode"]
-            if instruction["opcode"].startswith("PUSH"):
-                code += " " + "".join(str(instruction["argument"]))
-
-            code += "\\n"
-
-        return dict(
-            contract_name=self.contract_name,
-            start_addr=self.start_addr,
-            function_name=self.function_name,
-            code=code,
-        )
+        code_lines = [
+            f"{instruction['address']} {instruction['opcode']}"
+            + (
+                f" {instruction['argument']}"
+                if instruction["opcode"].startswith("PUSH")
+                and "argument" in instruction
+                else ""
+            )
+            for state in self.states
+            for instruction in [state.get_current_instruction()]
+        ]
+        code = "\\n".join(code_lines)
+
+        return {
+            "contract_name": self.contract_name,
+            "start_addr": self.start_addr,
+            "function_name": self.function_name,
+            "code": code,
+        }
 
 
 class Edge: