Rewrite loader logic

mythril/support/loader.py

@@ -6,36 +6,43 @@ class DynLoader:
 
     def __init__(self, eth):
         self.eth = eth
+        self.storage_cache = {}
 
-    def dynld(self, contract_address, dependency_address):
 
-        logging.info("Dynld at contract " + contract_address + ": "  + dependency_address)
+    def read_storage(self, contract_address, index):
 
-        m = re.match(r'^(0x[0-9a-fA-F]{40})$', dependency_address)
+        try:
+            contract_ref = self.storage_cache[contract_address]
+            data = contract_ref[index]
 
-        if (m):
-            dependency_address = m.group(1)
+        except KeyError:
 
-        else:
-            m = re.search(r'storage_(\d+)', dependency_address)
+            self.storage_cache[contract_address] = {}
+
+            data = self.eth.eth_getStorageAt(contract_address, position=index, block='latest')
+
+            self.storage_cache[contract_address][index] = data
+
+        except IndexError:
 
-            if (m):
-                idx = int(m.group(1))
-                logging.info("Dynamic contract address at storage index " + str(idx))
+            data = self.eth.eth_getStorageAt(contract_address, position=index, block='latest')
 
-                # testrpc simply returns the address, geth response is more elaborate.
+            self.storage_cache[contract_address][index] = data
 
-                dependency_address = self.eth.eth_getStorageAt(contract_address, position=idx, block='latest')
+        return data
 
-                if not re.match(r"^0x[0-9a-f]{40}$", dependency_address):
 
-                    dependency_address = "0x" + dependency_address[26:]
-
+    def dynld(self, contract_address, dependency_address):
 
-            else:
-                logging.info("Unable to resolve address.")
-                return None
+        logging.info("Dynld at contract " + contract_address + ": "  + dependency_address)
 
+        m = re.match(r'^(0x[0-9a-fA-F]{40})$', dependency_address)
+
+        if (m):
+            dependency_address = m.group(1)
+
+        else:
+            return None
 
         logging.info("Dependency address: " +  dependency_address)