Extended observables api (#949)

Co-authored-by: Thomas Chopitea <[email protected]>

Author: udgover

core/schemas/observable.py

@@ -3,7 +3,7 @@
 import datetime
 import re
 from enum import Enum
-from typing import ClassVar, Literal
+from typing import ClassVar, Literal, Type
 
 import validators
 from core import database_arango
@@ -53,16 +53,19 @@ class Observable(BaseModel, database_arango.ArangoYetiConnector):
     context: list[dict] = []
     last_analysis: dict[str, datetime.datetime] = {}
 
+
     @classmethod
-    def load(cls, object: dict) -> "Observable":
-        return cls(**object)
+    def load(cls, object: dict) -> "ObservableTypes":
+        if object["type"] in TYPE_MAPPING:
+            return TYPE_MAPPING[object["type"]](**object)
+        raise ValueError("Attempted to instantiate an undefined observable type.")
 
     @classmethod
     def is_valid(cls, object: dict) -> bool:
         return validate_observable(object)
 
     @classmethod
-    def add_text(cls, text: str, tags: list[str] = []) -> "Observable":
+    def add_text(cls, text: str, tags: list[str] = []) -> "ObservableTypes":
         """Adds and returns an observable for a given string.
 
         Args:
@@ -79,9 +82,8 @@ def add_text(cls, text: str, tags: list[str] = []) -> "Observable":
 
         observable = Observable.find(value=refanged)
         if not observable:
-            observable = Observable(
+            observable = TYPE_MAPPING[observable_type](
                 value=refanged,
-                type=observable_type,
                 created=datetime.datetime.now(datetime.timezone.utc),
             ).save()
         if tags:
@@ -90,7 +92,7 @@ def add_text(cls, text: str, tags: list[str] = []) -> "Observable":
 
     def add_context(
         self, source: str, context: dict, skip_compare: set = set()
-    ) -> "Observable":
+    ) -> "ObservableTypes":
         """Adds context to an observable."""
         compare_fields = set(context.keys()) - skip_compare - {"source"}
         for idx, db_context in enumerate(list(self.context)):
@@ -111,7 +113,7 @@ def add_context(
 
     def delete_context(
         self, source: str, context: dict, skip_compare: set = set()
-    ) -> "Observable":
+    ) -> "ObservableTypes":
         """Deletes context from an observable."""
         compare_fields = set(context.keys()) - skip_compare - {"source"}
         for idx, db_context in enumerate(list(self.context)):
@@ -140,9 +142,9 @@ def delete_context(
 
 REGEXES_OBSERVABLES = {
     # Unix
-    ObservableType.path : [
+    ObservableType.path: [
         re.compile(r"^(\/[^\/\0]+)+$"),
-        re.compile(r"^(?:[a-zA-Z]\:|\\\\[\w\.]+\\[\w.$]+)\\(?:[\w]+\\)*\w([\w.])+")
+        re.compile(r"^(?:[a-zA-Z]\:|\\\\[\w\.]+\\[\w.$]+)\\(?:[\w]+\\)*\w([\w.])+"),
     ]
 }
 
@@ -170,16 +172,14 @@ def find_type(value: str) -> ObservableType | None:
     return None
 
 
-TYPE_MAPPING = {
-    'observable': Observable,
-    'observables': Observable
-}
+TYPE_MAPPING = {"observable": Observable, "observables": Observable}
+
 
 # Import all observable types, as these register themselves in the TYPE_MAPPING
 # disable: pylint=wrong-import-position
 from core.schemas.observables import (asn, bitcoin_wallet, certificate, cidr,
                                       command_line, docker_image, email, file,
-                                      generic_observable, hostname, imphash, ipv4, ipv6,
-                                      mac_address, md5, path, registry_key,
-                                      sha1, sha256, ssdeep, tlsh, url,
-                                      user_agent)
+                                      generic_observable, hostname, imphash,
+                                      ipv4, ipv6, mac_address, md5, path,
+                                      registry_key, sha1, sha256, ssdeep, tlsh,
+                                      url, user_agent)

core/web/apiv2/observables.py

@@ -1,20 +1,26 @@
 import datetime
 from typing import Iterable
 
+from core.schemas import graph
+from core.schemas.observable import TYPE_MAPPING, Observable, ObservableType
 from fastapi import APIRouter, HTTPException
 from pydantic import BaseModel, ConfigDict, field_validator
 
-from core.schemas import graph
-from core.schemas.observable import Observable, ObservableType
+ObservableTypes = ()
 
+for key in TYPE_MAPPING:
+    if key in ["observable", "observables"]:
+        continue
+    cls = TYPE_MAPPING[key]
+    if not ObservableTypes:
+        ObservableTypes = cls
+    else:
+        ObservableTypes |= cls
 
-# Request schemas
-class NewObservableRequest(BaseModel):
-    model_config = ConfigDict(extra='forbid')
 
-    value: str
+class TagRequestMixin(BaseModel):
+
     tags: list[str] = []
-    type: ObservableType
 
     @field_validator("tags")
     @classmethod
@@ -25,17 +31,30 @@ def validate_tags(cls, value) -> list[str]:
         return value
 
 
+# Request schemas
+class NewObservableRequest(TagRequestMixin):
+    model_config = ConfigDict(extra='forbid')
+
+    value: str
+    type: ObservableType
+
+
+class NewExtendedObservableRequest(TagRequestMixin):
+    model_config = ConfigDict(extra='forbid')
+
+    observable: ObservableTypes
+
+
 class NewBulkObservableAddRequest(BaseModel):
     model_config = ConfigDict(extra='forbid')
 
     observables: list[NewObservableRequest]
 
 
-class AddTextRequest(BaseModel):
+class AddTextRequest(TagRequestMixin):
     model_config = ConfigDict(extra='forbid')
 
     text: str
-    tags: list[str] = []
 
 
 class AddContextRequest(BaseModel):
@@ -62,15 +81,14 @@ class ObservableSearchRequest(BaseModel):
 class ObservableSearchResponse(BaseModel):
     model_config = ConfigDict(extra='forbid')
 
-    observables: list[Observable]
+    observables: list[ObservableTypes]
     total: int
 
 
-class ObservableTagRequest(BaseModel):
+class ObservableTagRequest(TagRequestMixin):
     model_config = ConfigDict(extra='forbid')
 
     ids: list[str]
-    tags: list[str]
     strict: bool = False
 
 class ObservableTagResponse(BaseModel):
@@ -90,7 +108,7 @@ async def observables_root() -> Iterable[Observable]:
 
 
 @router.post("/")
-async def new(request: NewObservableRequest) -> Observable:
+async def new(request: NewObservableRequest) -> ObservableTypes:
     """Creates a new observable in the database.
 
     Raises:
@@ -113,8 +131,28 @@ async def new(request: NewObservableRequest) -> Observable:
     return new
 
 
+@router.post("/extended")
+async def new(request: NewExtendedObservableRequest) -> ObservableTypes:
+    """Creates a new observable in the database with extended properties.
+
+    Raises:
+        HTTPException(400) if observable already exists.
+    """
+    observable = Observable.find(value=request.observable.value, type=request.observable.type)
+    if observable:
+        raise HTTPException(
+            status_code=400,
+            detail=f"Observable with value {request.observable.value} already exists",
+        )
+    cls = TYPE_MAPPING[request.observable.type]
+    new = cls(**request.observable.model_dump()).save()
+    if request.tags:
+        new.tag(request.tags)
+    return new
+
+
 @router.post("/bulk")
-async def bulk_add(request: NewBulkObservableAddRequest) -> list[Observable]:
+async def bulk_add(request: NewBulkObservableAddRequest) -> list[ObservableTypes]:
     """Bulk-creates new observables in the database."""
     added = []
     for new_observable in request.observables:
@@ -123,19 +161,16 @@ async def bulk_add(request: NewBulkObservableAddRequest) -> list[Observable]:
                 new_observable.value, tags=new_observable.tags
             )
         else:
