implement bundle filters #36 (#54)

muchdogesec · Nov 12, 2024 · ab8514b · ab8514b
1 parent 60f4a0a
commit ab8514b
Show file tree

Hide file tree

Showing 2 changed files with 79 additions and 17 deletions.
diff --git a/vulmatch/server/arango_helpers.py b/vulmatch/server/arango_helpers.py
@@ -118,6 +118,16 @@
 
 CPE_RELATIONSHIP_TYPES = {"vulnerable-to": "is-vulnerable", "in-pattern": "pattern-contains"}
 CPE_REL_SORT_FIELDS = ["modified_descending", "modified_ascending", "created_descending", "created_ascending"]
+CVE_BUNDLE_TYPES = set([
+  "vulnerability",
+  "indicator",
+  "relationship",
+  "note",
+  "sighting",
+  "software",
+  "weakness",
+  "attack-pattern"
+])
 
 class ArangoDBHelper:
     max_page_size = settings.MAXIMUM_PAGE_SIZE
@@ -338,6 +348,56 @@ def get_vulnerabilities(self):
         return self.execute_query(query, bind_vars=binds)
 
     def get_cve_bundle(self, cve_id: str):
+        cve_rels_types = []
+        binds = dict(cve_id=cve_id.upper(), cve_edge_types=cve_rels_types)
+
+        more_queries = {}
+
+        include_attack = self.query_as_bool('include_attack', True)
+        include_capec = self.query_as_bool('include_capec', True) or include_attack
+        include_cwe = self.query_as_bool('include_cwe', True) or include_capec
+
+
+
+
+        if include_capec:
+            include_cwe = True
+            more_queries['cwe_capec'] = """
+            LET cwe_capec = FLATTEN(
+                FOR doc IN mitre_cwe_edge_collection
+                FILTER doc.relationship_type == 'exploited-using' AND [doc._from, doc._to] ANY IN cve_rels[*]._id
+                RETURN [doc, DOCUMENT(doc._from), DOCUMENT(doc._to)]
+            )
+            """
+
+
+        if include_attack:
+            include_capec = True
+            more_queries["capec_attack"] = """
+            LET capec_attack = FLATTEN(
+                FOR doc IN mitre_capec_edge_collection
+                FILTER doc.relationship_type == 'technique' AND [doc._from, doc._to] ANY IN cwe_capec[*]._id
+                RETURN [doc, DOCUMENT(doc._from), DOCUMENT(doc._to)]
+            )
+            """
+
+        if self.query_as_bool('include_cpe', True):
+            cve_rels_types.append('pattern-contains')
+        if self.query_as_bool('include_cpe_vulnerable', True):
+            cve_rels_types.append('is-vulnerable')
+
+        if include_cwe:
+            cve_rels_types.append('exploited-using')
+
+        if self.query_as_bool('include_epss', True):
+            cve_rels_types.append('object')
+        if self.query_as_bool('include_kev', True):
+            cve_rels_types.append('sighting-of')
+
+        types = self.query_as_array('object_type') or CVE_BUNDLE_TYPES
+        binds['types'] = list(CVE_BUNDLE_TYPES.intersection(types))
+
+
         query = '''
 LET cve_data = (
   FOR doc IN nvd_cve_vertex_collection
@@ -346,30 +406,23 @@ def get_cve_bundle(self, cve_id: str):
 )
 LET cve_rels = FLATTEN(
     FOR doc IN nvd_cve_edge_collection
-    FILTER [doc._from, doc._to] ANY IN cve_data[*]._id
-    RETURN [doc, DOCUMENT(doc._from), DOCUMENT(doc._to)]
-    )
-    
-LET cwe_capec = FLATTEN(
-    FOR doc IN mitre_cwe_edge_collection
-    FILTER [doc._from, doc._to] ANY IN cve_rels[*]._id
+    FILTER [doc._from, doc._to] ANY IN cve_data[*]._id AND doc.relationship_type IN @cve_edge_types
+
     RETURN [doc, DOCUMENT(doc._from), DOCUMENT(doc._to)]
     )
     
-LET capec_attack = FLATTEN(
-    FOR doc IN mitre_capec_edge_collection
-    FILTER [doc._from, doc._to] ANY IN cwe_capec[*]._id
-    RETURN [doc, DOCUMENT(doc._from), DOCUMENT(doc._to)]
-    )
+@@@more_queries
 
     
-FOR d in UNION_DISTINCT(cve_data, cve_rels, cwe_capec, capec_attack)
-
+FOR d in UNION_DISTINCT(cve_data, cve_rels, @@@extra_rels)
+FILTER d.type IN @types
 LIMIT @offset, @count
 //RETURN KEEP(d, 'id', '_stix2arango_note', '_arango_cti_processor_note', 'description')
 RETURN KEEP(d, KEYS(d, TRUE))
 '''
