Add CodeQL SARIF support to viz

Author: mschwager

CHANGELOG.md

@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 - Official support for Python 3.13
+- CodeQL SARIF file support, pass `--codeql` to `viz`
 
 ### Removed
 

routes/commands/viz.py

@@ -1,5 +1,4 @@
 import collections
-import json
 import logging
 import pathlib
 import webbrowser
@@ -16,7 +15,7 @@ def get_global(global_results, _global):
         return None
 
     global_locations = " ".join(
-        f"{r.check_id}:{r.path}:{r.start_line}" for r in global_results
+        f"{r.id}:{r.path}:{r.start_line}" for r in global_results
     )
 
     if len(global_results) > 1:
@@ -44,12 +43,7 @@ def d3ify(parts, output, result, _global):
         d3ify(parts, new_output, result, _global)
     else:
         name = f"ln {result.start_line}: {result.first_line}"
-
-        if _global:
-            fill = _global.rd_fill
-        else:
-            fill = result.rd_fill
-
+        fill = _global.rd_fill if _global else result.rd_fill
         check_node = {"name": name, "fill": fill, "title": result.lines}
         new_node.setdefault("children", []).append(check_node)
 
@@ -66,17 +60,18 @@ def merge_d3_results(d1s, d2s):
 
 
 def main(args):
-    logger.info("Reading input file %s", args.input.name)
-    data = json.load(args.input)
+    logger.info("Processing input file %s", args.input.name)
+
+    output_cls = types.CodeQLOutput if args.codeql else types.SemgrepOutput
+    output = output_cls.from_fd(args.input)
+    results = output.results
 
-    semgrep_results = [types.SemgrepResult(r) for r in data["results"]]
-    counts = collections.Counter([r.check_id for r in semgrep_results])
-    count_output = " ".join(f"{k}={v}" for k, v in counts.items())
-    logger.info("Finding rule counts: %s", count_output)
+    for result_id, count in collections.Counter([r.id for r in results]).items():
+        logger.info("Found %d results for id %s", count, result_id)
 
     results_by_type = {
         key: list(group)
-        for key, group in util.sorted_groupby(semgrep_results, key=lambda r: r.rd_type)
+        for key, group in util.sorted_groupby(results, key=lambda r: r.rd_type)
     }
 
     global_results = results_by_type.get(types.ResultType.GLOBAL.value, {})
@@ -86,7 +81,7 @@ def main(args):
     d3_results = []
     for result in results_by_type.get(types.ResultType.ROUTE.value, []):
         path = pathlib.PurePath(result.path)
-        logger.debug("Processing %s:%s:%s", result.check_id, path, result.start_line)
+        logger.debug("Processing %s:%s:%s", result.id, path, result.start_line)
         root, *_ = path.parts
         root_paths.add(root)
         output = []

routes/main.py

@@ -66,15 +66,16 @@ def parse_args(args=None):
         help="Open HTML output file in browser",
     )
     viz_parser.add_argument(
-        "--global",
-        dest="_global",
+        "-c",
+        "--codeql",
         action="store_true",
-        help="Expiremental: enable global security configuration detection",
+        help="Parse input file as CodeQL SARIF format",
     )
     viz_parser.add_argument(
-        "--interprocedural",
+        "--global",
+        dest="_global",
         action="store_true",
-        help="Expiremental: enable interprocedural security configuration detection",
+        help="Expiremental: enable global security configuration detection",
     )
 
     return p.parse_args(args=args)

routes/types.py

@@ -1,4 +1,6 @@
+import abc
 import enum
+import json
 
 from routes import const
 
@@ -8,12 +10,61 @@ class ResultType(enum.Enum):
     GLOBAL = "global"
 
 
-class SemgrepResult:
+class BaseOutput:
+    def __init__(self, output):
+        self.output = output
+
+    @classmethod
+    def from_fd(cls, fd):
+        return cls(json.load(fd))
+
+    @property
+    @abc.abstractmethod
+    def results(self):
+        pass
+
+
+class BaseResult(abc.ABC):
+    @property
+    @abc.abstractmethod
+    def id(self):
+        pass
+
+    @property
+    @abc.abstractmethod
+    def path(self):
+        pass
+
+    @property
+    @abc.abstractmethod
+    def start_line(self):
+        pass
+
+    @property
+    @abc.abstractmethod
+    def lines(self):
+        pass
+
+    @property
+    @abc.abstractmethod
+    def first_line(self):
+        pass
+
+    @property
+    def rd_type(self):
+        return ResultType.ROUTE.value
+
+    @property
+    def rd_fill(self):
+        return const.DEFAULT_FILL_COLOR
+
+
+class SemgrepResult(BaseResult):
     def __init__(self, result):
         self.result = result
 
     @property
-    def check_id(self):
+    def id(self):
         return self.result["check_id"]
 
     @property
@@ -55,3 +106,74 @@ def rd_type(self):
     @property
     def rd_fill(self):
         return self.rd_metadata.get("fill", const.DEFAULT_FILL_COLOR)
+
+
+class SemgrepOutput(BaseOutput):
+    @property
+    def results(self):
+        return [SemgrepResult(r) for r in self.output["results"]]
+
+
+class CodeQLResult(BaseResult):
+    def __init__(self, result, output):
+        self.result = result
+        self.output = output
+
+        # Assume last location provides the relevant code snippet
+        self.location = self.result["locations"][-1]["physicalLocation"]
+
+    @property
+    def id(self):
+        return self.result["rule"]["id"]
+
+    @property
+    def path(self):
+        # TODO find more robust way to ensure all result paths have a single root directory
+        return "repo" + "/" + self.location["artifactLocation"]["uri"]
+
+    @property
+    def artifact_index(self):
+        return self.location["artifactLocation"]["index"]
+
+    @property
+    def start_line(self):
+        return self.location["region"]["startLine"]
+
+    @property
+    def end_line(self):
+        return self.location["region"].get("endLine", self.start_line)
+
+    @property
+    def lines(self):
+        artifact = self.output.first_run["artifacts"][self.artifact_index]
+
+        # --sarif-add-file-contents provides this data
+        contents = artifact.get("contents", {}).get("text", "")
+        if not contents:
+            return []
+
+        context = 1
+
+        # -1 for 0-based indexing
+        start = max(self.start_line - context - 1, 0)
+        end = self.end_line + context
+        lines = contents.split("\n")
+
+        return "\n".join(lines[start:end])
+
+    @property
+    def first_line(self):
+        # We provide 1 line of context above and below the result's line, so
+        # the result's line should be the middle list index of 3 elements
+        return self.lines.split("\n")[1]
+
+
+class CodeQLOutput(BaseOutput):
+    @property
+    def first_run(self):
+        # Assume we only have a single run
+        return self.output["runs"][0]
+
+    @property
+    def results(self):
+        return [CodeQLResult(r, self) for r in self.first_run["results"]]

tests/test_commands/test_viz.py

@@ -218,7 +218,7 @@ def test_viz_basic(data, expected):
         output=output,
         template=template,
         browser=False,
-        interprocedural=True,
+        codeql=False,
         _global=True,
     )
 
@@ -249,7 +249,7 @@ def test_viz_multiple_root():
         output=output,
         template=template,
         browser=False,
-        interprocedural=True,
+        codeql=False,
         _global=True,
     )