Load schemas dynamically

core/schemas/__init__.py

@@ -0,0 +1,25 @@
+import importlib
+import inspect
+from pathlib import Path
+
+import aenum
+
+from core.schemas import observable
+
+print("Registering observable types")
+
+for observable_file in Path(__file__).parent.glob("observables/**/*.py"):
+    if observable_file.stem == "__init__":
+        continue
+    print(f"Registering observable type {observable_file.stem}")
+    if observable_file.parent.stem == "observables":
+        module_name = f"core.schemas.observables.{observable_file.stem}"
+    elif observable_file.parent.stem == "private":
+        module_name = f"core.schemas.observables.private.{observable_file.stem}"
+    aenum.extend_enum(
+        observable.ObservableType, observable_file.stem, observable_file.stem
+    )
+    module = importlib.import_module(module_name)
+    for _, obj in inspect.getmembers(module, inspect.isclass):
+        if issubclass(obj, observable.Observable):
+            observable.TYPE_MAPPING[observable_file.stem] = obj

core/schemas/observable.py

@@ -2,9 +2,13 @@
 
 import datetime
 import re
-from enum import Enum
+
+# from enum import Enum, EnumMeta
 from typing import ClassVar, Literal
 
+# Data Schema
+# Dynamically register all observable types
+import aenum
 import validators
 from pydantic import Field, computed_field
 
@@ -13,39 +17,11 @@
 from core.schemas.model import YetiTagModel
 
 
-# Data Schema
-class ObservableType(str, Enum):
-    asn = "asn"
-    bic = "bic"
-    certificate = "certificate"
-    cidr = "cidr"
-    command_line = "command_line"
-    docker_image = "docker_image"
-    email = "email"
-    file = "file"
-    guess = "guess"
-    hostname = "hostname"
-    iban = "iban"
-    imphash = "imphash"
-    ipv4 = "ipv4"
-    ipv6 = "ipv6"
-    ja3 = "ja3"
-    jarm = "jarm"
-    mac_address = "mac_address"
-    md5 = "md5"
-    generic = "generic"
-    path = "path"
-    registry_key = "registry_key"
-    sha1 = "sha1"
-    sha256 = "sha256"
-    ssdeep = "ssdeep"
-    tlsh = "tlsh"
-    url = "url"
-    user_agent = "user_agent"
-    user_account = "user_account"
-    wallet = "wallet"
-    mutex = "mutex"
-    named_pipe = "named_pipe"
+class ObservableType(str, aenum.Enum):
+    pass
+
+
+TYPE_MAPPING = {}
 
 
 class Observable(YetiTagModel, database_arango.ArangoYetiConnector):
