Refactor Sarif Result (#583)

clavedeluna · web-flow · commit 8cb84ac1e725 · 2024-05-22T18:53:47.000Z
* create base class sarif result

* more refactor

* create sarif loc
diff --git a/src/codemodder/codeql.py b/src/codemodder/codeql.py
@@ -3,7 +3,7 @@
 
 from typing_extensions import Self
 
-from codemodder.result import LineInfo, Location, LocationWithMessage, Result, ResultSet
+from codemodder.result import LineInfo, ResultSet, SarifLocation, SarifResult
 from codemodder.sarifs import AbstractSarifToolDetector
 
 
@@ -13,7 +13,7 @@ def detect(cls, run_data: dict) -> bool:
         return "tool" in run_data and "CodeQL" in run_data["tool"]["driver"]["name"]
 
 
-class CodeQLLocation(Location):
+class CodeQLLocation(SarifLocation):
     @classmethod
     def from_sarif(cls, sarif_location) -> Self:
         artifact_location = sarif_location["physicalLocation"]["artifactLocation"]
@@ -35,46 +35,8 @@ def from_sarif(cls, sarif_location) -> Self:
         return cls(file=file, start=start, end=end)
 
 
-class CodeQLResult(Result):
-    @classmethod
-    def from_sarif(
-        cls, sarif_result, sarif_run, rule_extensions, truncate_rule_id: bool = False
-    ) -> Self:
-        extension_index = sarif_result["rule"]["toolComponent"]["index"]
-        tool_index = sarif_result["rule"]["index"]
-        rule_data = rule_extensions[extension_index]["rules"][tool_index]
-
-        locations: list[Location] = []
-        for location in sarif_result["locations"]:
-            try:
-                codeql_location = CodeQLLocation.from_sarif(location)
-            except KeyError:
-                continue
-
-            locations.append(codeql_location)
-        all_flows: list[list[Location]] = [
-            [
-                CodeQLLocation.from_sarif(locations.get("location"))
-                for locations in threadflow.get("locations", {})
-            ]
-            for codeflow in sarif_result.get("codeFlows", {})
-            for threadflow in codeflow.get("threadFlows", {})
-        ]
-        related_locations: list[LocationWithMessage] = []
-        if "relatedLocations" in sarif_result:
-            related_locations = [
-                LocationWithMessage(
-                    message=rel_location.get("message", {}).get("text", ""),
-                    location=CodeQLLocation.from_sarif(rel_location),
-                )
-                for rel_location in sarif_result.get("relatedLocations", [])
-            ]
-        return cls(
-            rule_id=rule_data["id"],
-            locations=locations,
-            codeflows=all_flows,
-            related_locations=related_locations,
-        )
+class CodeQLResult(SarifResult):
+    location_type = CodeQLLocation
 
 
 class CodeQLResultSet(ResultSet):
@@ -85,11 +47,10 @@ def from_sarif(cls, sarif_file: str | Path, truncate_rule_id: bool = False) -> S
 
         result_set = cls()
         for sarif_run in data["runs"]:
-            rule_extensions = sarif_run["tool"]["extensions"]
             if CodeQLSarifToolDetector.detect(sarif_run):
                 for sarif_result in sarif_run["results"]:
                     codeql_result = CodeQLResult.from_sarif(
-                        sarif_result, sarif_run, rule_extensions, truncate_rule_id
+                        sarif_result, sarif_run, truncate_rule_id
                     )
                     result_set.add_result(codeql_result)
         return result_set
diff --git a/src/codemodder/result.py b/src/codemodder/result.py
@@ -1,11 +1,13 @@
 from __future__ import annotations
 
+from abc import abstractmethod
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, ClassVar, Type
 
 import libcst as cst
 from libcst._position import CodeRange
+from typing_extensions import Self
 
 from codemodder.codetf import Finding
 
@@ -29,6 +31,13 @@ class Location(ABCDataclass):
     end: LineInfo
 
 
+class SarifLocation(Location):
+    @classmethod
+    @abstractmethod
+    def from_sarif(cls, sarif_location) -> Self:
+        pass
+
+
 @dataclass
 class LocationWithMessage:
     location: Location
@@ -58,6 +67,65 @@ def match_location(self, pos: CodeRange, node: cst.CSTNode) -> bool:
         )
 
 
