Merge pull request cleanunicorn#5 from cleanunicorn/exploit-description

Display exploit description on startup.

theo/exploit/exploit.py

@@ -11,10 +11,17 @@ def __init__(
         contract: str,
         account: str,
         account_pk: str,
+        title: str = None,
+        description: str = None,
+        swc_id: int = 0,
         verbosity: int = logging.INFO,
     ):
         # Web3 instance (can be HTTP, WebSockets, IPC)
         self.w3 = w3
+        # Meta information
+        self.title = title
+        self.description = description
+        self.swc_id = swc_id
         # Transaction list
         self.txs = txs
 
@@ -51,7 +58,15 @@ def __init__(
             self.logger.addHandler(logger_stream)
 
     def __repr__(self):
-        return "Exploit: (txs={})".format(self.txs)
+        return """Exploit: {title}
+Description: {description}
+SWC ID: {swc_id}
+Transacion list: {txs}""".format(
+            title=self.title,
+            description=self.description,
+            swc_id=self.swc_id,
+            txs=self.txs,
+        )
 
     def execute(self, nonce=None):
         receipts = []
@@ -216,4 +231,4 @@ def dump_to_file(self, file=None):
         for tx in self.txs:
             exploit_object.append(tx.__dict__)
 
-        dump(ob=[exploit_object], filename=file)
+        dump(ob=[exploit_object], filename=file)

theo/exploit/tx.py

@@ -1,5 +1,5 @@
 class Tx:
-    def __init__(self, data: str, value: str):
+    def __init__(self, data: str, value: str, name: str = None):
         # Transaction input (data)
         self.data = data
         # Transaction value
@@ -11,8 +11,13 @@ def __init__(self, data: str, value: str):
                 self.value = int(value)
         else:
             self.value = int(value)
+        # Name
+        self.name = name
 
     def __repr__(self):
-        return "Transaction {{Data: {input}, Value: {value_eth:.2f} ether ({value})}}".format(
-            input=self.data, value_eth=self.value / 10 ** 18, value=self.value
+        return "Transaction {{Name: {name}, Data: {input}, Value: {value_eth:.2f} ether ({value})}}".format(
+            name=self.name,
+            input=self.data,
+            value_eth=self.value / 10 ** 18,
+            value=self.value,
         )

theo/interfaces/cli.py

@@ -100,7 +100,9 @@ def start_repl(args):
     if len(exploits) == 0:
         print("No exploits found. You're going to need to load some exploits.")
     else:
-        print("Found exploits(s):\n", exploits)
+        print("")
+        print("Found exploits(s):")
+        print(exploits)
 
     # Create a web3 instance
     w3 = Web3(Web3.HTTPProvider(args.rpc_http, request_kwargs={"timeout": 60}))
@@ -123,8 +125,6 @@ def start_repl(args):
     del os, atexit, readline, rlcompleter
     code.InteractiveConsole(vars).interact(
         banner="""
-Theo version {version}.
-
 Tools available in the console:
 - `exploits` is an array of loaded exploits found by Mythril or read from a file
 - `w3` an initialized instance of web3py for the provided HTTP RPC endpoint
@@ -133,6 +133,8 @@ def start_repl(args):
 Check the readme for more info:
 https://github.com/cleanunicorn/theo
 
+Theo version {version}.
+
 """.format(
             version=__version__
         )

theo/scanner/__init__.py

@@ -71,12 +71,13 @@ def exploits_from_mythril(
         w3 = Web3(Web3.HTTPProvider(rpcHTTP))
 
     exploits = []
+
     for ri in report.issues:
         txs = []
         issue = report.issues[ri]
 
         for si in issue.transaction_sequence["steps"]:
-            txs.append(Tx(data=si["input"], value=si["value"]))
+            txs.append(Tx(data=si["input"], value=si["value"], name=si["name"]))
 
         exploits.append(
             Exploit(
@@ -85,6 +86,9 @@ def exploits_from_mythril(
                 contract=contract,
                 account=private_key_to_account(account_pk),
                 account_pk=account_pk,
+                title=issue.title,
+                description=issue.description,
+                swc_id=issue.swc_id,
             )
         )