Fix: error when forcing type to "guess" (#971)

core/web/apiv2/graph.py

@@ -1,11 +1,13 @@
 import datetime
 from enum import Enum
 
-from core.schemas import entity, graph, indicator, observable, tag
 from fastapi import APIRouter, HTTPException
 from pydantic import BaseModel, ConfigDict, ValidationInfo
 from pydantic.functional_validators import field_validator
 
+from core.schemas import entity, graph, indicator, observable, tag
+from core.schemas.observable import ObservableType
+
 GRAPH_TYPE_MAPPINGS = (
     {}
 )  # type: dict[str, Type[entity.Entity] | Type[observable.Observable] | Type[indicator.Indicator]]
@@ -176,14 +178,15 @@ async def match(request: AnalysisRequest) -> AnalysisResponse:
     known = {}  # type: dict[str, observable.Observable]
     if request.add_unknown:
         for value in request.observables:
-            if request.add_type:
-                obs = observable.TYPE_MAPPING[request.add_type](value=value).save()
-                obs.tag(request.add_tags)
-            else:
+            if request.add_type == ObservableType.guess or not request.add_type:
                 try:
                     observable.Observable.add_text(value, tags=request.add_tags)
                 except ValueError:
-                    pass
+                    continue
+            elif request.add_type:
+                obs = observable.TYPE_MAPPING[request.add_type](value=value).save()
+                obs.tag(request.add_tags)
+
             unknown.discard(value)
 
     db_observables, _ = observable.Observable.filter(

tests/apiv2/auth.py

@@ -125,29 +125,3 @@ def test_api_key_bearer(self) -> None:
         data = response.json()
         self.assertEqual(response.status_code, 200)
         self.assertEqual(data["username"], "tomchop")
-
-    def test_logout(self) -> None:
-        response = client.post(
-            "/api/v2/auth/token", data={"username": "tomchop", "password": "test"}
-        )
-        data = response.json()
-        token = data["access_token"]
-
-        response = client.get(
-            "/api/v2/auth/me", headers={"cookie": "yeti_session=" + token}
-        )
-        data = response.json()
-        self.assertEqual(response.status_code, 200)
-        self.assertEqual(data["username"], "tomchop")
-
-        response = client.post(
-            "/api/v2/auth/logout", headers={"cookie": "yeti_session=" + token}
-        )
-        self.assertEqual(response.status_code, 200)
-
-        response = client.get(
-            "/api/v2/auth/me", headers={"cookie": "yeti_session=" + token}
-        )
-        self.assertEqual(response.status_code, 401)
-        data = response.json()
-        self.assertEqual(data["detail"], "Could not validate credentials")

tests/apiv2/graph.py

@@ -323,3 +323,33 @@ def test_match_known_observables_have_tags(self):
         self.assertEqual(len(data["known"]), 1)
         self.assertEqual(data["known"][0]["value"], "test1.com")
         self.assertEqual(sorted(data["known"][0]["tags"].keys()), ["tag1", "tag2"])
+
+    def test_match_guessing_type(self):
+        response = client.post(
+            "/api/v2/graph/match",
+            json={
+                "observables": ["test3.com"],
+                "add_unknown": True,
+                "add_type": "guess",
+            },
+        )
+
+        data = response.json()
+        self.assertEqual(response.status_code, 200, data)
+        self.assertEqual(len(data["known"]), 1)
+        self.assertEqual(data["known"][0]["value"], "test3.com")
+        self.assertEqual(data["known"][0]["type"], "hostname")
+
+        response = client.post(
+            "/api/v2/graph/match",
+            json={
+                "observables": ["test4.com"],
+                "add_unknown": True
+            },
+        )
+
+        data = response.json()
+        self.assertEqual(response.status_code, 200, data)
+        self.assertEqual(len(data["known"]), 1)
+        self.assertEqual(data["known"][0]["value"], "test4.com")
+        self.assertEqual(data["known"][0]["type"], "hostname")

yeti.conf.sample

@@ -28,7 +28,7 @@ module = local
 # SECRET_KEY = SECRET
 # ALGORITHM = HS256
 # ACCESS_TOKEN_EXPIRE_MINUTES = 30
-# enabled = False
+enabled = True
 
 # OIDC
 #