@@ -151,6 +127,12 @@ def delete_context(
         return self.save()
 
 
+TYPE_MAPPING.update({"observable": Observable, "observables": Observable})
+
+
+TYPE_VALIDATOR_MAP = {}
+
+
 TYPE_VALIDATOR_MAP = {
     ObservableType.ipv4: validators.ipv4,
     ObservableType.ipv6: validators.ipv6,
@@ -202,35 +184,35 @@ def find_type(value: str) -> ObservableType | None:
 # Import all observable types, as these register themselves in the TYPE_MAPPING
 # disable: pylint=wrong-import-position
 # noqa: F401, E402
-from core.schemas.observables import (
-    asn,
-    bic,
-    certificate,
-    cidr,
-    command_line,
-    docker_image,
-    email,
-    file,
-    generic_observable,
-    hostname,
-    iban,
-    imphash,
-    ipv4,
-    ipv6,
-    ja3,
-    jarm,
-    mac_address,
-    md5,
-    mutex,
-    named_pipe,
-    path,
-    registry_key,
-    sha1,
-    sha256,
-    ssdeep,
-    tlsh,
-    url,
-    user_account,
-    user_agent,
-    wallet,
-)
+# from core.schemas.observables import (
+#     asn,
+#     bic,
+#     certificate,
+#     cidr,
+#     command_line,
+#     docker_image,
+#     email,
+#     file,
+#     generic_observable,
+#     hostname,
+#     iban,
+#     imphash,
+#     ipv4,
+#     ipv6,
+#     ja3,
+#     jarm,
+#     mac_address,
+#     md5,
+#     mutex,
+#     named_pipe,
+#     path,
+#     registry_key,
+#     sha1,
+#     sha256,
+#     ssdeep,
+#     tlsh,
+#     url,
+#     user_account,
+#     user_agent,
+#     wallet,
+# )

core/schemas/observables/__init__.py

@@ -0,0 +1 @@
+

core/schemas/observables/asn.py

@@ -1,12 +1,7 @@
-from typing import Literal
-
 from core.schemas import observable
 
 
 class ASN(observable.Observable):
-    type: Literal[observable.ObservableType.asn] = observable.ObservableType.asn
+    type: observable.ObservableType = observable.ObservableType.asn
     country: str | None = None
     description: str | None = None
-
-
-observable.TYPE_MAPPING[observable.ObservableType.asn] = ASN

core/schemas/observables/bic.py

@@ -1,10 +1,5 @@
-from typing import Literal
-
 from core.schemas import observable
 
 
 class BIC(observable.Observable):
-    type: Literal[observable.ObservableType.bic] = observable.ObservableType.bic
-
-
-observable.TYPE_MAPPING[observable.ObservableType.bic] = BIC
+    type: observable.ObservableType = observable.ObservableType.bic

core/schemas/observables/certificate.py

@@ -1,6 +1,5 @@
 import datetime
 import hashlib
-from typing import Literal
 
 from pydantic import Field
 
@@ -22,9 +21,7 @@ class Certificate(observable.Observable):
         fingerprint: the certificate fingerprint.
     """
 
-    type: Literal[observable.ObservableType.certificate] = (
-        observable.ObservableType.certificate
-    )
+    type: observable.ObservableType = observable.ObservableType.certificate
     last_seen: datetime.datetime = Field(default_factory=now)
     first_seen: datetime.datetime = Field(default_factory=now)
     issuer: str | None = None
@@ -38,6 +35,3 @@ class Certificate(observable.Observable):
     def from_data(cls, data: bytes):
         hash_256 = hashlib.sha256(data).hexdigest()
         return cls(value=f"CERT:{hash_256}")
-
-
-observable.TYPE_MAPPING[observable.ObservableType.certificate] = Certificate

core/schemas/observables/cidr.py

@@ -1,10 +1,5 @@
-from typing import Literal
-
 from core.schemas import observable
 
 
 class CIDR(observable.Observable):
-    type: Literal[observable.ObservableType.cidr] = observable.ObservableType.cidr
-
-
-observable.TYPE_MAPPING[observable.ObservableType.cidr] = CIDR
+    type: observable.ObservableType = observable.ObservableType.cidr

core/schemas/observables/command_line.py

@@ -1,12 +1,5 @@
-from typing import Literal
-
 from core.schemas import observable
 
 
 class CommandLine(observable.Observable):
-    type: Literal[observable.ObservableType.command_line] = (
-        observable.ObservableType.command_line
-    )
-
-
-observable.TYPE_MAPPING[observable.ObservableType.command_line] = CommandLine
+    type: observable.ObservableType = observable.ObservableType.command_line

core/schemas/observables/docker_image.py

@@ -1,12 +1,5 @@
-from typing import Literal
-
 from core.schemas import observable
 
 
 class DockerImage(observable.Observable):
-    type: Literal[observable.ObservableType.docker_image] = (
-        observable.ObservableType.docker_image
-    )
-
-
-observable.TYPE_MAPPING[observable.ObservableType.docker_image] = DockerImage
+    type: observable.ObservableType = observable.ObservableType.docker_image

core/schemas/observables/email.py

@@ -1,10 +1,5 @@
-from typing import Literal
-
 from core.schemas import observable
 
 
 class Email(observable.Observable):
-    type: Literal[observable.ObservableType.email] = observable.ObservableType.email
-
-
-observable.TYPE_MAPPING[observable.ObservableType.email] = Email
+    type: observable.ObservableType = observable.ObservableType.email

core/schemas/observables/file.py

@@ -1,5 +1,3 @@
-from typing import Literal
-
 from core.schemas import observable
 
 
@@ -10,13 +8,10 @@ class File(observable.Observable):
     Value should to be in the form FILE:<HASH>.
     """
 
-    type: Literal[observable.ObservableType.file] = observable.ObservableType.file
+    type: observable.ObservableType = observable.ObservableType.file
     name: str | None = None
     size: int | None = None
     sha256: str | None = None
     md5: str | None = None
     sha1: str | None = None
     mime_type: str | None = None
-
-
-observable.TYPE_MAPPING[observable.ObservableType.file] = File

core/schemas/observables/generic.py

@@ -0,0 +1,9 @@
+from typing import Literal
+
+from core.schemas import observable
+
+
+class Generic(observable.Observable):
+    """Use this type of Observable for any type of observable that doesn't fit into any other category."""
+
+    type: observable.ObservableType = observable.ObservableType.generic

core/schemas/observables/hostname.py

@@ -1,12 +1,5 @@
-from typing import Literal
-
 from core.schemas import observable
 
 
 class Hostname(observable.Observable):
-    type: Literal[observable.ObservableType.hostname] = (
-        observable.ObservableType.hostname
-    )
-
-
-observable.TYPE_MAPPING[observable.ObservableType.hostname] = Hostname
+    type: observable.ObservableType = observable.ObservableType.hostname

core/schemas/observables/iban.py

@@ -1,10 +1,5 @@
-from typing import Literal
-
 from core.schemas import observable
 
 
 class IBAN(observable.Observable):
-    type: Literal[observable.ObservableType.iban] = observable.ObservableType.iban
-
-
-observable.TYPE_MAPPING[observable.ObservableType.iban] = IBAN
+    type: observable.ObservableType = observable.ObservableType.iban

core/schemas/observables/imphash.py

@@ -1,10 +1,5 @@
-from typing import Literal
-
 from core.schemas import observable
 
 
 class Imphash(observable.Observable):
-    type: Literal[observable.ObservableType.imphash] = observable.ObservableType.imphash
-
-
-observable.TYPE_MAPPING[observable.ObservableType.imphash] = Imphash
+    type: observable.ObservableType = observable.ObservableType.imphash

core/schemas/observables/ipv4.py

@@ -1,10 +1,5 @@
-from typing import Literal
-
 from core.schemas import observable
 
 
 class IPv4(observable.Observable):
-    type: Literal[observable.ObservableType.ipv4] = observable.ObservableType.ipv4
-
-
-observable.TYPE_MAPPING[observable.ObservableType.ipv4] = IPv4
+    type: observable.ObservableType = observable.ObservableType.ipv4

core/schemas/observables/ipv6.py

@@ -1,10 +1,5 @@
-from typing import Literal
-
 from core.schemas import observable
 
 
 class IPv6(observable.Observable):
-    type: Literal[observable.ObservableType.ipv6] = observable.ObservableType.ipv6
-
-
-observable.TYPE_MAPPING[observable.ObservableType.ipv6] = IPv6
+    type: observable.ObservableType = observable.ObservableType.ipv6

core/schemas/observables/ja3.py

@@ -1,10 +1,5 @@
-from typing import Literal
-
 from core.schemas import observable
 
 
 class JA3(observable.Observable):
-    type: Literal[observable.ObservableType.ja3] = observable.ObservableType.ja3
-
-
-observable.TYPE_MAPPING[observable.ObservableType.ja3] = JA3
+    type: observable.ObservableType = observable.ObservableType.ja3