Add get_root_span helper (#4206)

Author: smola

tests/appsec/api_security/test_schemas.py

@@ -15,12 +15,9 @@
 
 def get_schema(request, address):
     """Get api security schema from spans"""
-    for _, span in interfaces.library.get_root_spans(request):
-        meta = span.get("meta", {})
-        payload = meta.get("_dd.appsec.s." + address)
-        if payload is not None:
-            return payload
-    return None
+    span = interfaces.library.get_root_span(request)
+    meta = span.get("meta", {})
+    return meta.get("_dd.appsec.s." + address)
 
 
 # can be used to match any value in a schema

tests/appsec/iast/utils.py

@@ -17,9 +17,7 @@ def _get_expectation(d):
 
 
 def _get_span_meta(request):
-    spans = [span for _, span in interfaces.library.get_root_spans(request=request)]
-    assert spans, "No root span found"
-    span = spans[0]
+    span = interfaces.library.get_root_span(request)
     meta = span.get("meta", {})
     meta_struct = span.get("meta_struct", {})
     return meta, meta_struct
@@ -199,9 +197,7 @@ def assert_no_iast_event(request, tested_vulnerability_type=None) -> None:
 
 
 def validate_stack_traces(request):
-    spans = [span for _, span in interfaces.library.get_root_spans(request=request)]
-    assert spans, "No root span found"
-    span = spans[0]
+    span = interfaces.library.get_root_span(request)
     meta = span.get("meta", {})
     meta_struct = span.get("meta_struct", {})
     iast = meta.get("_dd.iast.json") or meta_struct.get("iast")
@@ -290,10 +286,7 @@ def validate_stack_traces(request):
 
 
 def validate_extended_location_data(request, vulnerability_type, is_expected_location_required=True):
-    spans = [span for _, span in interfaces.library.get_root_spans(request=request)]
-    assert spans, "No root span found"
-    span = spans[0]
-
+    span = interfaces.library.get_root_span(request)
     iast = span.get("meta", {}).get("_dd.iast.json")
     assert iast, "Expected at least one vulnerability"
     assert iast["vulnerabilities"], "Expected at least one vulnerability"

tests/appsec/rasp/utils.py

@@ -9,17 +9,14 @@
 
 def validate_span_tags(request, expected_meta=(), expected_metrics=()):
     """Validate RASP span tags are added when an event is generated"""
-    spans = [s for _, s in interfaces.library.get_root_spans(request=request)]
-    assert spans, "No spans to validate"
-
-    for span in spans:
-        meta = span["meta"]
-        for m in expected_meta:
-            assert m in meta, f"missing span meta tag `{m}` in {meta}"
-
-        metrics = span["metrics"]
-        for m in expected_metrics:
-            assert m in metrics, f"missing span metric tag `{m}` in {metrics}"
+    span = interfaces.library.get_root_span(request)
+    meta = span["meta"]
+    for m in expected_meta:
+        assert m in meta, f"missing span meta tag `{m}` in {meta}"
+
+    metrics = span["metrics"]
+    for m in expected_metrics:
+        assert m in metrics, f"missing span metric tag `{m}` in {metrics}"
 
 
 def validate_stack_traces(request):

tests/appsec/test_blocking_addresses.py

@@ -695,26 +695,25 @@ def setup_request_block_attack(self):
     @bug(context.library < "ruby@2.10.0-dev", reason="APPSEC-56464")
     def test_request_block_attack(self):
         assert self.r_attack.status_code == 403
-
-        for _, span in interfaces.library.get_root_spans(request=self.r_attack):
-            meta = span.get("meta", {})
-            meta_struct = span.get("meta_struct", {})
-            assert meta["appsec.event"] == "true"
-            assert ("_dd.appsec.json" in meta) ^ ("appsec" in meta_struct)
-            appsec = meta.get("_dd.appsec.json", {}) or meta_struct.get("appsec", {})
-            rule_triggered = appsec["triggers"][0]
-            parameters = rule_triggered["rule_matches"][0]["parameters"][0]
-            assert (
-                parameters["address"] == "graphql.server.all_resolvers"
-                or parameters["address"] == "graphql.server.resolver"
-            )
-            assert rule_triggered["rule"]["id"] == "block-resolvers"
-            assert parameters["key_path"] == (
-                ["userByName", "name"]
-                if parameters["address"] == "graphql.server.resolver"
-                else ["userByName", "0", "name"]
-            )
-            assert parameters["value"] == "testblockresolver"
+        span = interfaces.library.get_root_span(request=self.r_attack)
+        meta = span.get("meta", {})
+        meta_struct = span.get("meta_struct", {})
+        assert meta["appsec.event"] == "true"
+        assert ("_dd.appsec.json" in meta) ^ ("appsec" in meta_struct)
+        appsec = meta.get("_dd.appsec.json", {}) or meta_struct.get("appsec", {})
+        rule_triggered = appsec["triggers"][0]
+        parameters = rule_triggered["rule_matches"][0]["parameters"][0]
+        assert (
+            parameters["address"] == "graphql.server.all_resolvers"
+            or parameters["address"] == "graphql.server.resolver"
+        )
+        assert rule_triggered["rule"]["id"] == "block-resolvers"
+        assert parameters["key_path"] == (
+            ["userByName", "name"]
+            if parameters["address"] == "graphql.server.resolver"
+            else ["userByName", "0", "name"]
+        )
+        assert parameters["value"] == "testblockresolver"
 
     def setup_request_block_attack_directive(self):
         """Currently only monitoring is implemented"""
@@ -734,23 +733,22 @@ def setup_request_block_attack_directive(self):
     @bug(context.library < "ruby@2.10.0-dev", reason="APPSEC-56464")
     def test_request_block_attack_directive(self):
         assert self.r_attack.status_code == 403
-
-        for _, span in interfaces.library.get_root_spans(request=self.r_attack):
-            meta = span.get("meta", {})
-            meta_struct = span.get("meta_struct", {})
-            assert meta["appsec.event"] == "true"
-            assert ("_dd.appsec.json" in meta) ^ ("appsec" in meta_struct)
-            appsec = meta.get("_dd.appsec.json", {}) or meta_struct.get("appsec", {})
-            rule_triggered = appsec["triggers"][0]
-            assert rule_triggered["rule"]["id"] == "block-resolvers"
-            parameters = rule_triggered["rule_matches"][0]["parameters"][0]
-            assert (
-                parameters["address"] == "graphql.server.all_resolvers"
-                or parameters["address"] == "graphql.server.resolver"
-            )
-            assert (
-                parameters["key_path"] == ["userByName", "case", "format"]
-                if parameters["address"] == "graphql.server.resolver"
-                else ["userByName", "0", "case", "format"]
-            )
-            assert parameters["value"] == "testblockresolver"
+        span = interfaces.library.get_root_span(request=self.r_attack)
+        meta = span.get("meta", {})
+        meta_struct = span.get("meta_struct", {})
+        assert meta["appsec.event"] == "true"
+        assert ("_dd.appsec.json" in meta) ^ ("appsec" in meta_struct)
+        appsec = meta.get("_dd.appsec.json", {}) or meta_struct.get("appsec", {})
+        rule_triggered = appsec["triggers"][0]
+        assert rule_triggered["rule"]["id"] == "block-resolvers"
+        parameters = rule_triggered["rule_matches"][0]["parameters"][0]
+        assert (
+            parameters["address"] == "graphql.server.all_resolvers"
+            or parameters["address"] == "graphql.server.resolver"
+        )
+        assert (
+            parameters["key_path"] == ["userByName", "case", "format"]
+            if parameters["address"] == "graphql.server.resolver"
+            else ["userByName", "0", "case", "format"]
+        )
+        assert parameters["value"] == "testblockresolver"

tests/appsec/test_metastruct.py

@@ -14,21 +14,18 @@ def setup_appsec_event_use_metastruct(self):
         self.r = weblog.get("/", headers={"User-Agent": "Arachni/v1"})
 
     def test_appsec_event_use_metastruct(self):
-        spans = [s for _, s in interfaces.library.get_root_spans(request=self.r)]
-        assert spans
+        span = interfaces.library.get_root_span(request=self.r)
+        meta = span.get("meta", {})
+        meta_struct = span.get("meta_struct", {})
+        assert meta["appsec.event"] == "true"
+        assert "_dd.appsec.json" not in meta
+        assert "appsec" in meta_struct
 
-        for span in spans:
-            meta = span.get("meta", {})
-            meta_struct = span.get("meta_struct", {})
-            assert meta["appsec.event"] == "true"
-            assert "_dd.appsec.json" not in meta
-            assert "appsec" in meta_struct
+        # The event is not null
+        assert meta_struct.get("appsec", {}) not in [None, {}]
 
-            # The event is not null
-            assert meta_struct.get("appsec", {}) not in [None, {}]
-
-            # There is at least one rule triggered
-            assert len(meta_struct["appsec"].get("triggers", [])) > 0
+        # There is at least one rule triggered
+        assert len(meta_struct["appsec"].get("triggers", [])) > 0
 
 
 @features.security_events_metastruct
@@ -40,21 +37,18 @@ def setup_iast_event_use_metastruct(self):
         self.r = weblog.get("/set_cookie", params={"name": "metastruct-yes", "value": "yes"})
 
     def test_iast_event_use_metastruct(self):
-        spans = [s for _, s in interfaces.library.get_root_spans(request=self.r)]
-        assert spans
-
-        for span in spans:
-            meta = span.get("meta", {})
-            meta_struct = span.get("meta_struct", {})
-            assert meta["_dd.iast.enabled"] == "1"
-            assert "_dd.iast.json" not in meta
-            assert "iast" in meta_struct
+        span = interfaces.library.get_root_span(request=self.r)
+        meta = span.get("meta", {})
+        meta_struct = span.get("meta_struct", {})
+        assert meta["_dd.iast.enabled"] == "1"
+        assert "_dd.iast.json" not in meta
+        assert "iast" in meta_struct
 
-            # The event is not null
-            assert meta_struct.get("iast", {}) not in [None, {}]
+        # The event is not null
+        assert meta_struct.get("iast", {}) not in [None, {}]
 
-            # There is at least one vulnerability detected
-            assert len(meta_struct["iast"].get("vulnerabilities", [])) > 0
+        # There is at least one vulnerability detected
+        assert len(meta_struct["iast"].get("vulnerabilities", [])) > 0
 
 
 @features.security_events_metastruct
@@ -66,21 +60,18 @@ def setup_appsec_event_fallback_json(self):
         self.r = weblog.get("/", headers={"User-Agent": "Arachni/v1"})
 
     def test_appsec_event_fallback_json(self):
-        spans = [s for _, s in interfaces.library.get_root_spans(request=self.r)]
-        assert spans
-
-        for span in spans:
-            meta = span.get("meta", {})
-            meta_struct = span.get("meta_struct", {})
-            assert meta["appsec.event"] == "true"
-            assert "_dd.appsec.json" in meta
-            assert "appsec" not in meta_struct
+        span = interfaces.library.get_root_span(request=self.r)
+        meta = span.get("meta", {})
+        meta_struct = span.get("meta_struct", {})
+        assert meta["appsec.event"] == "true"
+        assert "_dd.appsec.json" in meta
+        assert "appsec" not in meta_struct
 
-            # The event is not null
-            assert meta.get("_dd.appsec.json", {}) not in [None, {}]
+        # The event is not null
+        assert meta.get("_dd.appsec.json", {}) not in [None, {}]
 
-            # There is at least one rule triggered
-            assert len(meta["_dd.appsec.json"].get("triggers", [])) > 0
+        # There is at least one rule triggered
+        assert len(meta["_dd.appsec.json"].get("triggers", [])) > 0
 
 
 @features.security_events_metastruct
@@ -93,18 +84,15 @@ def setup_iast_event_fallback_json(self):
         self.r = weblog.get("/set_cookie", params={"name": "metastruct-no", "value": "no"})
 
     def test_iast_event_fallback_json(self):
-        spans = [s for _, s in interfaces.library.get_root_spans(request=self.r)]
-        assert spans
-
-        for span in spans:
-            meta = span.get("meta", {})
-            meta_struct = span.get("meta_struct", {})
-            assert meta["_dd.iast.enabled"] == "1"
-            assert "_dd.iast.json" in meta
-            assert "iast" not in meta_struct
-
-            # The event is not null
-            assert meta.get("_dd.iast.json", {}) not in [None, {}]
-
-            # There is at least one vulnerability detected
-            assert len(meta["_dd.iast.json"].get("vulnerabilities", [])) > 0
+        span = interfaces.library.get_root_span(request=self.r)
+        meta = span.get("meta", {})
+        meta_struct = span.get("meta_struct", {})
+        assert meta["_dd.iast.enabled"] == "1"
+        assert "_dd.iast.json" in meta
+        assert "iast" not in meta_struct
+
+        # The event is not null
+        assert meta.get("_dd.iast.json", {}) not in [None, {}]
+
+        # There is at least one vulnerability detected
+        assert len(meta["_dd.iast.json"].get("vulnerabilities", [])) > 0

tests/appsec/test_user_blocking_full_denylist.py

@@ -43,9 +43,7 @@ def test_blocking_test(self):
         for r in self.r_blocked_requests:
             assert r.status_code == 403
             interfaces.library.assert_waf_attack(r, rule="blk-001-002", address="usr.id")
-            spans = [s for _, s in interfaces.library.get_root_spans(r)]
-            assert len(spans) == 1
-            span = spans[0]
+            span = interfaces.library.get_root_span(r)
             assert span["meta"]["appsec.event"] == "true"
             assert span["meta"]["appsec.blocked"] == "true"
             assert span["meta"]["http.status_code"] == "403"

tests/external_processing/test_apm.py

@@ -9,10 +9,9 @@ def setup_correct_span_structure(self):
 
     def test_correct_span_structure(self):
         assert self.r.status_code == 200
-
         interfaces.library.assert_trace_exists(self.r)
-        for _, span in interfaces.library.get_root_spans(request=self.r):
-            assert span["type"] == "web"
-            assert span["meta"]["span.kind"] == "server"
-            assert span["meta"]["http.url"] == "http://localhost:7777/"
-            assert span["meta"]["http.host"] == "localhost:7777"
+        span = interfaces.library.get_root_span(self.r)
+        assert span["type"] == "web"
+        assert span["meta"]["span.kind"] == "server"
+        assert span["meta"]["http.url"] == "http://localhost:7777/"
+        assert span["meta"]["http.host"] == "localhost:7777"

tests/test_span_events.py

@@ -22,10 +22,8 @@ def setup_v04_v07_default_format(self):
     def test_v04_v07_default_format(self):
         """For traces that default to the v0.4 or v0.7 format, send events as a top-level `span_events` field"""
         interfaces.library.assert_trace_exists(self.r)
-
-        span = self._get_span(self.r)
-        meta = self._get_root_span_meta(self.r)
-
+        span = interfaces.library.get_root_span(self.r)
+        meta = span.get("meta", {})
         assert "span_events" in span
         assert "events" not in meta
 
@@ -38,21 +36,11 @@ def test_v05_default_format(self):
         given this format does not support native serialization.
         """
         interfaces.library.assert_trace_exists(self.r)
-
-        span = self._get_span(self.r)
-        meta = self._get_root_span_meta(self.r)
-
+        span = interfaces.library.get_root_span(self.r)
+        meta = span.get("meta", {})
         assert "span_events" not in span
         assert "events" in meta
 
-    def _get_root_span_meta(self, request):
-        return self._get_span(request).get("meta", {})
-
-    def _get_span(self, request):
-        root_spans = [s for _, s in interfaces.library.get_root_spans(request=request)]
-        assert len(root_spans) == 1
-        return root_spans[0]
-
 
 @features.span_events
 @scenarios.agent_not_supporting_span_events
@@ -69,17 +57,7 @@ def setup_send_as_a_tag(self):
     def test_send_as_a_tag(self):
         """Send span events as the tag `events` when the agent does not support native serialization"""
         interfaces.library.assert_trace_exists(self.r)
-
-        span = self._get_span(self.r)
-        meta = self._get_root_span_meta(self.r)
-
+        span = interfaces.library.get_root_span(self.r)
+        meta = span.get("meta", {})
         assert "span_events" not in span
         assert "events" in meta
-
-    def _get_root_span_meta(self, request):
-        return self._get_span(request).get("meta", {})
-
-    def _get_span(self, request):
-        root_spans = [s for _, s in interfaces.library.get_root_spans(request=request)]
-        assert len(root_spans) == 1
-        return root_spans[0]

tests/test_standard_tags.py

@@ -338,7 +338,5 @@ def test_client_ip_with_appsec_event_and_vendor_headers(self):
             assert meta[tag] == value
 
     def _get_root_span_meta(self, request):
-        root_spans = [s for _, s in interfaces.library.get_root_spans(request=request)]
-        assert len(root_spans) == 1
-        span = root_spans[0]
+        span = interfaces.library.get_root_span(request)
         return span.get("meta", {})