sort by first values #439; add kb_type #442 (#448)

Author: fqrious

obstracts/server/models.py

@@ -377,7 +377,7 @@ class ObjectValue(models.Model):
     """
     stix_id = models.CharField(max_length=256, db_index=True)
     type = models.CharField(max_length=256, db_index=True)
-    ttp_type = models.CharField(max_length=64, null=True, blank=True, db_index=True)
+    knowledgebase = models.CharField(max_length=64, null=True, blank=True, db_index=True)
     values = models.JSONField()
     file = models.ForeignKey(File, on_delete=models.CASCADE, related_name='object_values')
     created = models.DateTimeField(default=None, null=True)
@@ -390,7 +390,7 @@ class Meta:
         unique_together = [['stix_id', 'file']]
 
     def __str__(self):
-        return f'ObjectValue(stix_id={self.stix_id}, ttp_type={self.ttp_type})'
+        return f'ObjectValue(stix_id={self.stix_id}, knowledgebase={self.knowledgebase})'
 
 
 class Job(models.Model):

obstracts/server/values/filters.py

@@ -1,5 +1,7 @@
 
 from django.db.models import JSONField, Lookup
+from django.db.models import Func, CharField
+
 
 
 @JSONField.register_lookup
@@ -68,4 +70,13 @@ def as_sql(self, compiler, connection):
 
         # Use jsonb_each_text to iterate over key-value pairs and check if any value matches exactly
         sql = f"EXISTS (SELECT 1 FROM jsonb_each_text({lhs}) WHERE value = %s)"
-        return sql, lhs_params + rhs_params
+        return sql, lhs_params + rhs_params
+
+class NormalizeDict(Func):
+    """
+    Returns the value for the first key (alphabetically) from a JSONB object.
+    """
+
+    template = "(SELECT kv.value FROM jsonb_each_text(%(expressions)s) AS kv ORDER BY kv.key LIMIT 1)"
+    output_field = CharField()
+

obstracts/server/values/serializers.py

@@ -6,7 +6,7 @@ class ObjectValueSerializer(serializers.Serializer):
 
     id = serializers.CharField(source='stix_id')
     type = serializers.CharField()
-    ttp_type = serializers.CharField(required=False)
+    kb_name = serializers.CharField(read_only=True, required=False, source='knowledgebase')
     values = serializers.JSONField(read_only=True)
     matched_posts = serializers.ListField(child=serializers.UUIDField())
     created = serializers.DateTimeField(required=False)

obstracts/server/values/values.py

@@ -20,6 +20,27 @@ def external_id(obj):
 def hashes(obj):
     return obj.get("hashes", {})
 
+KB_TYPES = {
+    "Tactic": [dict(type='x-mitre-tactic')],
+    "Analytic": [dict(type='x-mitre-analytic')],
+    "Detection Strategy": [dict(type='x-mitre-detection-strategy')],
+    "Technique": [dict(type='attack-pattern', x_mitre_is_subtechnique=False), dict(type='attack-pattern', x_mitre_is_subtechnique=None)],
+    "Sub-technique": [dict(type='attack-pattern', x_mitre_is_subtechnique=True)],
+    "Mitigation": [dict(type='course-of-action')],
+    "Group": [dict(type='intrusion-set')],
+    "Software": [dict(type='malware'), dict(type='tool')],
+    "Campaign": [dict(type='campaign')],
+    "Data Source": [dict(type='x-mitre-data-source')],
+    "Data Component": [dict(type='x-mitre-data-component')],
+    "Asset": [dict(type='x-mitre-asset')],
+}
+
+def get_kb_type(obj):
+    for form, criteria_list in KB_TYPES.items():
+        for criteria in criteria_list:
+            if all(obj.get(k) == v for k, v in criteria.items()):
+                return form
+    return None
 
 def get_file_values(obj):
     values = {}
@@ -133,36 +154,36 @@ def get_values(obj: dict, value_keys: list[str] | dict[str, str] | Callable):
 }
 
 
-def get_ttp_type(obj: dict) -> str | None:
+def guess_kb_data(obj: dict) -> str | None:
     """
-    Determine the TTP type of a STIX object based on its properties.
+    Determine the KnowledgeBase type of a STIX object based on its properties.
 
     Returns:
         - "cve" for vulnerability objects
         - "cwe" for weakness objects
         - "location" for location objects
         - "enterprise-attack", "mobile-attack", "ics-attack" based on x_mitre_domains
         - "capec", "atlas", "disarm", "sector" based on external_references source_name
-        - None if not a TTP object
+        - None if not a KnowledgeBase object
     """
     obj_type = obj["type"]
-    ttp_type = None
+    kb_name = None
     extra = {}
 
     # Check for CVE (vulnerability)
     match obj_type:
         case "vulnerability":
-            ttp_type = "cve"
+            kb_name = "cve"
         case "weakness":
-            ttp_type = "cwe"
+            kb_name = "cwe"
         case "location":
-            ttp_type = "location"
+            kb_name = "location"
     # Check for MITRE ATT&CK domains
     x_mitre_domains = obj.get("x_mitre_domains", [])
     if x_mitre_domains:
         domain = x_mitre_domains[0]
         if domain in ["enterprise-attack", "mobile-attack", "ics-attack"]:
-            ttp_type = domain
+            kb_name = domain
 
     # Check external references for other TTP types
     external_refs = obj.get("external_references", [])
@@ -176,10 +197,12 @@ def get_ttp_type(obj: dict) -> str | None:
             "sector2stix": "sector",
         }
         if source_name in ttp_source_name_mapping:
-            ttp_type = ttp_source_name_mapping[source_name]
-    if ttp_type and (ttp_ids := external_id(obj)):
-        extra["ttp_id"] = ttp_ids[0]
-    return ttp_type, extra
+            kb_name = ttp_source_name_mapping[source_name]
+    if kb_name and (kb_ids := external_id(obj)):
+        extra["kb_id"] = kb_ids[0]
+    if kb_name and (kb_type := get_kb_type(obj)):
+        extra["kb_type"] = kb_type    
+    return kb_name, extra
 
 
 def get_visibility(obj: dict) -> str:
@@ -209,12 +232,12 @@ def extract_object_metadata(obj: dict) -> dict:
         A dictionary containing:
         - id: The STIX object ID
         - type: The STIX object type
-        - ttp_type: The TTP type if applicable (None otherwise)
+        - knowledgebase: The source KnowledgeBase type if applicable (None otherwise)
         - values: The extracted values based on the object type
     """
     obj_id = obj["id"]
     obj_type = obj["type"]
-    ttp_type, ttp_extra = get_ttp_type(obj)
+    kb_name, kb_extra = guess_kb_data(obj)
 
     # Get the value configuration for this object type
     type_config = type_value_map.get(obj_type, {})
@@ -223,11 +246,11 @@ def extract_object_metadata(obj: dict) -> dict:
     # Extract values using get_values function
     values = get_values(obj, value_keys) or {}
 
-    values.update(ttp_extra)
+    values.update(kb_extra)
     return {
         "stix_id": obj_id,
         "type": obj_type,
-        "ttp_type": ttp_type,
+        "knowledgebase": kb_name,
         "values": values,
         "modified": obj.get("modified"),
         "created": obj.get("created"),

obstracts/server/values/views.py

@@ -1,29 +1,26 @@
 import textwrap
-from rest_framework import viewsets, filters, mixins
-from rest_framework.response import Response
+from rest_framework import viewsets, mixins
 from django_filters.rest_framework import (
     DjangoFilterBackend,
     FilterSet,
     CharFilter,
-    MultipleChoiceFilter,
     BooleanFilter,
     BaseCSVFilter,
 )
 from django_filters.fields import ChoiceField
 from obstracts.server import autoschema as api_schema
 
-from drf_spectacular.utils import extend_schema, extend_schema_view, OpenApiParameter
-from drf_spectacular.types import OpenApiTypes
-from django.db.models import F, Value, JSONField as DjangoJSONField, Min, Max
-from django.db.models.functions import JSONObject
+from drf_spectacular.utils import extend_schema, extend_schema_view
+from django.db.models import F, Min, Max
 from django.contrib.postgres.aggregates import ArrayAgg
 
 
 from obstracts.server.models import ObjectValue
 from obstracts.server.utils import Pagination
-from obstracts.server.values.values import sco_value_map, sdo_value_map
+from obstracts.server.values.values import sco_value_map, sdo_value_map, KB_TYPES
 from .serializers import ObjectValueSerializer
 from dogesec_commons.utils.ordering import Ordering
+from .filters import NormalizeDict
 
 TTP_TYPES = [
     "cve",
@@ -111,7 +108,7 @@ class BaseObjectValueView(mixins.ListModelMixin, viewsets.GenericViewSet):
     pagination_class = Pagination("values")
     filter_backends = [DjangoFilterBackend, Ordering]
     filterset_class = ObjectValueFilterSet
-    ordering_fields = ["stix_id", "type", "ttp_type"]
+    ordering_fields = ["stix_id", "type", "knowledgebase", "value"]
     ordering = "stix_id_descending"
     openapi_tags = ["Object Values"]
 
@@ -128,11 +125,12 @@ def get_queryset(self):
         # Aggregate all post_ids for each unique stix_id
         queryset = queryset.values("stix_id").annotate(
             type=F("type"),
-            ttp_type=F("ttp_type"),
+            knowledgebase=F("knowledgebase"),
             values=F("values"),
             matched_posts=ArrayAgg("file__post_id", distinct=True),
             created=Min("created"),
             modified=Max("modified"),
+            value=NormalizeDict(F("values")),
         )
 
         return queryset
@@ -175,8 +173,8 @@ class SCOValueView(BaseObjectValueView):
     """View for STIX Cyber Observable Objects (SCOs) only."""
 
     allowed_types = list(sco_value_map.keys())
-    ordering_fields = ["values", "stix_id", "type"]
-    ordering = "values_ascending"
+    ordering_fields = ["value", "stix_id", "type"]
+    ordering = "value_ascending"
 
     class filterset_class(ObjectValueFilterSet):
         types = ChoiceCSVFilter(
@@ -225,11 +223,11 @@ class SDOValueView(BaseObjectValueView):
     """View for STIX Domain Objects (SDOs) only."""
 
     allowed_types = list(sdo_value_map.keys())
-    ordering_fields = ["stix_id", "type", "ttp_type", "values", "created", "modified"]
+    ordering_fields = ["stix_id", "type", "knowledgebase", "value", "created", "modified"]
 
     class filterset_class(ObjectValueFilterSet):
-        ttp_types = ChoiceCSVFilter(
-            field_name="ttp_type",
+        knowledgebases = ChoiceCSVFilter(
+            field_name="knowledgebase",
             help_text="Filter results by source of TTP object (cve, cwe, enterprise-attack, mobile-attack, ics-attack, capec, location, disarm, atlas, sector)",
             choices=[(c, c) for c in TTP_TYPES],
         )
@@ -238,3 +236,9 @@ class filterset_class(ObjectValueFilterSet):
             help_text="Filter the results by one or more STIX Domain Object types",
             choices=[(c, c) for c in sdo_value_map.keys()],
         )
+
+        kb_type = ChoiceCSVFilter(
+            field_name="values__kb_type",
+            help_text="Filter results by knowledge base type.",
+            choices=[(c, c) for c in KB_TYPES.keys()],
+        )