-            observable = Observable(
-                value=new_observable.value,
-                type=new_observable.type,
-                created=datetime.datetime.now(datetime.timezone.utc),
-            ).save()
+            cls = TYPE_MAPPING[new_observable.type]
+            observable = cls(value=new_observable.value).save()
             if new_observable.tags:
                 observable = observable.tag(new_observable.tags)
         added.append(observable)
     return added
 
 
 @router.get("/{observable_id}")
-async def details(observable_id) -> Observable:
+async def details(observable_id) -> ObservableTypes:
     """Returns details about an observable."""
     observable = Observable.get(observable_id)
     if not observable:
@@ -145,7 +180,7 @@ async def details(observable_id) -> Observable:
 
 
 @router.post("/{observable_id}/context")
-async def add_context(observable_id, request: AddContextRequest) -> Observable:
+async def add_context(observable_id, request: AddContextRequest) -> ObservableTypes:
     """Adds context to an observable."""
     observable = Observable.get(observable_id)
     if not observable:
@@ -160,7 +195,7 @@ async def add_context(observable_id, request: AddContextRequest) -> Observable:
 
 
 @router.post("/{observable_id}/context/delete")
-async def delete_context(observable_id, request: DeleteContextRequest) -> Observable:
+async def delete_context(observable_id, request: DeleteContextRequest) -> ObservableTypes:
     """Removes context to an observable."""
     observable = Observable.get(observable_id)
     if not observable:
@@ -193,7 +228,7 @@ async def search(request: ObservableSearchRequest) -> ObservableSearchResponse:
 
 
 @router.post("/add_text")
-async def add_text(request: AddTextRequest) -> Observable:
+async def add_text(request: AddTextRequest) -> ObservableTypes:
     """Adds and returns an observable for a given string, attempting to guess
     its type."""
     try: