add capec_id and attack_id filters #63

muchdogesec · Dec 11, 2024 · 997301e · 997301e
1 parent f47ac45
commit 997301e
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 5 deletions.
diff --git a/vulmatch/server/arango_helpers.py b/vulmatch/server/arango_helpers.py
@@ -344,7 +344,19 @@ def get_vulnerabilities(self):
         if q := self.query_as_array('weakness_id'):
             binds['weakness_ids'] = q
             filters.append('''
-                FILTER LENGTH(FOR d IN nvd_cve_edge_collection FILTER doc._id == d._from AND d.relationship_type == 'exploited-using' AND LAST(SPLIT(d.description, ' ')) IN @weakness_ids LIMIT 1 RETURN TRUE) > 0
+                FILTER doc.external_references[? ANY FILTER CURRENT.source_name=='cwe' AND CURRENT.external_id IN @weakness_ids]
+                ''')
+
+        if q := self.query_as_array('attack_id'):
+            binds['attack_ids'] = q
+            filters.append('''
+                FILTER LENGTH(FOR d IN nvd_cve_edge_collection FILTER doc._id == d._from AND d.relationship_type == 'exploited-using' AND d._arango_cve_processor_note == "cve-attack" AND LAST(SPLIT(d.description, ' ')) IN @attack_ids LIMIT 1 RETURN TRUE) > 0
+                ''')
+
+        if q := self.query_as_array('capec_id'):
+            binds['capec_ids'] = q
+            filters.append('''
+                FILTER LENGTH(FOR d IN nvd_cve_edge_collection FILTER doc._id == d._from AND d.relationship_type == 'exploited-using' AND d._arango_cve_processor_note == "cve-capec" AND LAST(SPLIT(d.description, ' ')) IN @capec_ids LIMIT 1 RETURN TRUE) > 0
                 ''')
 
         query = """
@@ -367,7 +379,7 @@ def get_vulnerabilities(self):
                 },
             ),
         )
-        # return Response(query)
+        # return Response([query, binds])
         return self.execute_query(query, bind_vars=binds)
 
     def get_cve_bundle(self, cve_id: str):

diff --git a/vulmatch/server/views.py b/vulmatch/server/views.py
@@ -204,17 +204,17 @@ class filterset_class(FilterSet):
             Is the maximum `created` value (`YYYY-MM-DDThh:mm:ss.sssZ`)
             """
         ))
-        modified_min = DateTimeFilter(label=textwrap.dedent(
+        modified_min = DateTimeFilter(help_text=textwrap.dedent(
             """
             Is the minimum `modified` value (`YYYY-MM-DDThh:mm:ss.sssZ`)
             """
         ))
-        modified_max = DateTimeFilter(label=textwrap.dedent(
+        modified_max = DateTimeFilter(help_text=textwrap.dedent(
             """
             Is the maximum `modified` value (`YYYY-MM-DDThh:mm:ss.sssZ`)
             """
         ))
-        sort = ChoiceFilter(choices=[(v, v) for v in CVE_SORT_FIELDS], label=textwrap.dedent(
+        sort = ChoiceFilter(choices=[(v, v) for v in CVE_SORT_FIELDS], help_text=textwrap.dedent(
             """
             Sort results by
             """
@@ -225,6 +225,19 @@ class filterset_class(FilterSet):
             """
         ))
 
+        attack_id = BaseCSVFilter(help_text=textwrap.dedent(
+            """
+            Filter results by weakness (ATT&CK ID). e.g. `T1223`.\n\n
+            filters using the `description` property of `cve-attack` relationship object
+            """
+        ))
+        capec_id = BaseCSVFilter(help_text=textwrap.dedent(
+            """
+            Filter results by weakness (CAPEC ID). e.g. `CAPEC-665`.\n\n
+            filters using the `description` property of `cve-capec` relationship object
+            """
+        ))
+
 
     def create(self, request, *args, **kwargs):
         serializer = serializers.NVDTaskSerializer(data=request.data)