Add get_source_info convenience method to SolidityContract

myth

@@ -420,23 +420,14 @@ elif (args.graph) or (args.fire_lasers):
 
                 if (type(contract) == SolidityContract):
 
-                    disassembly = contract.get_disassembly()
-
                     for issue in issues:
 
                         if (issue.pc):
 
-                            index = helper.get_instruction_index(disassembly.instruction_list, issue.pc)
+                            codeinfo = contract.get_source_info(issue.pc)
-
-                            solidity_file = contract.solidity_files[contract.mappings[index].solidity_file_idx]
-
-                            issue.filename = solidity_file.filename
-
-                            offset = contract.mappings[index].offset
-                            length = contract.mappings[index].length
 
-                            issue.code = solidity_file.data[offset:offset + length]
+                            issue.code = codeinfo.code
-                            issue.lineno = contract.mappings[index].lineno
+                            issue.lineno = codeinfo.lineno
 
                 for issue in issues:
                     report.append_issue(issue)  # For text and markdown output

mythril/ether/soliditycontract.py

@@ -1,6 +1,7 @@
 from mythril.ether.ethcontract import ETHContract
 from mythril.ether.util import *
 from mythril.exceptions import NoContractFoundError
+from laser.ethereum import helper
 
 class SourceMapping:
 
@@ -18,6 +19,14 @@ class SolidityFile:
         self.data = data
 
 
+class SourceCodeInfo:
+
+    def __init__(self, filename, lineno, code):
+        self.filename = filename
+        self.lineno = lineno
+        self.code = code
+
+
 class SolidityContract(ETHContract):
 
     def __init__(self, input_file, name=None, solc_args=None):
@@ -82,3 +91,19 @@ class SolidityContract(ETHContract):
             self.mappings.append(SourceMapping(idx, offset, length, lineno))
 
         super().__init__(code, creation_code, name=name)
+
+    def get_source_info(self, address):
+
+        index = helper.get_instruction_index(self.get_disassembly().instruction_list, address)
+
+        solidity_file = self.solidity_files[self.mappings[index].solidity_file_idx]
+
+        filename = solidity_file.filename
+
+        offset = self.mappings[index].offset
+        length = self.mappings[index].length
+
+        code = solidity_file.data[offset:offset + length]
+        lineno = self.mappings[index].lineno
+
+        return SourceCodeInfo(filename, lineno, code)