+@dataclass(kw_only=True)
+class SarifResult(Result, ABCDataclass):
+    location_type: ClassVar[Type[SarifLocation]]
+
+    @classmethod
+    def from_sarif(
+        cls, sarif_result, sarif_run, truncate_rule_id: bool = False
+    ) -> Self:
+        return cls(
+            rule_id=cls.extract_rule_id(sarif_result, sarif_run, truncate_rule_id),
+            locations=cls.extract_locations(sarif_result),
+            codeflows=cls.extract_code_flows(sarif_result),
+            related_locations=cls.extract_related_locations(sarif_result),
+        )
+
+    @classmethod
+    def extract_locations(cls, sarif_result) -> list[Location]:
+        return [
+            cls.location_type.from_sarif(location)
+            for location in sarif_result["locations"]
+        ]
+
+    @classmethod
+    def extract_related_locations(cls, sarif_result) -> list[LocationWithMessage]:
+        return [
+            LocationWithMessage(
+                message=rel_location.get("message", {}).get("text", ""),
+                location=cls.location_type.from_sarif(rel_location),
+            )
+            for rel_location in sarif_result.get("relatedLocations", [])
+        ]
+
+    @classmethod
+    def extract_code_flows(cls, sarif_result) -> list[list[Location]]:
+        return [
+            [
+                cls.location_type.from_sarif(locations.get("location"))
+                for locations in threadflow.get("locations", {})
+            ]
+            for codeflow in sarif_result.get("codeFlows", {})
+            for threadflow in codeflow.get("threadFlows", {})
+        ]
+
+    @classmethod
+    def extract_rule_id(cls, result, sarif_run, truncate_rule_id: bool = False) -> str:
+        if rule_id := result.get("ruleId"):
+            return rule_id.split(".")[-1] if truncate_rule_id else rule_id
+
+        # it may be contained in the 'rule' field through the tool component in the sarif file
+        if "rule" in result:
+            tool_index = result["rule"]["toolComponent"]["index"]
+            rule_index = result["rule"]["index"]
+            return sarif_run["tool"]["extensions"][tool_index]["rules"][rule_index][
+                "id"
+            ]
+
+        raise ValueError("Could not extract rule id from sarif result.")
+
+
 @dataclass(kw_only=True)
 class SASTResult(Result):
     finding_id: str
diff --git a/src/codemodder/semgrep.py b/src/codemodder/semgrep.py
@@ -9,7 +9,7 @@
 
 from codemodder.context import CodemodExecutionContext
 from codemodder.logging import logger
-from codemodder.result import LineInfo, Location, LocationWithMessage, Result, ResultSet
+from codemodder.result import LineInfo, Result, ResultSet, SarifLocation, SarifResult
 from codemodder.sarifs import AbstractSarifToolDetector
 
 
@@ -22,7 +22,7 @@ def detect(cls, run_data: dict) -> bool:
         )
 
 
-class SemgrepLocation(Location):
+class SemgrepLocation(SarifLocation):
     @classmethod
     def from_sarif(cls, sarif_location) -> Self:
         artifact_location = sarif_location["physicalLocation"]["artifactLocation"]
@@ -40,60 +40,8 @@ def from_sarif(cls, sarif_location) -> Self:
         return cls(file=file, start=start, end=end)
 
 
-class SemgrepResult(Result):
-    @classmethod
-    def extract_rule_id(
-        cls, result, sarif_run, truncate_rule_id: bool = False
-    ) -> Optional[str]:
-        if rule_id := result.get("ruleId"):
-            return rule_id.split(".")[-1] if truncate_rule_id else rule_id
-
-        # it may be contained in the 'rule' field through the tool component in the sarif file
-        if "rule" in result:
-            tool_index = result["rule"]["toolComponent"]["index"]
-            rule_index = result["rule"]["index"]
-            return sarif_run["tool"]["extensions"][tool_index]["rules"][rule_index][
-                "id"
-            ]
-
-        return None
-
-    @classmethod
-    def from_sarif(
-        cls, sarif_result, sarif_run, truncate_rule_id: bool = False
-    ) -> Self:
-        if not (
-            rule_id := cls.extract_rule_id(sarif_result, sarif_run, truncate_rule_id)
-        ):
-            raise ValueError("Could not extract rule id from sarif result.")
-
-        locations: list[Location] = []
-        for location in sarif_result["locations"]:
-            artifact_location = SemgrepLocation.from_sarif(location)
-            locations.append(artifact_location)
-        all_flows: list[list[Location]] = [
-            [
-                SemgrepLocation.from_sarif(locations.get("location"))
-                for locations in threadflow.get("locations", {})
-            ]
-            for codeflow in sarif_result.get("codeFlows", {})
-            for threadflow in codeflow.get("threadFlows", {})
-        ]
-        related_locations: list[LocationWithMessage] = []
-        if "relatedLocations" in sarif_result:
-            related_locations = [
-                LocationWithMessage(
-                    message=rel_location.get("message", {}).get("text", ""),
-                    location=SemgrepLocation.from_sarif(rel_location),
-                )
-                for rel_location in sarif_result.get("relatedLocations", [])
-            ]
-        return cls(
-            rule_id=rule_id,
-            locations=locations,
-            codeflows=all_flows,
-            related_locations=related_locations,
-        )
+class SemgrepResult(SarifResult):
+    location_type = SemgrepLocation
 
 
 class SemgrepResultSet(ResultSet):
