Rewrite loader logic

mythril/support/loader.py

@@ -6,37 +6,44 @@ 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):
+        except KeyError:
-            dependency_address = m.group(1)
 
-        else:
+            self.storage_cache[contract_address] = {}
-            m = re.search(r'storage_(\d+)', dependency_address)
 
-            if (m):
+            data = self.eth.eth_getStorageAt(contract_address, position=index, block='latest')
-                idx = int(m.group(1))
+
-                logging.info("Dynamic contract address at storage index " + str(idx))
+            self.storage_cache[contract_address][index] = data
+
+        except IndexError:
 
-                # testrpc simply returns the address, geth response is more elaborate.
+            data = self.eth.eth_getStorageAt(contract_address, position=index, block='latest')
 
-                dependency_address = self.eth.eth_getStorageAt(contract_address, position=idx, block='latest')
+            self.storage_cache[contract_address][index] = data
 
-                if not re.match(r"^0x[0-9a-f]{40}$", dependency_address):
+        return data
 
-                    dependency_address = "0x" + dependency_address[26:]
 
+    def dynld(self, contract_address, dependency_address):
+
+        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:
-                logging.info("Unable to resolve address.")
             return None
 
-
         logging.info("Dependency address: " +  dependency_address)
 
         code = self.eth.eth_getCode(dependency_address)