diff --git a/sapp/pipeline/__init__.py b/sapp/pipeline/__init__.py index b5ce3633..a0fd3e31 100644 --- a/sapp/pipeline/__init__.py +++ b/sapp/pipeline/__init__.py @@ -107,6 +107,14 @@ class ParseTypeInterval(NamedTuple): preserves_type_context: bool +class ParseTraceAnnotationSubtrace(NamedTuple): + callee: str + port: str + position: SourceLocation + features: List["ParseTraceFeature"] = [] + annotations: List["ParseTraceAnnotation"] = [] + + class ParseTraceAnnotation(NamedTuple): location: SourceLocation kind: str @@ -117,9 +125,7 @@ class ParseTraceAnnotation(NamedTuple): link: Optional[str] trace_key: Optional[str] titos: List[SourceLocation] - subtraces: List[ - Dict[str, Any] - ] # callee, port, features, annotations (TODO figure what exactly this shape is) + subtraces: List[ParseTraceAnnotationSubtrace] @staticmethod def from_json(j: Dict[str, Any]) -> "ParseTraceAnnotation": diff --git a/sapp/pipeline/mariana_trench_parser.py b/sapp/pipeline/mariana_trench_parser.py index 2a8800df..2b4282c2 100644 --- a/sapp/pipeline/mariana_trench_parser.py +++ b/sapp/pipeline/mariana_trench_parser.py @@ -342,11 +342,11 @@ def from_taint_json(caller: Method, extra_trace: Dict[str, Any]) -> "ExtraTrace" def to_sapp(self) -> sapp.ParseTraceAnnotation: subtraces = ( [ - { - "callee": self.callee.method.name, - "port": self.callee.port.value, - "position": self.callee.position.to_sapp(), - } + sapp.ParseTraceAnnotationSubtrace( + callee=self.callee.method.name, + port=self.callee.port.value, + position=self.callee.position.to_sapp(), + ) ] if not self.callee.method.is_leaf() else [] diff --git a/sapp/pipeline/model_generator.py b/sapp/pipeline/model_generator.py index 497cb6da..46f16cfd 100644 --- a/sapp/pipeline/model_generator.py +++ b/sapp/pipeline/model_generator.py @@ -8,7 +8,7 @@ import json import logging from collections import defaultdict -from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Union +from typing import Dict, Iterable, List, Optional, Set, Tuple, Union from ..models import ( DBID, @@ -39,6 +39,7 @@ ParseIssueTuple, ParseLeaf, ParseTraceAnnotation, + ParseTraceAnnotationSubtrace, ParseTraceFeature, ParseTypeInterval, PipelineStep, @@ -578,17 +579,18 @@ def _generate_annotation_trace( run: Run, parent_filename: str, parent_caller: str, - trace: Dict[str, Any], + trace: ParseTraceAnnotationSubtrace, annotation: ParseTraceAnnotation, ) -> TraceFrame: # Generates the first-hop trace frames from the annotation and # all dependencies of these sub traces. If this gets called, it is # assumed that the annotation leads to traces, and that the leaf kind # and depth are specified. - callee = trace["callee"] - callee_port = trace["port"] - features = trace.get("features", []) - nested_annotations = trace.get("annotations", []) + callee = trace.callee + callee_port = trace.port + features = trace.features + nested_annotations = trace.annotations + position = trace.position titos = self._generate_tito(parent_filename, annotation, parent_caller) call_tf = self._generate_raw_trace_frame( trace_kind, @@ -598,7 +600,7 @@ def _generate_annotation_trace( "root", callee, callee_port, - annotation.location, + position, titos, [(annotation.leaf_kind or "", annotation.leaf_depth)], annotation.type_interval, diff --git a/sapp/pipeline/pysa_taint_parser.py b/sapp/pipeline/pysa_taint_parser.py index 90666c52..d57e8bb3 100644 --- a/sapp/pipeline/pysa_taint_parser.py +++ b/sapp/pipeline/pysa_taint_parser.py @@ -35,6 +35,7 @@ ParseIssueTuple, ParsePosition, ParseTraceAnnotation, + ParseTraceAnnotationSubtrace, ParseTypeInterval, SourceLocation, ) @@ -444,11 +445,13 @@ def _parse_extra_traces(self, trace: Dict[str, Any]) -> List[ParseTraceAnnotatio if "call" in extra_trace: call = extra_trace["call"] first_hops = [ - { - "callee": resolved, - "port": call["port"], - "position": self._adjust_location(call["position"]), - } + ParseTraceAnnotationSubtrace( + callee=resolved, + port=call["port"], + position=SourceLocation.from_typed_dict( + self._adjust_location(call["position"]) + ), + ) for resolved in call["resolves_to"] ] if len(first_hops) == 0: diff --git a/sapp/pipeline/tests/test_pysa_taint_parser.py b/sapp/pipeline/tests/test_pysa_taint_parser.py index 389bf054..6de9b016 100644 --- a/sapp/pipeline/tests/test_pysa_taint_parser.py +++ b/sapp/pipeline/tests/test_pysa_taint_parser.py @@ -14,6 +14,7 @@ ParseIssueConditionTuple, ParseIssueTuple, ParseTraceAnnotation, + ParseTraceAnnotationSubtrace, ParseTraceFeature, ParseTypeInterval, SourceLocation, @@ -241,15 +242,15 @@ def testIssueV3(self) -> None: trace_key=None, titos=[], subtraces=[ - { - "callee": "extra_trace.transform_yz", - "port": "formal(arg)", - "position": { - "line": 117, - "start": 23, - "end": 24, - }, - } + ParseTraceAnnotationSubtrace( + callee="extra_trace.transform_yz", + port="formal(arg)", + position=SourceLocation( + line_no=117, + begin_column=23, + end_column=24, + ), + ) ], ) ], @@ -288,15 +289,15 @@ def testIssueV3(self) -> None: trace_key=None, titos=[], subtraces=[ - { - "callee": "extra_trace.transform_yz", - "port": "formal(arg)", - "position": { - "line": 117, - "start": 23, - "end": 24, - }, - } + ParseTraceAnnotationSubtrace( + callee="extra_trace.transform_yz", + port="formal(arg)", + position=SourceLocation( + line_no=117, + begin_column=23, + end_column=24, + ), + ) ], ) ], @@ -1593,11 +1594,13 @@ def testSourceModelV3(self) -> None: trace_key=None, titos=[], subtraces=[ - { - "callee": "extra_trace.nested_transform_x", - "port": "formal(arg)", - "position": {"line": 59, "start": 33, "end": 34}, - } + ParseTraceAnnotationSubtrace( + callee="extra_trace.nested_transform_x", + port="formal(arg)", + position=SourceLocation( + line_no=59, begin_column=33, end_column=34 + ), + ) ], ) ], @@ -1637,11 +1640,13 @@ def testSourceModelV3(self) -> None: trace_key=None, titos=[], subtraces=[ - { - "callee": "extra_trace.nested_transform_x", - "port": "formal(arg)", - "position": {"line": 59, "start": 33, "end": 34}, - } + ParseTraceAnnotationSubtrace( + callee="extra_trace.nested_transform_x", + port="formal(arg)", + position=SourceLocation( + line_no=59, begin_column=33, end_column=34 + ), + ) ], ) ], @@ -1681,11 +1686,13 @@ def testSourceModelV3(self) -> None: trace_key=None, titos=[], subtraces=[ - { - "callee": "extra_trace.nested_transform_x", - "port": "formal(arg)", - "position": {"line": 59, "start": 33, "end": 34}, - } + ParseTraceAnnotationSubtrace( + callee="extra_trace.nested_transform_x", + port="formal(arg)", + position=SourceLocation( + line_no=59, begin_column=33, end_column=34 + ), + ), ], ) ], @@ -1725,11 +1732,13 @@ def testSourceModelV3(self) -> None: trace_key=None, titos=[], subtraces=[ - { - "callee": "extra_trace.nested_transform_x", - "port": "formal(arg)", - "position": {"line": 59, "start": 33, "end": 34}, - } + ParseTraceAnnotationSubtrace( + callee="extra_trace.nested_transform_x", + port="formal(arg)", + position=SourceLocation( + line_no=59, begin_column=33, end_column=34 + ), + ) ], ) ], @@ -2253,11 +2262,13 @@ def testSinkModelV3(self) -> None: trace_key=None, titos=[], subtraces=[ - { - "callee": "extra_trace.nested_transform_x", - "port": "formal(arg)", - "position": {"line": 59, "start": 33, "end": 34}, - } + ParseTraceAnnotationSubtrace( + callee="extra_trace.nested_transform_x", + port="formal(arg)", + position=SourceLocation( + line_no=59, begin_column=33, end_column=34 + ), + ) ], ) ],