Add main source support to CLI

pythx/cli/main.py

@@ -171,7 +171,8 @@ def top(config, staging, interval):
 @opts.source_file_opt
 @opts.solc_path_opt
 @opts.no_cache_opt
-def check(config, staging, bytecode_file, source_file, solc_path, no_cache):
+@opts.entrypoint_opt
+def check(config, staging, bytecode_file, source_file, solc_path, no_cache, entrypoint):
     """Submit a new analysis job based on source code, byte code, or both.
 
     :param no_cache: Disable the API's result cache
@@ -180,6 +181,7 @@ def check(config, staging, bytecode_file, source_file, solc_path, no_cache):
     :param bytecode_file: A file specifying the bytecode to analyse
     :param source_file: A file specifying the source code to analyse
     :param solc_path: The path to the solc compiler to compile the source code
+    :param entrypoint: The main Solidity file (e.g. passed to solc)
     """
     c = utils.recover_client(config_path=config, staging=staging, no_cache=no_cache)
     if bytecode_file:
@@ -191,7 +193,16 @@ def check(config, staging, bytecode_file, source_file, solc_path, no_cache):
         with open(source_file, "r") as source_f:
             source_content = source_f.read().strip()
         compiled = utils.compile_from_source(source_file, solc_path=solc_path)
-        if len(compiled["contracts"]) > 1:
+
+        # try to get main source file
+        if entrypoint is None:
+            if len(compiled["sourceList"]) == 1:
+                entrypoint = compiled["sourceList"][0]
+            else:
+                click.echo("Please provide an entrypoint with --entrypoint/-e")
+                sys.exit(1)
+
+        if len(compiled["contracts"]) != 1:
             click.echo(
                 (
                     "The PythX CLI currently does not support sending multiple contracts. "
@@ -202,18 +213,24 @@ def check(config, staging, bytecode_file, source_file, solc_path, no_cache):
             )
             sys.exit(1)
 
-        for _, contract_data in compiled["contracts"].items():
-            bytecode = contract_data["bin"]
-            source_map = contract_data["srcmap"]
+        _, contract_data = compiled["contracts"].popitem()
+        bytecode = contract_data["bin"]
+        deployed_bytecode = contract_data["bin-runtime"]
+        source_map = contract_data["srcmap"]
+        deployed_source_map = contract_data["srcmap-runtime"]
 
         sources_dict = copy(compiled["sources"])
+
         sources_dict[source_file]["source"] = source_content
         resp = c.analyze(
             bytecode=bytecode,
+            deployed_bytecode=deployed_bytecode,
             source_map=source_map,
+            deployed_source_map=deployed_source_map,
             source_list=compiled["sourceList"],
             sources=sources_dict,
             solc_version=compiled["version"],
+            main_source=entrypoint,
         )
     else:
         click.echo("Please pass a bytecode or a source code file")
@@ -284,6 +301,7 @@ def truffle(config, staging, no_cache):
                 }
             },
             source_list=[artifact.get("sourcePath")],
+            main_source=artifact.get("sourcePath"),
             solc_version=artifact["compiler"]["version"],
         )
         jobs.append(resp.uuid)

pythx/cli/opts.py

@@ -60,3 +60,9 @@
 no_cache_opt = click.option(
     "--no-cache", is_flag=True, default=False, help="Disable the API's request cache"
 )
+entrypoint_opt = click.option(
+    "--entrypoint",
+    "-e",
+    type=click.Path(exists=True, readable=True),
+    help="The main Solidity file called by the user"
+)

pythx/cli/utils.py

@@ -182,8 +182,14 @@ def compile_from_source(source_path: str, solc_path: str = None):
         "ast,bin,bin-runtime,srcmap,srcmap-runtime",
         source_path,
     ]
-    output = check_output(solc_command)
-    return json.loads(output)
+    output = json.loads(check_output(solc_command))
+
+    # we need the AST key to be lowercase
+    for fname, content in output["sources"].items():
+        content["ast"] = content["AST"]
+        del content["AST"]
+
+    return output
 
 
 def echo_report_as_table(resp):