-        return self.execute_query(query, bind_vars=dict(cve_id=cve_id.upper()))
+        query = query.replace('@@@more_queries', "\n".join(more_queries.values())) \
+                    .replace("@@@extra_rels", ", ".join(more_queries.keys()) or '[]')
+        return self.execute_query(query, bind_vars=binds)
 
     def get_attack_objects(self, matrix):
         filters = []

diff --git a/vulmatch/server/views.py b/vulmatch/server/views.py
@@ -2,7 +2,7 @@
 from django.shortcuts import render
 from rest_framework import viewsets, filters, status, decorators
 
-from vulmatch.server.arango_helpers import ATLAS_TYPES, CPE_REL_SORT_FIELDS, CPE_RELATIONSHIP_TYPES, CVE_SORT_FIELDS, LOCATION_TYPES, TLP_TYPES, ArangoDBHelper, ATTACK_TYPES, CWE_TYPES, SOFTWARE_TYPES, CAPEC_TYPES
+from vulmatch.server.arango_helpers import ATLAS_TYPES, CPE_REL_SORT_FIELDS, CPE_RELATIONSHIP_TYPES, CVE_BUNDLE_TYPES, CVE_SORT_FIELDS, LOCATION_TYPES, TLP_TYPES, ArangoDBHelper, ATTACK_TYPES, CWE_TYPES, SOFTWARE_TYPES, CAPEC_TYPES
 from vulmatch.server.autoschema import DEFAULT_400_ERROR
 from vulmatch.server.utils import Pagination, Response, Ordering, split_mitre_version
 from vulmatch.worker.tasks import new_task
@@ -105,7 +105,16 @@ class VulnerabilityStatus(models.models.TextChoices):
             """
         ),
         responses={200: ArangoDBHelper.get_paginated_response_schema('objects', 'vulnerability')},
-        parameters=ArangoDBHelper.get_schema_operation_parameters(),
+        parameters=ArangoDBHelper.get_schema_operation_parameters() + [
+            OpenApiParameter('object_type', description="The type of STIX object to be returned", enum=CVE_BUNDLE_TYPES, many=True, explode=False),
+            OpenApiParameter('include_cpe', description="will show all `software` objects related to this vulnerability (and the SROS linking cve-cpe)", type=OpenApiTypes.BOOL),
+            OpenApiParameter('include_cpe_vulnerable', description="will show `software` objects vulnerable to this vulnerability (and the SROS), if exist. Note `include_cpe` should be set to `false` if you only want to see vulnerable cpes (and the SROS linking cve-cpe)", type=OpenApiTypes.BOOL),
+            OpenApiParameter('include_cwe', description="will show `weakness` objects related to this vulnerability, if exist (and the SROS linking cve-cwe)", type=OpenApiTypes.BOOL),
+            OpenApiParameter('include_epss', description="will show `note` objects related to this vulnerability, if exist", type=OpenApiTypes.BOOL),
+            OpenApiParameter('include_kev', description="will show `sighthing` objects related to this vulnerability, if exist (and the SROS linking cve-sighting)", type=OpenApiTypes.BOOL),
+            OpenApiParameter('include_capec', description="will show CAPEC `attack-pattern` objects related to this vulnerability, if exist  (and the SROS linking cwe-capec)\n * note this mode will also show `include_cwe` outputs, due to the way CAPEC is linked to CVE", type=OpenApiTypes.BOOL),
+            OpenApiParameter('include_attack', description="will show ATT&CK `attack-pattern` objects (for Techniques/Sub-techniques) related to this vulnerability, if exist (and the SROS linking capec-attack)\n * note this mode will also show `include_capec` and `include_cwe` outputs, due to the way ATT&CK is linked to CVE", type=OpenApiTypes.BOOL),
+        ],
     ),
     versions=extend_schema(
         responses=serializers.StixVersionsSerializer,