From f6df03cd116d24cc17a7dd3a4707614543442ce7 Mon Sep 17 00:00:00 2001 From: Fred Baguelin Date: Sat, 14 Sep 2024 15:11:32 +0200 Subject: [PATCH 01/39] Load observables dynamically --- core/schemas/__init__.py | 23 ++++ core/schemas/observable.py | 112 +++++++----------- core/schemas/observables/__init__.py | 1 + core/schemas/observables/asn.py | 9 +- core/schemas/observables/bic.py | 7 +- core/schemas/observables/certificate.py | 10 +- core/schemas/observables/cidr.py | 7 +- core/schemas/observables/command_line.py | 9 +- core/schemas/observables/docker_image.py | 9 +- core/schemas/observables/email.py | 7 +- core/schemas/observables/file.py | 9 +- core/schemas/observables/generic.py | 8 ++ .../schemas/observables/generic_observable.py | 12 -- core/schemas/observables/hostname.py | 9 +- core/schemas/observables/iban.py | 7 +- core/schemas/observables/imphash.py | 7 +- core/schemas/observables/ipv4.py | 7 +- core/schemas/observables/ipv6.py | 7 +- core/schemas/observables/ja3.py | 7 +- core/schemas/observables/jarm.py | 7 +- core/schemas/observables/mac_address.py | 9 +- core/schemas/observables/md5.py | 7 +- core/schemas/observables/mutex.py | 7 +- core/schemas/observables/named_pipe.py | 9 +- core/schemas/observables/path.py | 7 +- core/schemas/observables/private/.gitignore | 3 + core/schemas/observables/private/README.md | 2 + core/schemas/observables/registry_key.py | 9 +- core/schemas/observables/sha1.py | 7 +- core/schemas/observables/sha256.py | 7 +- core/schemas/observables/ssdeep.py | 9 +- core/schemas/observables/tlsh.py | 7 +- core/schemas/observables/url.py | 7 +- core/schemas/observables/user_account.py | 9 +- core/schemas/observables/user_agent.py | 9 +- core/schemas/observables/wallet.py | 9 +- 36 files changed, 119 insertions(+), 272 deletions(-) create mode 100644 core/schemas/observables/generic.py delete mode 100644 core/schemas/observables/generic_observable.py create mode 100644 core/schemas/observables/private/.gitignore create mode 100644 core/schemas/observables/private/README.md diff --git a/core/schemas/__init__.py b/core/schemas/__init__.py index e69de29bb..20fbbfdc3 100644 --- a/core/schemas/__init__.py +++ b/core/schemas/__init__.py @@ -0,0 +1,23 @@ +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 \ No newline at end of file diff --git a/core/schemas/observable.py b/core/schemas/observable.py index b8b54f243..ec5f65310 100644 --- a/core/schemas/observable.py +++ b/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,10 @@ 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): @@ -150,6 +125,11 @@ def delete_context( break return self.save() +TYPE_MAPPING.update({"observable": Observable, "observables": Observable}) + + +TYPE_VALIDATOR_MAP = {} + TYPE_VALIDATOR_MAP = { ObservableType.ipv4: validators.ipv4, @@ -202,35 +182,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, +# ) diff --git a/core/schemas/observables/__init__.py b/core/schemas/observables/__init__.py index e69de29bb..8b1378917 100644 --- a/core/schemas/observables/__init__.py +++ b/core/schemas/observables/__init__.py @@ -0,0 +1 @@ + diff --git a/core/schemas/observables/asn.py b/core/schemas/observables/asn.py index 1aeba6027..5d6062b58 100644 --- a/core/schemas/observables/asn.py +++ b/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 + description: str | None = None \ No newline at end of file diff --git a/core/schemas/observables/bic.py b/core/schemas/observables/bic.py index 1eb80de98..c6aca2c69 100644 --- a/core/schemas/observables/bic.py +++ b/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 diff --git a/core/schemas/observables/certificate.py b/core/schemas/observables/certificate.py index 1e2127ab7..88357d25f 100644 --- a/core/schemas/observables/certificate.py +++ b/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 @@ -37,7 +34,4 @@ class Certificate(observable.Observable): @classmethod 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 + return cls(value=f"CERT:{hash_256}") \ No newline at end of file diff --git a/core/schemas/observables/cidr.py b/core/schemas/observables/cidr.py index b9e73b375..4a46d0941 100644 --- a/core/schemas/observables/cidr.py +++ b/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 \ No newline at end of file diff --git a/core/schemas/observables/command_line.py b/core/schemas/observables/command_line.py index b166eff27..c2119ab7f 100644 --- a/core/schemas/observables/command_line.py +++ b/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 diff --git a/core/schemas/observables/docker_image.py b/core/schemas/observables/docker_image.py index bdc342793..3d166e4eb 100644 --- a/core/schemas/observables/docker_image.py +++ b/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 diff --git a/core/schemas/observables/email.py b/core/schemas/observables/email.py index 8e51696d4..f19d48c9d 100644 --- a/core/schemas/observables/email.py +++ b/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 \ No newline at end of file diff --git a/core/schemas/observables/file.py b/core/schemas/observables/file.py index df700033a..e8a4d4d85 100644 --- a/core/schemas/observables/file.py +++ b/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:. """ - 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 + mime_type: str | None = None \ No newline at end of file diff --git a/core/schemas/observables/generic.py b/core/schemas/observables/generic.py new file mode 100644 index 000000000..a9c6ff800 --- /dev/null +++ b/core/schemas/observables/generic.py @@ -0,0 +1,8 @@ +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 diff --git a/core/schemas/observables/generic_observable.py b/core/schemas/observables/generic_observable.py deleted file mode 100644 index 64dfdc0c3..000000000 --- a/core/schemas/observables/generic_observable.py +++ /dev/null @@ -1,12 +0,0 @@ -from typing import Literal - -from core.schemas import observable - - -class GenericObservable(observable.Observable): - """Use this type of Observable for any type of observable that doesn't fit into any other category.""" - - type: Literal[observable.ObservableType.generic] = observable.ObservableType.generic - - -observable.TYPE_MAPPING[observable.ObservableType.generic] = GenericObservable diff --git a/core/schemas/observables/hostname.py b/core/schemas/observables/hostname.py index eae0cfffa..46e6aa1d9 100644 --- a/core/schemas/observables/hostname.py +++ b/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 \ No newline at end of file diff --git a/core/schemas/observables/iban.py b/core/schemas/observables/iban.py index a97db090e..dbedf72e9 100644 --- a/core/schemas/observables/iban.py +++ b/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 \ No newline at end of file diff --git a/core/schemas/observables/imphash.py b/core/schemas/observables/imphash.py index b9e6d477a..d53b2b9f5 100644 --- a/core/schemas/observables/imphash.py +++ b/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 \ No newline at end of file diff --git a/core/schemas/observables/ipv4.py b/core/schemas/observables/ipv4.py index 7aba9ed81..0520e0b03 100644 --- a/core/schemas/observables/ipv4.py +++ b/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 diff --git a/core/schemas/observables/ipv6.py b/core/schemas/observables/ipv6.py index 6568f3197..6f3e7ea70 100644 --- a/core/schemas/observables/ipv6.py +++ b/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 \ No newline at end of file diff --git a/core/schemas/observables/ja3.py b/core/schemas/observables/ja3.py index 1ca270462..c082b2824 100644 --- a/core/schemas/observables/ja3.py +++ b/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 \ No newline at end of file diff --git a/core/schemas/observables/jarm.py b/core/schemas/observables/jarm.py index 4b1001010..a766319a2 100644 --- a/core/schemas/observables/jarm.py +++ b/core/schemas/observables/jarm.py @@ -1,10 +1,5 @@ -from typing import Literal - from core.schemas import observable class JARM(observable.Observable): - type: Literal[observable.ObservableType.jarm] = observable.ObservableType.jarm - - -observable.TYPE_MAPPING[observable.ObservableType.ja3] = JARM + type: observable.ObservableType = observable.ObservableType.jarm \ No newline at end of file diff --git a/core/schemas/observables/mac_address.py b/core/schemas/observables/mac_address.py index 97c4fd018..9e18aaf33 100644 --- a/core/schemas/observables/mac_address.py +++ b/core/schemas/observables/mac_address.py @@ -1,12 +1,5 @@ -from typing import Literal - from core.schemas import observable class MacAddress(observable.Observable): - type: Literal[observable.ObservableType.mac_address] = ( - observable.ObservableType.mac_address - ) - - -observable.TYPE_MAPPING[observable.ObservableType.mac_address] = MacAddress + type: observable.ObservableType = observable.ObservableType.mac_address \ No newline at end of file diff --git a/core/schemas/observables/md5.py b/core/schemas/observables/md5.py index 24f13eaba..3e2df4426 100644 --- a/core/schemas/observables/md5.py +++ b/core/schemas/observables/md5.py @@ -1,10 +1,5 @@ -from typing import Literal - from core.schemas import observable class MD5(observable.Observable): - type: Literal[observable.ObservableType.md5] = observable.ObservableType.md5 - - -observable.TYPE_MAPPING[observable.ObservableType.md5] = MD5 + type: observable.ObservableType = observable.ObservableType.md5 \ No newline at end of file diff --git a/core/schemas/observables/mutex.py b/core/schemas/observables/mutex.py index 2ad41e614..e56ded500 100644 --- a/core/schemas/observables/mutex.py +++ b/core/schemas/observables/mutex.py @@ -1,10 +1,5 @@ -from typing import Literal - from core.schemas import observable class Mutex(observable.Observable): - type: Literal[observable.ObservableType.mutex] = observable.ObservableType.mutex - - -observable.TYPE_MAPPING[observable.ObservableType.mutex] = Mutex + type: observable.ObservableType = observable.ObservableType.mutex \ No newline at end of file diff --git a/core/schemas/observables/named_pipe.py b/core/schemas/observables/named_pipe.py index e75165386..b71e58968 100644 --- a/core/schemas/observables/named_pipe.py +++ b/core/schemas/observables/named_pipe.py @@ -1,12 +1,5 @@ -from typing import Literal - from core.schemas import observable class NamedPipe(observable.Observable): - type: Literal[observable.ObservableType.named_pipe] = ( - observable.ObservableType.named_pipe - ) - - -observable.TYPE_MAPPING[observable.ObservableType.named_pipe] = NamedPipe + type: observable.ObservableType = observable.ObservableType.named_pipe \ No newline at end of file diff --git a/core/schemas/observables/path.py b/core/schemas/observables/path.py index 4d40198d9..9e12447e6 100644 --- a/core/schemas/observables/path.py +++ b/core/schemas/observables/path.py @@ -1,10 +1,5 @@ -from typing import Literal - from core.schemas import observable class Path(observable.Observable): - type: Literal[observable.ObservableType.path] = observable.ObservableType.path - - -observable.TYPE_MAPPING[observable.ObservableType.path] = Path + type: observable.ObservableType = observable.ObservableType.path \ No newline at end of file diff --git a/core/schemas/observables/private/.gitignore b/core/schemas/observables/private/.gitignore new file mode 100644 index 000000000..e5af87e9b --- /dev/null +++ b/core/schemas/observables/private/.gitignore @@ -0,0 +1,3 @@ +* +!.gitignore +!README.md \ No newline at end of file diff --git a/core/schemas/observables/private/README.md b/core/schemas/observables/private/README.md new file mode 100644 index 000000000..d612910f1 --- /dev/null +++ b/core/schemas/observables/private/README.md @@ -0,0 +1,2 @@ +### Private observables +This directory is where you should place your private observables. It could be named anything else, but this one has a `.gitignore` so you don't mess things up. ;-) diff --git a/core/schemas/observables/registry_key.py b/core/schemas/observables/registry_key.py index 29298bc49..0de001cbc 100644 --- a/core/schemas/observables/registry_key.py +++ b/core/schemas/observables/registry_key.py @@ -26,13 +26,8 @@ class RegistryKey(observable.Observable): path_file: The filesystem path to the file that contains the registry key value. """ - type: Literal[observable.ObservableType.registry_key] = ( - observable.ObservableType.registry_key - ) + type: observable.ObservableType = observable.ObservableType.registry_key key: str data: bytes hive: RegistryHive - path_file: str | None = None - - -observable.TYPE_MAPPING[observable.ObservableType.registry_key] = RegistryKey + path_file: str | None = None \ No newline at end of file diff --git a/core/schemas/observables/sha1.py b/core/schemas/observables/sha1.py index 07e43b145..fe781ad54 100644 --- a/core/schemas/observables/sha1.py +++ b/core/schemas/observables/sha1.py @@ -1,10 +1,5 @@ -from typing import Literal - from core.schemas import observable class SHA1(observable.Observable): - type: Literal[observable.ObservableType.sha1] = observable.ObservableType.sha1 - - -observable.TYPE_MAPPING[observable.ObservableType.sha1] = SHA1 + type: observable.ObservableType = observable.ObservableType.sha1 \ No newline at end of file diff --git a/core/schemas/observables/sha256.py b/core/schemas/observables/sha256.py index c28d6f7e0..0bcbaaf41 100644 --- a/core/schemas/observables/sha256.py +++ b/core/schemas/observables/sha256.py @@ -1,10 +1,5 @@ -from typing import Literal - from core.schemas import observable class SHA256(observable.Observable): - type: Literal[observable.ObservableType.sha256] = observable.ObservableType.sha256 - - -observable.TYPE_MAPPING[observable.ObservableType.sha256] = SHA256 + type: observable.ObservableType = observable.ObservableType.sha256 \ No newline at end of file diff --git a/core/schemas/observables/ssdeep.py b/core/schemas/observables/ssdeep.py index d355d6ffd..35b68de5c 100644 --- a/core/schemas/observables/ssdeep.py +++ b/core/schemas/observables/ssdeep.py @@ -1,10 +1,5 @@ -from typing import Literal - from core.schemas import observable -class SsdeepHash(observable.Observable): - type: Literal[observable.ObservableType.ssdeep] = observable.ObservableType.ssdeep - - -observable.TYPE_MAPPING[observable.ObservableType.ssdeep] = SsdeepHash +class Ssdeep(observable.Observable): + type: observable.ObservableType = observable.ObservableType.ssdeep diff --git a/core/schemas/observables/tlsh.py b/core/schemas/observables/tlsh.py index 08cdcf7f8..ddbf4982d 100644 --- a/core/schemas/observables/tlsh.py +++ b/core/schemas/observables/tlsh.py @@ -1,10 +1,5 @@ -from typing import Literal - from core.schemas import observable class TLSH(observable.Observable): - type: Literal[observable.ObservableType.tlsh] = observable.ObservableType.tlsh - - -observable.TYPE_MAPPING[observable.ObservableType.tlsh] = TLSH + type: observable.ObservableType = observable.ObservableType.tlsh \ No newline at end of file diff --git a/core/schemas/observables/url.py b/core/schemas/observables/url.py index 76ca203ae..dfa6c8fcb 100644 --- a/core/schemas/observables/url.py +++ b/core/schemas/observables/url.py @@ -1,10 +1,5 @@ -from typing import Literal - from core.schemas import observable class Url(observable.Observable): - type: Literal[observable.ObservableType.url] = observable.ObservableType.url - - -observable.TYPE_MAPPING[observable.ObservableType.url] = Url + type: observable.ObservableType = observable.ObservableType.url \ No newline at end of file diff --git a/core/schemas/observables/user_account.py b/core/schemas/observables/user_account.py index fa8d172d7..fef5a6b31 100644 --- a/core/schemas/observables/user_account.py +++ b/core/schemas/observables/user_account.py @@ -14,9 +14,7 @@ class UserAccount(observable.Observable): Value should to be in the form :. """ - type: Literal[observable.ObservableType.user_account] = ( - observable.ObservableType.user_account - ) + type: observable.ObservableType = observable.ObservableType.user_account user_id: str | None = None credential: str | None = None account_login: str | None = None @@ -39,7 +37,4 @@ def check_timestamp_coherence(self) -> "UserAccount": raise ValueError( "Account created date is after account expiration date." ) - return self - - -observable.TYPE_MAPPING[observable.ObservableType.user_account] = UserAccount + return self \ No newline at end of file diff --git a/core/schemas/observables/user_agent.py b/core/schemas/observables/user_agent.py index b1d163c83..fa24ac68c 100644 --- a/core/schemas/observables/user_agent.py +++ b/core/schemas/observables/user_agent.py @@ -1,12 +1,5 @@ -from typing import Literal - from core.schemas import observable class UserAgent(observable.Observable): - type: Literal[observable.ObservableType.user_agent] = ( - observable.ObservableType.user_agent - ) - - -observable.TYPE_MAPPING[observable.ObservableType.user_agent] = UserAgent + type: observable.ObservableType = observable.ObservableType.user_agent \ No newline at end of file diff --git a/core/schemas/observables/wallet.py b/core/schemas/observables/wallet.py index 995f1094b..a39f0b932 100644 --- a/core/schemas/observables/wallet.py +++ b/core/schemas/observables/wallet.py @@ -1,5 +1,3 @@ -from typing import Literal - from core.schemas import observable @@ -10,9 +8,6 @@ class Wallet(observable.Observable): Value should be in the form :
. """ - type: Literal[observable.ObservableType.wallet] = observable.ObservableType.wallet + type: observable.ObservableType = observable.ObservableType.wallet coin: str | None = None - address: str | None = None - - -observable.TYPE_MAPPING[observable.ObservableType.wallet] = Wallet + address: str | None = None \ No newline at end of file From 76e4a38adcf201e2c75ca2b544dc5a7abed091ad Mon Sep 17 00:00:00 2001 From: Fred Baguelin Date: Sat, 14 Sep 2024 15:17:34 +0200 Subject: [PATCH 02/39] Ruff format --- core/schemas/__init__.py | 6 ++++-- core/schemas/observable.py | 4 +++- core/schemas/observables/asn.py | 2 +- core/schemas/observables/certificate.py | 2 +- core/schemas/observables/cidr.py | 2 +- core/schemas/observables/email.py | 2 +- core/schemas/observables/file.py | 2 +- core/schemas/observables/generic.py | 1 + core/schemas/observables/hostname.py | 2 +- core/schemas/observables/iban.py | 2 +- core/schemas/observables/imphash.py | 2 +- core/schemas/observables/ipv6.py | 2 +- core/schemas/observables/ja3.py | 2 +- core/schemas/observables/jarm.py | 2 +- core/schemas/observables/mac_address.py | 2 +- core/schemas/observables/md5.py | 2 +- core/schemas/observables/mutex.py | 2 +- core/schemas/observables/named_pipe.py | 2 +- core/schemas/observables/package.py | 7 +++++++ core/schemas/observables/path.py | 2 +- core/schemas/observables/registry_key.py | 2 +- core/schemas/observables/sha1.py | 2 +- core/schemas/observables/sha256.py | 2 +- core/schemas/observables/tlsh.py | 2 +- core/schemas/observables/url.py | 2 +- core/schemas/observables/user_account.py | 2 +- core/schemas/observables/user_agent.py | 2 +- core/schemas/observables/wallet.py | 2 +- 28 files changed, 39 insertions(+), 27 deletions(-) create mode 100644 core/schemas/observables/package.py diff --git a/core/schemas/__init__.py b/core/schemas/__init__.py index 20fbbfdc3..8d727e8fd 100644 --- a/core/schemas/__init__.py +++ b/core/schemas/__init__.py @@ -16,8 +16,10 @@ 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) + 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 \ No newline at end of file + observable.TYPE_MAPPING[observable_file.stem] = obj diff --git a/core/schemas/observable.py b/core/schemas/observable.py index ec5f65310..853bfc642 100644 --- a/core/schemas/observable.py +++ b/core/schemas/observable.py @@ -3,7 +3,7 @@ import datetime import re -#from enum import Enum, EnumMeta +# from enum import Enum, EnumMeta from typing import ClassVar, Literal # Data Schema @@ -20,6 +20,7 @@ class ObservableType(str, aenum.Enum): pass + TYPE_MAPPING = {} @@ -125,6 +126,7 @@ def delete_context( break return self.save() + TYPE_MAPPING.update({"observable": Observable, "observables": Observable}) diff --git a/core/schemas/observables/asn.py b/core/schemas/observables/asn.py index 5d6062b58..8866e010c 100644 --- a/core/schemas/observables/asn.py +++ b/core/schemas/observables/asn.py @@ -4,4 +4,4 @@ class ASN(observable.Observable): type: observable.ObservableType = observable.ObservableType.asn country: str | None = None - description: str | None = None \ No newline at end of file + description: str | None = None diff --git a/core/schemas/observables/certificate.py b/core/schemas/observables/certificate.py index 88357d25f..a9433988e 100644 --- a/core/schemas/observables/certificate.py +++ b/core/schemas/observables/certificate.py @@ -34,4 +34,4 @@ class Certificate(observable.Observable): @classmethod def from_data(cls, data: bytes): hash_256 = hashlib.sha256(data).hexdigest() - return cls(value=f"CERT:{hash_256}") \ No newline at end of file + return cls(value=f"CERT:{hash_256}") diff --git a/core/schemas/observables/cidr.py b/core/schemas/observables/cidr.py index 4a46d0941..97433ff8e 100644 --- a/core/schemas/observables/cidr.py +++ b/core/schemas/observables/cidr.py @@ -2,4 +2,4 @@ class CIDR(observable.Observable): - type: observable.ObservableType = observable.ObservableType.cidr \ No newline at end of file + type: observable.ObservableType = observable.ObservableType.cidr diff --git a/core/schemas/observables/email.py b/core/schemas/observables/email.py index f19d48c9d..b85bfebd6 100644 --- a/core/schemas/observables/email.py +++ b/core/schemas/observables/email.py @@ -2,4 +2,4 @@ class Email(observable.Observable): - type: observable.ObservableType = observable.ObservableType.email \ No newline at end of file + type: observable.ObservableType = observable.ObservableType.email diff --git a/core/schemas/observables/file.py b/core/schemas/observables/file.py index e8a4d4d85..de2d8e043 100644 --- a/core/schemas/observables/file.py +++ b/core/schemas/observables/file.py @@ -14,4 +14,4 @@ class File(observable.Observable): sha256: str | None = None md5: str | None = None sha1: str | None = None - mime_type: str | None = None \ No newline at end of file + mime_type: str | None = None diff --git a/core/schemas/observables/generic.py b/core/schemas/observables/generic.py index a9c6ff800..a8e9bbe68 100644 --- a/core/schemas/observables/generic.py +++ b/core/schemas/observables/generic.py @@ -5,4 +5,5 @@ 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 diff --git a/core/schemas/observables/hostname.py b/core/schemas/observables/hostname.py index 46e6aa1d9..0eb452b55 100644 --- a/core/schemas/observables/hostname.py +++ b/core/schemas/observables/hostname.py @@ -2,4 +2,4 @@ class Hostname(observable.Observable): - type: observable.ObservableType = observable.ObservableType.hostname \ No newline at end of file + type: observable.ObservableType = observable.ObservableType.hostname diff --git a/core/schemas/observables/iban.py b/core/schemas/observables/iban.py index dbedf72e9..f963dc3bb 100644 --- a/core/schemas/observables/iban.py +++ b/core/schemas/observables/iban.py @@ -2,4 +2,4 @@ class IBAN(observable.Observable): - type: observable.ObservableType = observable.ObservableType.iban \ No newline at end of file + type: observable.ObservableType = observable.ObservableType.iban diff --git a/core/schemas/observables/imphash.py b/core/schemas/observables/imphash.py index d53b2b9f5..316e4ce29 100644 --- a/core/schemas/observables/imphash.py +++ b/core/schemas/observables/imphash.py @@ -2,4 +2,4 @@ class Imphash(observable.Observable): - type: observable.ObservableType = observable.ObservableType.imphash \ No newline at end of file + type: observable.ObservableType = observable.ObservableType.imphash diff --git a/core/schemas/observables/ipv6.py b/core/schemas/observables/ipv6.py index 6f3e7ea70..a264310f4 100644 --- a/core/schemas/observables/ipv6.py +++ b/core/schemas/observables/ipv6.py @@ -2,4 +2,4 @@ class IPv6(observable.Observable): - type: observable.ObservableType = observable.ObservableType.ipv6 \ No newline at end of file + type: observable.ObservableType = observable.ObservableType.ipv6 diff --git a/core/schemas/observables/ja3.py b/core/schemas/observables/ja3.py index c082b2824..5b93f8755 100644 --- a/core/schemas/observables/ja3.py +++ b/core/schemas/observables/ja3.py @@ -2,4 +2,4 @@ class JA3(observable.Observable): - type: observable.ObservableType = observable.ObservableType.ja3 \ No newline at end of file + type: observable.ObservableType = observable.ObservableType.ja3 diff --git a/core/schemas/observables/jarm.py b/core/schemas/observables/jarm.py index a766319a2..bcf21a87c 100644 --- a/core/schemas/observables/jarm.py +++ b/core/schemas/observables/jarm.py @@ -2,4 +2,4 @@ class JARM(observable.Observable): - type: observable.ObservableType = observable.ObservableType.jarm \ No newline at end of file + type: observable.ObservableType = observable.ObservableType.jarm diff --git a/core/schemas/observables/mac_address.py b/core/schemas/observables/mac_address.py index 9e18aaf33..5165d0260 100644 --- a/core/schemas/observables/mac_address.py +++ b/core/schemas/observables/mac_address.py @@ -2,4 +2,4 @@ class MacAddress(observable.Observable): - type: observable.ObservableType = observable.ObservableType.mac_address \ No newline at end of file + type: observable.ObservableType = observable.ObservableType.mac_address diff --git a/core/schemas/observables/md5.py b/core/schemas/observables/md5.py index 3e2df4426..7281a1beb 100644 --- a/core/schemas/observables/md5.py +++ b/core/schemas/observables/md5.py @@ -2,4 +2,4 @@ class MD5(observable.Observable): - type: observable.ObservableType = observable.ObservableType.md5 \ No newline at end of file + type: observable.ObservableType = observable.ObservableType.md5 diff --git a/core/schemas/observables/mutex.py b/core/schemas/observables/mutex.py index e56ded500..5124aed83 100644 --- a/core/schemas/observables/mutex.py +++ b/core/schemas/observables/mutex.py @@ -2,4 +2,4 @@ class Mutex(observable.Observable): - type: observable.ObservableType = observable.ObservableType.mutex \ No newline at end of file + type: observable.ObservableType = observable.ObservableType.mutex diff --git a/core/schemas/observables/named_pipe.py b/core/schemas/observables/named_pipe.py index b71e58968..314e6e2ea 100644 --- a/core/schemas/observables/named_pipe.py +++ b/core/schemas/observables/named_pipe.py @@ -2,4 +2,4 @@ class NamedPipe(observable.Observable): - type: observable.ObservableType = observable.ObservableType.named_pipe \ No newline at end of file + type: observable.ObservableType = observable.ObservableType.named_pipe diff --git a/core/schemas/observables/package.py b/core/schemas/observables/package.py new file mode 100644 index 000000000..23cff722f --- /dev/null +++ b/core/schemas/observables/package.py @@ -0,0 +1,7 @@ +from core.schemas import observable + + +class Package(observable.Observable): + type: observable.ObservableType = observable.ObservableType.package + version: str = None + regitry_type: str = None diff --git a/core/schemas/observables/path.py b/core/schemas/observables/path.py index 9e12447e6..7fb227845 100644 --- a/core/schemas/observables/path.py +++ b/core/schemas/observables/path.py @@ -2,4 +2,4 @@ class Path(observable.Observable): - type: observable.ObservableType = observable.ObservableType.path \ No newline at end of file + type: observable.ObservableType = observable.ObservableType.path diff --git a/core/schemas/observables/registry_key.py b/core/schemas/observables/registry_key.py index 0de001cbc..c9d7b5956 100644 --- a/core/schemas/observables/registry_key.py +++ b/core/schemas/observables/registry_key.py @@ -30,4 +30,4 @@ class RegistryKey(observable.Observable): key: str data: bytes hive: RegistryHive - path_file: str | None = None \ No newline at end of file + path_file: str | None = None diff --git a/core/schemas/observables/sha1.py b/core/schemas/observables/sha1.py index fe781ad54..33cb0edcb 100644 --- a/core/schemas/observables/sha1.py +++ b/core/schemas/observables/sha1.py @@ -2,4 +2,4 @@ class SHA1(observable.Observable): - type: observable.ObservableType = observable.ObservableType.sha1 \ No newline at end of file + type: observable.ObservableType = observable.ObservableType.sha1 diff --git a/core/schemas/observables/sha256.py b/core/schemas/observables/sha256.py index 0bcbaaf41..f7f452172 100644 --- a/core/schemas/observables/sha256.py +++ b/core/schemas/observables/sha256.py @@ -2,4 +2,4 @@ class SHA256(observable.Observable): - type: observable.ObservableType = observable.ObservableType.sha256 \ No newline at end of file + type: observable.ObservableType = observable.ObservableType.sha256 diff --git a/core/schemas/observables/tlsh.py b/core/schemas/observables/tlsh.py index ddbf4982d..edb97be94 100644 --- a/core/schemas/observables/tlsh.py +++ b/core/schemas/observables/tlsh.py @@ -2,4 +2,4 @@ class TLSH(observable.Observable): - type: observable.ObservableType = observable.ObservableType.tlsh \ No newline at end of file + type: observable.ObservableType = observable.ObservableType.tlsh diff --git a/core/schemas/observables/url.py b/core/schemas/observables/url.py index dfa6c8fcb..accdbb507 100644 --- a/core/schemas/observables/url.py +++ b/core/schemas/observables/url.py @@ -2,4 +2,4 @@ class Url(observable.Observable): - type: observable.ObservableType = observable.ObservableType.url \ No newline at end of file + type: observable.ObservableType = observable.ObservableType.url diff --git a/core/schemas/observables/user_account.py b/core/schemas/observables/user_account.py index fef5a6b31..9c77009e5 100644 --- a/core/schemas/observables/user_account.py +++ b/core/schemas/observables/user_account.py @@ -37,4 +37,4 @@ def check_timestamp_coherence(self) -> "UserAccount": raise ValueError( "Account created date is after account expiration date." ) - return self \ No newline at end of file + return self diff --git a/core/schemas/observables/user_agent.py b/core/schemas/observables/user_agent.py index fa24ac68c..40c0d8cdc 100644 --- a/core/schemas/observables/user_agent.py +++ b/core/schemas/observables/user_agent.py @@ -2,4 +2,4 @@ class UserAgent(observable.Observable): - type: observable.ObservableType = observable.ObservableType.user_agent \ No newline at end of file + type: observable.ObservableType = observable.ObservableType.user_agent diff --git a/core/schemas/observables/wallet.py b/core/schemas/observables/wallet.py index a39f0b932..4670f5ca3 100644 --- a/core/schemas/observables/wallet.py +++ b/core/schemas/observables/wallet.py @@ -10,4 +10,4 @@ class Wallet(observable.Observable): type: observable.ObservableType = observable.ObservableType.wallet coin: str | None = None - address: str | None = None \ No newline at end of file + address: str | None = None From 12a0b970ab380efd34b94d164889f13e02429609 Mon Sep 17 00:00:00 2001 From: Fred Baguelin Date: Mon, 16 Sep 2024 22:49:08 +0200 Subject: [PATCH 03/39] Remove end of file observables imports --- core/schemas/observable.py | 39 +------------------------------------- 1 file changed, 1 insertion(+), 38 deletions(-) diff --git a/core/schemas/observable.py b/core/schemas/observable.py index 853bfc642..e0e4c5f43 100644 --- a/core/schemas/observable.py +++ b/core/schemas/observable.py @@ -178,41 +178,4 @@ def find_type(value: str) -> ObservableType | None: return None -TYPE_MAPPING = {"observable": Observable, "observables": Observable} - - -# 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, -# ) +TYPE_MAPPING = {"observable": Observable, "observables": Observable} \ No newline at end of file From 9ed714ba41a105630419c2716b39382f3db92492 Mon Sep 17 00:00:00 2001 From: Fred Baguelin Date: Mon, 16 Sep 2024 22:50:43 +0200 Subject: [PATCH 04/39] Add dependencies to poetry / pyproject --- poetry.lock | 14 +++++++++++++- pyproject.toml | 1 + 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/poetry.lock b/poetry.lock index d1ef79290..83f4943a0 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,5 +1,17 @@ # This file is automatically @generated by Poetry 1.8.3 and should not be changed by hand. +[[package]] +name = "aenum" +version = "3.1.15" +description = "Advanced Enumerations (compatible with Python's stdlib Enum), NamedTuples, and NamedConstants" +optional = false +python-versions = "*" +files = [ + {file = "aenum-3.1.15-py2-none-any.whl", hash = "sha256:27b1710b9d084de6e2e695dab78fe9f269de924b51ae2850170ee7e1ca6288a5"}, + {file = "aenum-3.1.15-py3-none-any.whl", hash = "sha256:e0dfaeea4c2bd362144b87377e2c61d91958c5ed0b4daf89cb6f45ae23af6288"}, + {file = "aenum-3.1.15.tar.gz", hash = "sha256:8cbd76cd18c4f870ff39b24284d3ea028fbe8731a58df3aa581e434c575b9559"}, +] + [[package]] name = "altair" version = "5.3.0" @@ -2714,4 +2726,4 @@ test = ["big-O", "importlib-resources", "jaraco.functools", "jaraco.itertools", [metadata] lock-version = "2.0" python-versions = ">=3.10,<3.12" -content-hash = "9faf47cc18eca6baf02abc450a8037b5dbbcb1977db0d241c2003e7f6b1f7710" +content-hash = "f5ab0f88602f650e5d643833d31e44764daa404b06b28a3353a5bd58357f6ba1" diff --git a/pyproject.toml b/pyproject.toml index 480085d47..42929086f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,6 +26,7 @@ pyyaml = "^6.0.1" parameterized = "^0.9.0" yara-python = "^4.5.0" idstools = "^0.6.5" +aenum = "^3.1.15" [tool.poetry.group.dev.dependencies] pylint = "^2.16.1" From b15222ea84b27aed098ffa258dcfb5227ddcc5d1 Mon Sep 17 00:00:00 2001 From: Fred Baguelin Date: Tue, 17 Sep 2024 00:23:20 +0200 Subject: [PATCH 05/39] Add load_observables function and handle guess type --- core/schemas/__init__.py | 39 +++++++++++++++++++++++---------------- 1 file changed, 23 insertions(+), 16 deletions(-) diff --git a/core/schemas/__init__.py b/core/schemas/__init__.py index 8d727e8fd..e75d333af 100644 --- a/core/schemas/__init__.py +++ b/core/schemas/__init__.py @@ -1,25 +1,32 @@ import importlib import inspect +import logging from pathlib import Path import aenum from core.schemas import observable -print("Registering observable types") +logger = logging.getLogger(__name__) -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 +logger.info("Registering observable types") + +def load_observables(): + for observable_file in Path(__file__).parent.glob("observables/**/*.py"): + if observable_file.stem == "__init__": + continue + logger.info(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 + aenum.extend_enum(observable.ObservableType, "guess", "guess") + +load_observables() \ No newline at end of file From 888ede4a8557659da0f3c290bccd8f8ffe3c821d Mon Sep 17 00:00:00 2001 From: Fred Baguelin Date: Tue, 17 Sep 2024 00:25:04 +0200 Subject: [PATCH 06/39] Rename GenericObservable with Generic in tests --- tests/schemas/fixture.py | 6 +++--- tests/schemas/observable.py | 8 ++++---- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/tests/schemas/fixture.py b/tests/schemas/fixture.py index 950298ee0..774a16879 100644 --- a/tests/schemas/fixture.py +++ b/tests/schemas/fixture.py @@ -5,7 +5,7 @@ from core.schemas.indicator import DiamondModel, Query, Regex from core.schemas.observables import ( bic, - generic_observable, + generic, hostname, iban, ipv4, @@ -40,10 +40,10 @@ def test_something(self): hacker = hostname.Hostname(value="hacker.com").save() sus_hacker = hostname.Hostname(value="sus.hacker.com").save() mac_address.MacAddress(value="00:11:22:33:44:55").save() - generic = generic_observable.GenericObservable( + generic_obs = generic.Generic( value="SomeInterestingString" ).save() - generic.add_context("test_source", {"test": "test"}) + generic_obs.add_context("test_source", {"test": "test"}) hacker.link_to(www_hacker, "domain", "Domain") hacker.link_to(c2_hacker, "domain", "Domain") diff --git a/tests/schemas/observable.py b/tests/schemas/observable.py index dae5e0d2d..d97a7b675 100644 --- a/tests/schemas/observable.py +++ b/tests/schemas/observable.py @@ -13,7 +13,7 @@ docker_image, email, file, - generic_observable, + generic, hostname, iban, imphash, @@ -74,7 +74,7 @@ def test_observable_update(self) -> None: self.assertEqual(result.context[0], {"source": "source1", "some": "info"}) def test_create_generic_observable(self): - result = generic_observable.GenericObservable(value="Some_String").save() + result = generic.Generic(value="Some_String").save() self.assertIsNotNone(result.id) self.assertEqual(result.value, "Some_String") self.assertEqual(result.type, "generic") @@ -465,10 +465,10 @@ def test_create_sha256(self) -> None: def test_create_ssdeep(self) -> None: """Tests creating an ssdeep.""" - observable = ssdeep.SsdeepHash(value="1234567890").save() + observable = ssdeep.Ssdeep(value="1234567890").save() self.assertIsNotNone(observable.id) self.assertEqual(observable.value, "1234567890") - self.assertIsInstance(observable, ssdeep.SsdeepHash) + self.assertIsInstance(observable, ssdeep.Ssdeep) def test_create_tlsh(self) -> None: """Tests creating a TLSH.""" From 2eb59c445f13c0e04322d127033d307671ccb5c5 Mon Sep 17 00:00:00 2001 From: Fred Baguelin Date: Tue, 17 Sep 2024 00:26:49 +0200 Subject: [PATCH 07/39] Remove validators and replace with static method is_valid --- core/schemas/observable.py | 62 +++++--------------------------------- 1 file changed, 7 insertions(+), 55 deletions(-) diff --git a/core/schemas/observable.py b/core/schemas/observable.py index e0e4c5f43..c476e7f9c 100644 --- a/core/schemas/observable.py +++ b/core/schemas/observable.py @@ -4,12 +4,11 @@ import re # from enum import Enum, EnumMeta -from typing import ClassVar, Literal +from typing import Any, ClassVar, Literal # Data Schema # Dynamically register all observable types import aenum -import validators from pydantic import Field, computed_field from core import database_arango @@ -20,10 +19,8 @@ class ObservableType(str, aenum.Enum): pass - TYPE_MAPPING = {} - class Observable(YetiTagModel, database_arango.ArangoYetiConnector): _collection_name: ClassVar[str] = "observables" _type_filter: ClassVar[str | None] = None @@ -46,9 +43,9 @@ def load(cls, object: dict) -> "ObservableTypes": # noqa: F821 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) + @staticmethod + def is_valid(value: Any) -> bool: + return False @classmethod def add_text(cls, text: str, tags: list[str] = []) -> "ObservableTypes": # noqa: F821 @@ -129,53 +126,8 @@ def delete_context( TYPE_MAPPING.update({"observable": Observable, "observables": Observable}) - -TYPE_VALIDATOR_MAP = {} - - -TYPE_VALIDATOR_MAP = { - ObservableType.ipv4: validators.ipv4, - ObservableType.ipv6: validators.ipv6, - ObservableType.sha256: validators.sha256, - ObservableType.sha1: validators.sha1, - ObservableType.md5: validators.md5, - ObservableType.hostname: validators.domain, - ObservableType.url: validators.url, - ObservableType.email: validators.email, - ObservableType.iban: validators.iban, -} - -REGEXES_OBSERVABLES = { - # Unix - ObservableType.path: [ - re.compile(r"^(\/[^\/\0]+)+$"), - re.compile(r"^(?:[a-zA-Z]\:|\\\\[\w\.]+\\[\w.$]+)\\(?:[\w]+\\)*\w([\w.])+"), - ], - ObservableType.bic: [re.compile("^[A-Z]{6}[A-Z0-9]{2}[A-Z0-9]{3}?")], -} - - -def validate_observable(obs: Observable) -> bool: - if obs.type in TYPE_VALIDATOR_MAP: - return TYPE_VALIDATOR_MAP[obs.type](obs.value) is True - elif obs.type in dict(REGEXES_OBSERVABLES): - for regex in REGEXES_OBSERVABLES[obs.type]: - if regex.match(obs.value): - return True - return False - else: - return False - - def find_type(value: str) -> ObservableType | None: - for obs_type, validator in TYPE_VALIDATOR_MAP.items(): - if validator(value): + for obs_type, obj in TYPE_MAPPING.items(): + if obj.is_valid(value): return obs_type - for obs_type, regexes in REGEXES_OBSERVABLES.items(): - for regex in regexes: - if regex.match(value): - return obs_type - return None - - -TYPE_MAPPING = {"observable": Observable, "observables": Observable} \ No newline at end of file + return None \ No newline at end of file From 1c5e3a9d0f05c1f098d643e84c11b96bef76c9c3 Mon Sep 17 00:00:00 2001 From: Fred Baguelin Date: Tue, 17 Sep 2024 00:29:39 +0200 Subject: [PATCH 08/39] Implement validators for previously handled observables --- core/schemas/observables/bic.py | 8 ++++++++ core/schemas/observables/email.py | 6 ++++++ core/schemas/observables/hostname.py | 6 ++++++ core/schemas/observables/iban.py | 6 ++++++ core/schemas/observables/ipv4.py | 6 ++++++ core/schemas/observables/ipv6.py | 6 ++++++ core/schemas/observables/md5.py | 6 ++++++ core/schemas/observables/path.py | 11 +++++++++++ core/schemas/observables/sha1.py | 6 ++++++ core/schemas/observables/sha256.py | 6 ++++++ core/schemas/observables/url.py | 6 ++++++ 11 files changed, 73 insertions(+) diff --git a/core/schemas/observables/bic.py b/core/schemas/observables/bic.py index c6aca2c69..68b7c8ea3 100644 --- a/core/schemas/observables/bic.py +++ b/core/schemas/observables/bic.py @@ -1,5 +1,13 @@ +import re + from core.schemas import observable +bic_matcher = re.compile("^[A-Z]{6}[A-Z0-9]{2}[A-Z0-9]{3}?") + class BIC(observable.Observable): type: observable.ObservableType = observable.ObservableType.bic + + @staticmethod + def is_valid(value: str) -> bool: + return bic_matcher.match(value) diff --git a/core/schemas/observables/email.py b/core/schemas/observables/email.py index b85bfebd6..b1c38491f 100644 --- a/core/schemas/observables/email.py +++ b/core/schemas/observables/email.py @@ -1,5 +1,11 @@ +import validators + from core.schemas import observable class Email(observable.Observable): type: observable.ObservableType = observable.ObservableType.email + + @staticmethod + def is_valid(value: str) -> bool: + return validators.email(value) \ No newline at end of file diff --git a/core/schemas/observables/hostname.py b/core/schemas/observables/hostname.py index 0eb452b55..9c443e53d 100644 --- a/core/schemas/observables/hostname.py +++ b/core/schemas/observables/hostname.py @@ -1,5 +1,11 @@ +import validators + from core.schemas import observable class Hostname(observable.Observable): type: observable.ObservableType = observable.ObservableType.hostname + + @staticmethod + def is_valid(value: str) -> bool: + return validators.domain(value) \ No newline at end of file diff --git a/core/schemas/observables/iban.py b/core/schemas/observables/iban.py index f963dc3bb..f3c477016 100644 --- a/core/schemas/observables/iban.py +++ b/core/schemas/observables/iban.py @@ -1,5 +1,11 @@ +import validators + from core.schemas import observable class IBAN(observable.Observable): type: observable.ObservableType = observable.ObservableType.iban + + @staticmethod + def is_valid(value: str) -> bool: + return validators.iban(value) \ No newline at end of file diff --git a/core/schemas/observables/ipv4.py b/core/schemas/observables/ipv4.py index 0520e0b03..760e3cd90 100644 --- a/core/schemas/observables/ipv4.py +++ b/core/schemas/observables/ipv4.py @@ -1,5 +1,11 @@ +import validators + from core.schemas import observable class IPv4(observable.Observable): type: observable.ObservableType = observable.ObservableType.ipv4 + + @staticmethod + def is_valid(value: str) -> bool: + return validators.ipv4(value) \ No newline at end of file diff --git a/core/schemas/observables/ipv6.py b/core/schemas/observables/ipv6.py index a264310f4..2daa2cbdc 100644 --- a/core/schemas/observables/ipv6.py +++ b/core/schemas/observables/ipv6.py @@ -1,5 +1,11 @@ +import validators + from core.schemas import observable class IPv6(observable.Observable): type: observable.ObservableType = observable.ObservableType.ipv6 + + @staticmethod + def is_valid(value: str) -> bool: + return validators.ipv6(value) \ No newline at end of file diff --git a/core/schemas/observables/md5.py b/core/schemas/observables/md5.py index 7281a1beb..d7849667a 100644 --- a/core/schemas/observables/md5.py +++ b/core/schemas/observables/md5.py @@ -1,5 +1,11 @@ +import validators + from core.schemas import observable class MD5(observable.Observable): type: observable.ObservableType = observable.ObservableType.md5 + + @staticmethod + def is_valid(value: str) -> bool: + return validators.md5(value) \ No newline at end of file diff --git a/core/schemas/observables/path.py b/core/schemas/observables/path.py index 7fb227845..12621eddf 100644 --- a/core/schemas/observables/path.py +++ b/core/schemas/observables/path.py @@ -1,5 +1,16 @@ +import re + from core.schemas import observable +linux_path_matcher = re.compile(r"^(\/[^\/\0]+)+$") +windows_path_matcher = re.compile(r"^(?:[a-zA-Z]\:|\\\\[\w\.]+\\[\w.$]+)\\(?:[\w]+\\)*\w([\w.])+") + +def path_validator(value): + return linux_path_matcher.match(value) or windows_path_matcher.match(value) class Path(observable.Observable): type: observable.ObservableType = observable.ObservableType.path + + @staticmethod + def is_valid(value: str) -> bool: + return path_validator(value) \ No newline at end of file diff --git a/core/schemas/observables/sha1.py b/core/schemas/observables/sha1.py index 33cb0edcb..a86f6a487 100644 --- a/core/schemas/observables/sha1.py +++ b/core/schemas/observables/sha1.py @@ -1,5 +1,11 @@ +import validators + from core.schemas import observable class SHA1(observable.Observable): type: observable.ObservableType = observable.ObservableType.sha1 + + @staticmethod + def is_valid(value: str) -> bool: + return validators.sha1(value) \ No newline at end of file diff --git a/core/schemas/observables/sha256.py b/core/schemas/observables/sha256.py index f7f452172..18171f1e6 100644 --- a/core/schemas/observables/sha256.py +++ b/core/schemas/observables/sha256.py @@ -1,5 +1,11 @@ +import validators + from core.schemas import observable class SHA256(observable.Observable): type: observable.ObservableType = observable.ObservableType.sha256 + + @staticmethod + def is_valid(value: str) -> bool: + return validators.sha256(value) \ No newline at end of file diff --git a/core/schemas/observables/url.py b/core/schemas/observables/url.py index accdbb507..594dbd74d 100644 --- a/core/schemas/observables/url.py +++ b/core/schemas/observables/url.py @@ -1,5 +1,11 @@ +import validators + from core.schemas import observable class Url(observable.Observable): type: observable.ObservableType = observable.ObservableType.url + + @staticmethod + def is_valid(value: str) -> bool: + return validators.url(value) \ No newline at end of file From df542ff6ff9e8f85dd6eb7efe51ca71502db6edf Mon Sep 17 00:00:00 2001 From: Fred Baguelin Date: Tue, 17 Sep 2024 00:30:56 +0200 Subject: [PATCH 09/39] Ruff format --- core/schemas/__init__.py | 4 +++- core/schemas/observable.py | 5 ++++- core/schemas/observables/email.py | 2 +- core/schemas/observables/hostname.py | 2 +- core/schemas/observables/iban.py | 2 +- core/schemas/observables/ipv4.py | 2 +- core/schemas/observables/ipv6.py | 2 +- core/schemas/observables/md5.py | 2 +- core/schemas/observables/path.py | 8 ++++++-- core/schemas/observables/sha1.py | 2 +- core/schemas/observables/sha256.py | 2 +- core/schemas/observables/url.py | 2 +- tests/apiv2/templates.py | 12 ++++++------ tests/schemas/fixture.py | 4 +--- 14 files changed, 29 insertions(+), 22 deletions(-) diff --git a/core/schemas/__init__.py b/core/schemas/__init__.py index e75d333af..b0cc0fb4d 100644 --- a/core/schemas/__init__.py +++ b/core/schemas/__init__.py @@ -11,6 +11,7 @@ logger.info("Registering observable types") + def load_observables(): for observable_file in Path(__file__).parent.glob("observables/**/*.py"): if observable_file.stem == "__init__": @@ -29,4 +30,5 @@ def load_observables(): observable.TYPE_MAPPING[observable_file.stem] = obj aenum.extend_enum(observable.ObservableType, "guess", "guess") -load_observables() \ No newline at end of file + +load_observables() diff --git a/core/schemas/observable.py b/core/schemas/observable.py index c476e7f9c..8f1392568 100644 --- a/core/schemas/observable.py +++ b/core/schemas/observable.py @@ -19,8 +19,10 @@ class ObservableType(str, aenum.Enum): pass + TYPE_MAPPING = {} + class Observable(YetiTagModel, database_arango.ArangoYetiConnector): _collection_name: ClassVar[str] = "observables" _type_filter: ClassVar[str | None] = None @@ -126,8 +128,9 @@ def delete_context( TYPE_MAPPING.update({"observable": Observable, "observables": Observable}) + def find_type(value: str) -> ObservableType | None: for obs_type, obj in TYPE_MAPPING.items(): if obj.is_valid(value): return obs_type - return None \ No newline at end of file + return None diff --git a/core/schemas/observables/email.py b/core/schemas/observables/email.py index b1c38491f..ba619ec3f 100644 --- a/core/schemas/observables/email.py +++ b/core/schemas/observables/email.py @@ -8,4 +8,4 @@ class Email(observable.Observable): @staticmethod def is_valid(value: str) -> bool: - return validators.email(value) \ No newline at end of file + return validators.email(value) diff --git a/core/schemas/observables/hostname.py b/core/schemas/observables/hostname.py index 9c443e53d..db0b892a6 100644 --- a/core/schemas/observables/hostname.py +++ b/core/schemas/observables/hostname.py @@ -8,4 +8,4 @@ class Hostname(observable.Observable): @staticmethod def is_valid(value: str) -> bool: - return validators.domain(value) \ No newline at end of file + return validators.domain(value) diff --git a/core/schemas/observables/iban.py b/core/schemas/observables/iban.py index f3c477016..9641080a2 100644 --- a/core/schemas/observables/iban.py +++ b/core/schemas/observables/iban.py @@ -8,4 +8,4 @@ class IBAN(observable.Observable): @staticmethod def is_valid(value: str) -> bool: - return validators.iban(value) \ No newline at end of file + return validators.iban(value) diff --git a/core/schemas/observables/ipv4.py b/core/schemas/observables/ipv4.py index 760e3cd90..9eab620c3 100644 --- a/core/schemas/observables/ipv4.py +++ b/core/schemas/observables/ipv4.py @@ -8,4 +8,4 @@ class IPv4(observable.Observable): @staticmethod def is_valid(value: str) -> bool: - return validators.ipv4(value) \ No newline at end of file + return validators.ipv4(value) diff --git a/core/schemas/observables/ipv6.py b/core/schemas/observables/ipv6.py index 2daa2cbdc..f1fae0be9 100644 --- a/core/schemas/observables/ipv6.py +++ b/core/schemas/observables/ipv6.py @@ -8,4 +8,4 @@ class IPv6(observable.Observable): @staticmethod def is_valid(value: str) -> bool: - return validators.ipv6(value) \ No newline at end of file + return validators.ipv6(value) diff --git a/core/schemas/observables/md5.py b/core/schemas/observables/md5.py index d7849667a..9e783c55e 100644 --- a/core/schemas/observables/md5.py +++ b/core/schemas/observables/md5.py @@ -8,4 +8,4 @@ class MD5(observable.Observable): @staticmethod def is_valid(value: str) -> bool: - return validators.md5(value) \ No newline at end of file + return validators.md5(value) diff --git a/core/schemas/observables/path.py b/core/schemas/observables/path.py index 12621eddf..5d90025e3 100644 --- a/core/schemas/observables/path.py +++ b/core/schemas/observables/path.py @@ -3,14 +3,18 @@ from core.schemas import observable linux_path_matcher = re.compile(r"^(\/[^\/\0]+)+$") -windows_path_matcher = re.compile(r"^(?:[a-zA-Z]\:|\\\\[\w\.]+\\[\w.$]+)\\(?:[\w]+\\)*\w([\w.])+") +windows_path_matcher = re.compile( + r"^(?:[a-zA-Z]\:|\\\\[\w\.]+\\[\w.$]+)\\(?:[\w]+\\)*\w([\w.])+" +) + def path_validator(value): return linux_path_matcher.match(value) or windows_path_matcher.match(value) + class Path(observable.Observable): type: observable.ObservableType = observable.ObservableType.path @staticmethod def is_valid(value: str) -> bool: - return path_validator(value) \ No newline at end of file + return path_validator(value) diff --git a/core/schemas/observables/sha1.py b/core/schemas/observables/sha1.py index a86f6a487..db60b5f09 100644 --- a/core/schemas/observables/sha1.py +++ b/core/schemas/observables/sha1.py @@ -8,4 +8,4 @@ class SHA1(observable.Observable): @staticmethod def is_valid(value: str) -> bool: - return validators.sha1(value) \ No newline at end of file + return validators.sha1(value) diff --git a/core/schemas/observables/sha256.py b/core/schemas/observables/sha256.py index 18171f1e6..51910b03c 100644 --- a/core/schemas/observables/sha256.py +++ b/core/schemas/observables/sha256.py @@ -8,4 +8,4 @@ class SHA256(observable.Observable): @staticmethod def is_valid(value: str) -> bool: - return validators.sha256(value) \ No newline at end of file + return validators.sha256(value) diff --git a/core/schemas/observables/url.py b/core/schemas/observables/url.py index 594dbd74d..25c7ceb4f 100644 --- a/core/schemas/observables/url.py +++ b/core/schemas/observables/url.py @@ -8,4 +8,4 @@ class Url(observable.Observable): @staticmethod def is_valid(value: str) -> bool: - return validators.url(value) \ No newline at end of file + return validators.url(value) diff --git a/tests/apiv2/templates.py b/tests/apiv2/templates.py index d5e9e2c4a..e1db94b4c 100644 --- a/tests/apiv2/templates.py +++ b/tests/apiv2/templates.py @@ -90,9 +90,9 @@ def test_render_template_by_id(self): }, ) data = response.text - response.headers["Content-Disposition"] = ( - "attachment; filename=FakeTemplate.txt" - ) + response.headers[ + "Content-Disposition" + ] = "attachment; filename=FakeTemplate.txt" self.assertEqual(response.status_code, 200, data) self.assertEqual(data, "\n1.1.1.1\n2.2.2.2\n3.3.3.3\n\n\n") @@ -106,8 +106,8 @@ def test_render_template_by_search(self): json={"template_id": self.template.id, "search_query": "yeti"}, ) data = response.text - response.headers["Content-Disposition"] = ( - "attachment; filename=FakeTemplate.txt" - ) + response.headers[ + "Content-Disposition" + ] = "attachment; filename=FakeTemplate.txt" self.assertEqual(response.status_code, 200, data) self.assertEqual(data, "\nyeti1.com\nyeti2.com\nyeti3.com\n\n\n") diff --git a/tests/schemas/fixture.py b/tests/schemas/fixture.py index 774a16879..4ba63a045 100644 --- a/tests/schemas/fixture.py +++ b/tests/schemas/fixture.py @@ -40,9 +40,7 @@ def test_something(self): hacker = hostname.Hostname(value="hacker.com").save() sus_hacker = hostname.Hostname(value="sus.hacker.com").save() mac_address.MacAddress(value="00:11:22:33:44:55").save() - generic_obs = generic.Generic( - value="SomeInterestingString" - ).save() + generic_obs = generic.Generic(value="SomeInterestingString").save() generic_obs.add_context("test_source", {"test": "test"}) hacker.link_to(www_hacker, "domain", "Domain") From bb0faeb8a9c47e8b0f92440c67bef594fab08e02 Mon Sep 17 00:00:00 2001 From: Fred Baguelin Date: Tue, 17 Sep 2024 02:23:33 +0200 Subject: [PATCH 10/39] Add entity loader --- core/schemas/__init__.py | 43 +++++++++++++++++++++++++++++++++------- 1 file changed, 36 insertions(+), 7 deletions(-) diff --git a/core/schemas/__init__.py b/core/schemas/__init__.py index b0cc0fb4d..a86d5dbc2 100644 --- a/core/schemas/__init__.py +++ b/core/schemas/__init__.py @@ -5,14 +5,12 @@ import aenum -from core.schemas import observable +from core.schemas import entity, observable logger = logging.getLogger(__name__) -logger.info("Registering observable types") - - def load_observables(): + logger.info("Registering observable types") for observable_file in Path(__file__).parent.glob("observables/**/*.py"): if observable_file.stem == "__init__": continue @@ -21,9 +19,10 @@ def load_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 - ) + if observable_file.stem not in observable.ObservableType.__members__: + 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): @@ -31,4 +30,34 @@ def load_observables(): aenum.extend_enum(observable.ObservableType, "guess", "guess") +def load_entities(): + logger.info("Registering entity types") + for entity_file in Path(__file__).parent.glob("entities/**/*.py"): + if entity_file.stem == "__init__": + continue + logger.info(f"Registering entity type {entity_file.stem}") + if entity_file.parent.stem == "entities": + module_name = f"core.schemas.entities.{entity_file.stem}" + elif entity_file.parent.stem == "private": + module_name = f"core.schemas.entities.private.{entity_file.stem}" + enum_value = entity_file.stem.replace("_", "-") + if entity_file.stem not in entity.EntityType.__members__: + aenum.extend_enum( + entity.EntityType, entity_file.stem, enum_value + ) + module = importlib.import_module(module_name) + for _, obj in inspect.getmembers(module, inspect.isclass): + if issubclass(obj, entity.Entity): + entity.TYPE_MAPPING[enum_value] = obj + for key in entity.TYPE_MAPPING: + if key in ["entity", "entities"]: + continue + cls = entity.TYPE_MAPPING[key] + if not entity.EntityTypes: + entity.EntityTypes = cls + else: + entity.EntityTypes |= cls + + load_observables() +load_entities() From b03b7a2eeb7948935f70ad983d1372bb7839a5e7 Mon Sep 17 00:00:00 2001 From: Fred Baguelin Date: Tue, 17 Sep 2024 02:24:37 +0200 Subject: [PATCH 11/39] Move entities under entities folder --- core/schemas/entities/__init__.py | 0 core/schemas/entities/attack_pattern.py | 10 + core/schemas/entities/campaign.py | 16 ++ core/schemas/entities/company.py | 8 + core/schemas/entities/course_of_action.py | 8 + core/schemas/entities/identity.py | 12 + core/schemas/entities/intrusion_set.py | 16 ++ core/schemas/entities/investigation.py | 10 + core/schemas/entities/malware.py | 12 + core/schemas/entities/note.py | 8 + core/schemas/entities/phone.py | 8 + core/schemas/entities/private/.gitignore | 3 + core/schemas/entities/private/README.md | 2 + core/schemas/entities/threat_actor.py | 17 ++ core/schemas/entities/tool.py | 12 + core/schemas/entities/vulnerability.py | 43 ++++ core/schemas/entity.py | 253 ++++------------------ 17 files changed, 230 insertions(+), 208 deletions(-) create mode 100644 core/schemas/entities/__init__.py create mode 100644 core/schemas/entities/attack_pattern.py create mode 100644 core/schemas/entities/campaign.py create mode 100644 core/schemas/entities/company.py create mode 100644 core/schemas/entities/course_of_action.py create mode 100644 core/schemas/entities/identity.py create mode 100644 core/schemas/entities/intrusion_set.py create mode 100644 core/schemas/entities/investigation.py create mode 100644 core/schemas/entities/malware.py create mode 100644 core/schemas/entities/note.py create mode 100644 core/schemas/entities/phone.py create mode 100644 core/schemas/entities/private/.gitignore create mode 100644 core/schemas/entities/private/README.md create mode 100644 core/schemas/entities/threat_actor.py create mode 100644 core/schemas/entities/tool.py create mode 100644 core/schemas/entities/vulnerability.py diff --git a/core/schemas/entities/__init__.py b/core/schemas/entities/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/core/schemas/entities/attack_pattern.py b/core/schemas/entities/attack_pattern.py new file mode 100644 index 000000000..6074257ba --- /dev/null +++ b/core/schemas/entities/attack_pattern.py @@ -0,0 +1,10 @@ +from typing import ClassVar + +from core.schemas import entity + + +class AttackPattern(entity.Entity): + _type_filter: ClassVar[str] = entity.EntityType.attack_pattern + type: entity.EntityType = entity.EntityType.attack_pattern + aliases: list[str] = [] + kill_chain_phases: list[str] = [] diff --git a/core/schemas/entities/campaign.py b/core/schemas/entities/campaign.py new file mode 100644 index 000000000..16012ea9a --- /dev/null +++ b/core/schemas/entities/campaign.py @@ -0,0 +1,16 @@ +import datetime +from typing import ClassVar + +from pydantic import Field + +from core.helpers import now +from core.schemas import entity + + +class Campaign(entity.Entity): + _type_filter: ClassVar[str] = entity.EntityType.campaign + type: entity.EntityType = entity.EntityType.campaign + + aliases: list[str] = [] + first_seen: datetime.datetime = Field(default_factory=now) + last_seen: datetime.datetime = Field(default_factory=now) diff --git a/core/schemas/entities/company.py b/core/schemas/entities/company.py new file mode 100644 index 000000000..2e3f60c68 --- /dev/null +++ b/core/schemas/entities/company.py @@ -0,0 +1,8 @@ +from typing import ClassVar + +from core.schemas import entity + + +class Company(entity.Entity): + type: entity.EntityType = entity.EntityType.company + _type_filter: ClassVar[str] = entity.EntityType.company diff --git a/core/schemas/entities/course_of_action.py b/core/schemas/entities/course_of_action.py new file mode 100644 index 000000000..395ee1073 --- /dev/null +++ b/core/schemas/entities/course_of_action.py @@ -0,0 +1,8 @@ +from typing import ClassVar + +from core.schemas import entity + + +class CourseOfAction(entity.Entity): + _type_filter: ClassVar[str] = entity.EntityType.course_of_action + type: entity.EntityType = entity.EntityType.course_of_action diff --git a/core/schemas/entities/identity.py b/core/schemas/entities/identity.py new file mode 100644 index 000000000..d3f5b20c8 --- /dev/null +++ b/core/schemas/entities/identity.py @@ -0,0 +1,12 @@ +from typing import ClassVar + +from core.schemas import entity + + +class Identity(entity.Entity): + _type_filter: ClassVar[str] = entity.EntityType.identity + type: entity.EntityType = entity.EntityType.identity + + identity_class: str = "" + sectors: list[str] = [] + contact_information: str = "" diff --git a/core/schemas/entities/intrusion_set.py b/core/schemas/entities/intrusion_set.py new file mode 100644 index 000000000..35dfa4501 --- /dev/null +++ b/core/schemas/entities/intrusion_set.py @@ -0,0 +1,16 @@ +import datetime +from typing import ClassVar + +from pydantic import Field + +from core.helpers import now +from core.schemas import entity + + +class IntrusionSet(entity.Entity): + _type_filter: ClassVar[str] = entity.EntityType.intrusion_set + type: entity.EntityType = entity.EntityType.intrusion_set + + aliases: list[str] = [] + first_seen: datetime.datetime = Field(default_factory=now) + last_seen: datetime.datetime = Field(default_factory=now) diff --git a/core/schemas/entities/investigation.py b/core/schemas/entities/investigation.py new file mode 100644 index 000000000..b993f6a9b --- /dev/null +++ b/core/schemas/entities/investigation.py @@ -0,0 +1,10 @@ +from typing import ClassVar + +from core.schemas import entity + + +class Investigation(entity.Entity): + _type_filter: ClassVar[str] = entity.EntityType.investigation + type: entity.EntityType = entity.EntityType.investigation + + reference: str = "" diff --git a/core/schemas/entities/malware.py b/core/schemas/entities/malware.py new file mode 100644 index 000000000..dbdcdaafa --- /dev/null +++ b/core/schemas/entities/malware.py @@ -0,0 +1,12 @@ +from typing import ClassVar + +from core.schemas import entity + + +class Malware(entity.Entity): + _type_filter: ClassVar[str] = entity.EntityType.malware + type: entity.EntityType = entity.EntityType.malware + + kill_chain_phases: list[str] = [] + aliases: list[str] = [] + family: str = "" diff --git a/core/schemas/entities/note.py b/core/schemas/entities/note.py new file mode 100644 index 000000000..706e42e85 --- /dev/null +++ b/core/schemas/entities/note.py @@ -0,0 +1,8 @@ +from typing import ClassVar + +from core.schemas import entity + + +class Note(entity.Entity): + type: entity.EntityType = entity.EntityType.note + _type_filter: ClassVar[str] = entity.EntityType.note diff --git a/core/schemas/entities/phone.py b/core/schemas/entities/phone.py new file mode 100644 index 000000000..809eb3745 --- /dev/null +++ b/core/schemas/entities/phone.py @@ -0,0 +1,8 @@ +from typing import ClassVar + +from core.schemas import entity + + +class Phone(entity.Entity): + type: entity.EntityType = entity.EntityType.phone + _type_filter: ClassVar[str] = entity.EntityType.phone diff --git a/core/schemas/entities/private/.gitignore b/core/schemas/entities/private/.gitignore new file mode 100644 index 000000000..e5af87e9b --- /dev/null +++ b/core/schemas/entities/private/.gitignore @@ -0,0 +1,3 @@ +* +!.gitignore +!README.md \ No newline at end of file diff --git a/core/schemas/entities/private/README.md b/core/schemas/entities/private/README.md new file mode 100644 index 000000000..4f9a2488b --- /dev/null +++ b/core/schemas/entities/private/README.md @@ -0,0 +1,2 @@ +### Private entities +This directory is where you should place your private entities. It could be named anything else, but this one has a `.gitignore` so you don't mess things up. ;-) diff --git a/core/schemas/entities/threat_actor.py b/core/schemas/entities/threat_actor.py new file mode 100644 index 000000000..2f9841d00 --- /dev/null +++ b/core/schemas/entities/threat_actor.py @@ -0,0 +1,17 @@ +import datetime +from typing import ClassVar + +from pydantic import Field + +from core.helpers import now +from core.schemas import entity + + +class ThreatActor(entity.Entity): + _type_filter: ClassVar[str] = entity.EntityType.threat_actor + type: entity.EntityType = entity.EntityType.threat_actor + + threat_actor_types: list[str] = [] + aliases: list[str] = [] + first_seen: datetime.datetime = Field(default_factory=now) + last_seen: datetime.datetime = Field(default_factory=now) diff --git a/core/schemas/entities/tool.py b/core/schemas/entities/tool.py new file mode 100644 index 000000000..c184561f7 --- /dev/null +++ b/core/schemas/entities/tool.py @@ -0,0 +1,12 @@ +from typing import ClassVar + +from core.schemas import entity + + +class Tool(entity.Entity): + _type_filter: ClassVar[str] = entity.EntityType.tool + type: entity.EntityType = entity.EntityType.tool + + aliases: list[str] = [] + kill_chain_phases: list[str] = [] + tool_version: str = "" diff --git a/core/schemas/entities/vulnerability.py b/core/schemas/entities/vulnerability.py new file mode 100644 index 000000000..58fbe26e8 --- /dev/null +++ b/core/schemas/entities/vulnerability.py @@ -0,0 +1,43 @@ +import re +from enum import Enum +from typing import ClassVar + +from pydantic import Field + +from core.schemas import entity + +vulnerability_matcher = re.compile(r"(?P
\W?)(?PCVE-\d{4}-\d{4,7})(?P\W?)")
+
+class SeverityType(str, Enum):
+    none = "none"
+    low = "low"
+    medium = "medium"
+    high = "high"
+    critical = "critical"
+
+
+class Vulnerability(entity.Entity):
+    """
+    This class represents a vulnerability in the schema.
+
+    Attributes:
+        title: title of the vulnerability.
+        base_score : base score of the vulnerability obtained from CVSS metric
+                     ranging from 0.0 to 10.0.
+        severity: represents the severity of a vulnerability. One of none, low,
+                  medium, high, critical.
+    """
+
+    _type_filter: ClassVar[str] = entity.EntityType.vulnerability
+    type: entity.EntityType = entity.EntityType.vulnerability
+
+    title: str = ""
+    base_score: float = Field(ge=0.0, le=10.0, default=0.0)
+    severity: SeverityType = "none"
+    reference: str = ""
+
+    @classmethod
+    def is_valid(cls, ent: entity.Entity) -> bool:
+        if vulnerability_matcher.match(ent.name):
+            return True
+        return False
diff --git a/core/schemas/entity.py b/core/schemas/entity.py
index 484d006d2..81be7dbad 100644
--- a/core/schemas/entity.py
+++ b/core/schemas/entity.py
@@ -11,19 +11,48 @@
 
 
 class EntityType(str, Enum):
-    attack_pattern = "attack-pattern"
-    campaign = "campaign"
-    company = "company"
-    identity = "identity"
-    intrusion_set = "intrusion-set"
-    investigation = "investigation"
-    malware = "malware"
-    note = "note"
-    phone = "phone"
-    threat_actor = "threat-actor"
-    tool = "tool"
-    vulnerability = "vulnerability"
-    course_of_action = "course-of-action"
+    pass
+
+TYPE_MAPPING = {}
+
+EntityTypes = ()
+
+# EntityTypes = Annotated[
+#     Union[
+#         AttackPattern,
+#         Campaign,
+#         Company,
+#         CourseOfAction,
+#         Identity,
+#         IntrusionSet,
+#         Investigation,
+#         Malware,
+#         Note,
+#         Phone,
+#         ThreatActor,
+#         Tool,
+#         Vulnerability,
+#     ],
+#     Field(discriminator="type"),
+# ]
+
+
+EntityClasses = (
+    # Type[AttackPattern]
+    # | Type[Campaign]
+    # | Type[Company]
+    # | Type[CourseOfAction]
+    # | Type[Identity]
+    # | Type[IntrusionSet]
+    # | Type[Investigation]
+    # | Type[Malware]
+    # | Type[Note]
+    # | Type[Phone]
+    # | Type[ThreatActor]
+    # | Type[Tool]
+    # | Type[Vulnerability]
+)
+
 
 
 class Entity(YetiTagModel, database_arango.ArangoYetiConnector):
@@ -64,8 +93,8 @@ def load(cls, object: dict) -> "EntityTypes":
         return loader(**object)
 
     @classmethod
-    def is_valid(cls, object: dict) -> bool:
-        return validate_entity(object)
+    def is_valid(cls, object: "Entity") -> bool:
+        return False
 
     def add_context(
         self, source: str, context: dict, skip_compare: set = set()
@@ -89,196 +118,4 @@ def add_context(
         return self.save()
 
 
-class Note(Entity):
-    type: Literal[EntityType.note] = EntityType.note
-    _type_filter: ClassVar[str] = EntityType.note
-
-
-class Phone(Entity):
-    _type_filter: ClassVar[str] = EntityType.phone
-    type: Literal[EntityType.phone] = EntityType.phone
-
-
-class Company(Entity):
-    type: Literal[EntityType.company] = EntityType.company
-    _type_filter: ClassVar[str] = EntityType.company
-
-
-class ThreatActor(Entity):
-    _type_filter: ClassVar[str] = EntityType.threat_actor
-    type: Literal[EntityType.threat_actor] = EntityType.threat_actor
-
-    threat_actor_types: list[str] = []
-    aliases: list[str] = []
-    first_seen: datetime.datetime = Field(default_factory=now)
-    last_seen: datetime.datetime = Field(default_factory=now)
-
-
-class IntrusionSet(Entity):
-    _type_filter: ClassVar[str] = EntityType.intrusion_set
-    type: Literal[EntityType.intrusion_set] = EntityType.intrusion_set
-
-    aliases: list[str] = []
-    first_seen: datetime.datetime = Field(default_factory=now)
-    last_seen: datetime.datetime = Field(default_factory=now)
-
-
-class Tool(Entity):
-    _type_filter: ClassVar[str] = EntityType.tool
-    type: Literal[EntityType.tool] = EntityType.tool
-
-    aliases: list[str] = []
-    kill_chain_phases: list[str] = []
-    tool_version: str = ""
-
-
-class AttackPattern(Entity):
-    _type_filter: ClassVar[str] = EntityType.attack_pattern
-    type: Literal[EntityType.attack_pattern] = EntityType.attack_pattern
-    aliases: list[str] = []
-    kill_chain_phases: list[str] = []
-
-
-class Malware(Entity):
-    _type_filter: ClassVar[str] = EntityType.malware
-    type: Literal[EntityType.malware] = EntityType.malware
-
-    kill_chain_phases: list[str] = []
-    aliases: list[str] = []
-    family: str = ""
-
-
-class Campaign(Entity):
-    _type_filter: ClassVar[str] = EntityType.campaign
-    type: Literal[EntityType.campaign] = EntityType.campaign
-
-    aliases: list[str] = []
-    first_seen: datetime.datetime = Field(default_factory=now)
-    last_seen: datetime.datetime = Field(default_factory=now)
-
-
-class Identity(Entity):
-    _type_filter: ClassVar[str] = EntityType.identity
-    type: Literal[EntityType.identity] = EntityType.identity
-
-    identity_class: str = ""
-    sectors: list[str] = []
-    contact_information: str = ""
-
-
-class Investigation(Entity):
-    _type_filter: ClassVar[str] = EntityType.investigation
-    type: Literal[EntityType.investigation] = EntityType.investigation
-
-    reference: str = ""
-
-
-class SeverityType(str, Enum):
-    none = "none"
-    low = "low"
-    medium = "medium"
-    high = "high"
-    critical = "critical"
-
-
-class Vulnerability(Entity):
-    """
-    This class represents a vulnerability in the schema.
-
-    Attributes:
-        title: title of the vulnerability.
-        base_score : base score of the vulnerability obtained from CVSS metric
-                     ranging from 0.0 to 10.0.
-        severity: represents the severity of a vulnerability. One of none, low,
-                  medium, high, critical.
-    """
-
-    _type_filter: ClassVar[str] = EntityType.vulnerability
-    type: Literal[EntityType.vulnerability] = EntityType.vulnerability
-
-    title: str = ""
-    base_score: float = Field(ge=0.0, le=10.0, default=0.0)
-    severity: SeverityType = "none"
-    reference: str = ""
-
-
-class CourseOfAction(Entity):
-    _type_filter: ClassVar[str] = EntityType.course_of_action
-    type: Literal[EntityType.course_of_action] = EntityType.course_of_action
-
-
-TYPE_MAPPING = {
-    "entities": Entity,
-    "entity": Entity,
-    EntityType.attack_pattern: AttackPattern,
-    EntityType.campaign: Campaign,
-    EntityType.company: Company,
-    EntityType.course_of_action: CourseOfAction,
-    EntityType.identity: Identity,
-    EntityType.intrusion_set: IntrusionSet,
-    EntityType.investigation: Investigation,
-    EntityType.malware: Malware,
-    EntityType.note: Note,
-    EntityType.phone: Phone,
-    EntityType.threat_actor: ThreatActor,
-    EntityType.tool: Tool,
-    EntityType.vulnerability: Vulnerability,
-}
-
-TYPE_VALIDATOR_MAP = {}
-
-REGEXES_ENTITIES = {
-    EntityType.vulnerability: (
-        "name",
-        re.compile(r"(?P
\W?)(?PCVE-\d{4}-\d{4,7})(?P\W?)"),
-    )
-}
-
-
-def validate_entity(ent: Entity) -> bool:
-    if ent.type in TYPE_VALIDATOR_MAP:
-        return TYPE_VALIDATOR_MAP[ent.type](ent) is True
-    elif ent.type in REGEXES_ENTITIES:
-        field, regex = REGEXES_ENTITIES[ent.type]
-        if regex.match(getattr(ent, field)):
-            return True
-        else:
-            return False
-    return True
-
-
-EntityTypes = Annotated[
-    Union[
-        AttackPattern,
-        Campaign,
-        Company,
-        CourseOfAction,
-        Identity,
-        IntrusionSet,
-        Investigation,
-        Malware,
-        Note,
-        Phone,
-        ThreatActor,
-        Tool,
-        Vulnerability,
-    ],
-    Field(discriminator="type"),
-]
-
-
-EntityClasses = (
-    Type[AttackPattern]
-    | Type[Campaign]
-    | Type[Company]
-    | Type[CourseOfAction]
-    | Type[Identity]
-    | Type[IntrusionSet]
-    | Type[Investigation]
-    | Type[Malware]
-    | Type[Note]
-    | Type[Phone]
-    | Type[ThreatActor]
-    | Type[Tool]
-    | Type[Vulnerability]
-)
+TYPE_MAPPING.update({"entities": Entity, "entity": Entity})
\ No newline at end of file

From 16e6d0d517d07e0ba9c9fb857897dca6b1e97f6a Mon Sep 17 00:00:00 2001
From: Fred Baguelin 
Date: Tue, 17 Sep 2024 08:33:48 +0200
Subject: [PATCH 12/39] Use enum.Enum for ObservableType definition

---
 core/schemas/observable.py | 14 ++++++--------
 1 file changed, 6 insertions(+), 8 deletions(-)

diff --git a/core/schemas/observable.py b/core/schemas/observable.py
index 8f1392568..9cd64fa6e 100644
--- a/core/schemas/observable.py
+++ b/core/schemas/observable.py
@@ -3,12 +3,13 @@
 import datetime
 import re
 
+# Data Schema
+# Dynamically register all observable types
+from enum import Enum
+
 # from enum import Enum, EnumMeta
 from typing import Any, ClassVar, Literal
 
-# Data Schema
-# Dynamically register all observable types
-import aenum
 from pydantic import Field, computed_field
 
 from core import database_arango
@@ -16,13 +17,10 @@
 from core.schemas.model import YetiTagModel
 
 
-class ObservableType(str, aenum.Enum):
-    pass
-
-
+# forward declarations
+class ObservableType(str, Enum): ...
 TYPE_MAPPING = {}
 
-
 class Observable(YetiTagModel, database_arango.ArangoYetiConnector):
     _collection_name: ClassVar[str] = "observables"
     _type_filter: ClassVar[str | None] = None

From ce32276851ef8aa9b0930c8cd73f2f3a140f8b51 Mon Sep 17 00:00:00 2001
From: Fred Baguelin 
Date: Tue, 17 Sep 2024 08:34:55 +0200
Subject: [PATCH 13/39] Add guess ObservableType if not already part of enum

---
 core/schemas/__init__.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/core/schemas/__init__.py b/core/schemas/__init__.py
index a86d5dbc2..a8093a8fc 100644
--- a/core/schemas/__init__.py
+++ b/core/schemas/__init__.py
@@ -27,7 +27,8 @@ def load_observables():
         for _, obj in inspect.getmembers(module, inspect.isclass):
             if issubclass(obj, observable.Observable):
                 observable.TYPE_MAPPING[observable_file.stem] = obj
-    aenum.extend_enum(observable.ObservableType, "guess", "guess")
+    if "guess" not in observable.ObservableType.__members__:
+        aenum.extend_enum(observable.ObservableType, "guess", "guess")
 
 
 def load_entities():

From c6f3f8d6fba3343faa2c792b6ad0368474bb7ff3 Mon Sep 17 00:00:00 2001
From: Fred Baguelin 
Date: Tue, 17 Sep 2024 08:35:52 +0200
Subject: [PATCH 14/39] Remove unused declarations

---
 core/schemas/entity.py | 49 ++++--------------------------------------
 1 file changed, 4 insertions(+), 45 deletions(-)

diff --git a/core/schemas/entity.py b/core/schemas/entity.py
index 81be7dbad..6b3f39982 100644
--- a/core/schemas/entity.py
+++ b/core/schemas/entity.py
@@ -1,7 +1,6 @@
 import datetime
-import re
 from enum import Enum
-from typing import Annotated, ClassVar, Literal, Type, Union
+from typing import ClassVar, Literal
 
 from pydantic import Field, computed_field
 
@@ -10,50 +9,10 @@
 from core.schemas.model import YetiTagModel
 
 
-class EntityType(str, Enum):
-    pass
-
-TYPE_MAPPING = {}
-
+# forward declarations
+class EntityType(str, Enum): ...
 EntityTypes = ()
-
-# EntityTypes = Annotated[
-#     Union[
-#         AttackPattern,
-#         Campaign,
-#         Company,
-#         CourseOfAction,
-#         Identity,
-#         IntrusionSet,
-#         Investigation,
-#         Malware,
-#         Note,
-#         Phone,
-#         ThreatActor,
-#         Tool,
-#         Vulnerability,
-#     ],
-#     Field(discriminator="type"),
-# ]
-
-
-EntityClasses = (
-    # Type[AttackPattern]
-    # | Type[Campaign]
-    # | Type[Company]
-    # | Type[CourseOfAction]
-    # | Type[Identity]
-    # | Type[IntrusionSet]
-    # | Type[Investigation]
-    # | Type[Malware]
-    # | Type[Note]
-    # | Type[Phone]
-    # | Type[ThreatActor]
-    # | Type[Tool]
-    # | Type[Vulnerability]
-)
-
-
+TYPE_MAPPING = {}
 
 class Entity(YetiTagModel, database_arango.ArangoYetiConnector):
     _exclude_overwrite: list[str] = ["related_observables_count"]

From 2a9fc2bf1bcb24ca26eeb51d6c423df6e1971183 Mon Sep 17 00:00:00 2001
From: Fred Baguelin 
Date: Tue, 17 Sep 2024 08:36:39 +0200
Subject: [PATCH 15/39] Remove Field(discriminator=type) which is only relevant
 when subclassing models

---
 core/web/apiv2/entities.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/core/web/apiv2/entities.py b/core/web/apiv2/entities.py
index 0597b22cc..42e6c60e6 100644
--- a/core/web/apiv2/entities.py
+++ b/core/web/apiv2/entities.py
@@ -10,14 +10,14 @@
 class NewEntityRequest(BaseModel):
     model_config = ConfigDict(extra="forbid")
 
-    entity: EntityTypes = Field(discriminator="type")
+    entity: EntityTypes
     tags: conlist(str, max_length=MAX_TAGS_REQUEST) = []
 
 
 class PatchEntityRequest(BaseModel):
     model_config = ConfigDict(extra="forbid")
 
-    entity: EntityTypes = Field(discriminator="type")
+    entity: EntityTypes
 
 
 class EntitySearchRequest(BaseModel):

From 71ecd7e3073a3e548cf7ced443d07b72281ba508 Mon Sep 17 00:00:00 2001
From: Fred Baguelin 
Date: Tue, 17 Sep 2024 08:38:02 +0200
Subject: [PATCH 16/39] Update tests to reflect entities changes

---
 tests/apiv2/entities.py  |  9 +++---
 tests/apiv2/graph.py     | 14 ++++-----
 tests/schemas/entity.py  | 62 ++++++++++++++++++++--------------------
 tests/schemas/fixture.py |  8 +++---
 tests/schemas/graph.py   |  9 +++---
 5 files changed, 52 insertions(+), 50 deletions(-)

diff --git a/tests/apiv2/entities.py b/tests/apiv2/entities.py
index 541550a20..458879e15 100644
--- a/tests/apiv2/entities.py
+++ b/tests/apiv2/entities.py
@@ -7,6 +7,7 @@
 
 from core import database_arango
 from core.schemas import entity
+from core.schemas.entities import attack_pattern, malware, threat_actor
 from core.schemas.user import UserSensitive
 from core.web import webapp
 
@@ -23,11 +24,11 @@ def setUp(self) -> None:
             "/api/v2/auth/api-token", headers={"x-yeti-apikey": user.api_key}
         ).json()
         client.headers = {"Authorization": "Bearer " + token_data["access_token"]}
-        self.entity1 = entity.ThreatActor(
+        self.entity1 = threat_actor.ThreatActor(
             name="ta1", aliases=["badactor"], created=datetime.datetime(2020, 1, 1)
         ).save()
         self.entity1.tag(["ta1"])
-        self.entity2 = entity.ThreatActor(name="bears").save()
+        self.entity2 = threat_actor.ThreatActor(name="bears").save()
 
     def tearDown(self) -> None:
         database_arango.db.clear()
@@ -91,8 +92,8 @@ def test_search_entities_subfields(self):
         self.assertEqual(data["entities"][0]["type"], "threat-actor")
 
     def test_search_entities_multiple_types(self):
-        entity.AttackPattern(name="ttp1").save()
-        entity.Malware(name="malware1").save()  # this won't show up
+        attack_pattern.AttackPattern(name="ttp1").save()
+        malware.Malware(name="malware1").save()  # this won't show up
         response = client.post(
             "/api/v2/entities/search",
             json={"query": {"type__in": ["threat-actor", "attack-pattern"]}},
diff --git a/tests/apiv2/graph.py b/tests/apiv2/graph.py
index 53ee43875..639adafb0 100644
--- a/tests/apiv2/graph.py
+++ b/tests/apiv2/graph.py
@@ -5,7 +5,7 @@
 from fastapi.testclient import TestClient
 
 from core import database_arango
-from core.schemas.entity import AttackPattern, Malware, ThreatActor
+from core.schemas.entities import attack_pattern, malware, threat_actor
 from core.schemas.graph import Relationship
 from core.schemas.indicator import DiamondModel, ForensicArtifact, Query, Regex
 from core.schemas.observables import hostname, ipv4, url
@@ -27,7 +27,7 @@ def setUp(self) -> None:
         client.headers = {"Authorization": "Bearer " + token_data["access_token"]}
         self.observable1 = hostname.Hostname(value="tomchop.me").save()
         self.observable2 = ipv4.IPv4(value="127.0.0.1").save()
-        self.entity1 = ThreatActor(name="actor0").save()
+        self.entity1 = threat_actor.ThreatActor(name="actor0").save()
         self.indicator1 = Query(
             name="query1",
             query_type="opensearch",
@@ -283,8 +283,8 @@ def test_neighbors_filter(self):
 
     def test_neighbor_filter_in(self):
         self.entity1.link_to(self.observable1, "uses", "asd")
-        malware = Malware(name="malware1", aliases=["blah"]).save()
-        self.entity1.link_to(malware, "uses", "asd")
+        malware_entity = malware.Malware(name="malware1", aliases=["blah"]).save()
+        self.entity1.link_to(malware_entity, "uses", "asd")
         response = client.post(
             "/api/v2/graph/search",
             json={
@@ -299,7 +299,7 @@ def test_neighbor_filter_in(self):
         data = response.json()
         self.assertEqual(response.status_code, 200, data)
         self.assertEqual(len(data["vertices"]), 1)
-        self.assertEqual(data["vertices"][malware.extended_id]["name"], "malware1")
+        self.assertEqual(data["vertices"][malware_entity.extended_id]["name"], "malware1")
 
     def test_add_link(self):
         response = client.post(
@@ -375,7 +375,7 @@ def setUp(self) -> None:
         self.observable1.tag(["tag1", "tag2"])
         self.observable2 = hostname.Hostname(value="test2.com").save()
         self.observable3 = url.Url(value="http://test1.com/admin").save()
-        self.entity1 = ThreatActor(name="tester").save()
+        self.entity1 = threat_actor.ThreatActor(name="tester").save()
         self.indicator1 = Regex(
             name="test c2",
             pattern="test[0-9].com",
@@ -576,7 +576,7 @@ def setUp(self) -> None:
         ).json()
         client.headers = {"Authorization": "Bearer " + token_data["access_token"]}
 
-        self.persistence = AttackPattern(name="persistence").save()
+        self.persistence = attack_pattern.AttackPattern(name="persistence").save()
         self.persistence.tag(["triage"])
         self.persistence_artifact = ForensicArtifact.from_yaml_string(
             """doc: Crontab files.
diff --git a/tests/schemas/entity.py b/tests/schemas/entity.py
index ea33cf99c..b6d38208c 100644
--- a/tests/schemas/entity.py
+++ b/tests/schemas/entity.py
@@ -3,14 +3,14 @@
 
 from core import database_arango
 from core.schemas import tag
-from core.schemas.entity import (
-    AttackPattern,
-    Entity,
-    Malware,
-    ThreatActor,
-    Tool,
-    Vulnerability,
+from core.schemas.entities import (
+    attack_pattern,
+    malware,
+    threat_actor,
+    tool,
+    vulnerability,
 )
+from core.schemas.entity import Entity
 from core.schemas.observables import hostname
 
 
@@ -18,18 +18,18 @@ class EntityTest(unittest.TestCase):
     def setUp(self) -> None:
         database_arango.db.connect(database="yeti_test")
         database_arango.db.clear()
-        self.ta1 = ThreatActor(name="APT123", aliases=["CrazyFrog"]).save()
-        self.vuln1 = Vulnerability(name="CVE-2018-1337", title="elite exploit").save()
-        self.malware1 = Malware(
+        self.ta1 = threat_actor.ThreatActor(name="APT123", aliases=["CrazyFrog"]).save()
+        self.vuln1 = vulnerability.Vulnerability(name="CVE-2018-1337", title="elite exploit").save()
+        self.malware1 = malware.Malware(
             name="zeus", created=datetime.datetime(2020, 1, 1)
         ).save()
-        self.tool1 = Tool(name="mimikatz").save()
+        self.tool1 = tool.Tool(name="mimikatz").save()
 
     def tearDown(self) -> None:
         database_arango.db.clear()
 
     def test_create_entity(self) -> None:
-        result = ThreatActor(name="APT0").save()
+        result = threat_actor.ThreatActor(name="APT0").save()
         self.assertIsNotNone(result.id)
         self.assertIsNotNone(result.created)
         self.assertEqual(result.name, "APT0")
@@ -41,11 +41,11 @@ def test_entity_get_correct_type(self) -> None:
         result = Entity.get(self.ta1.id)
         assert result is not None
         self.assertIsNotNone(result)
-        self.assertIsInstance(result, ThreatActor)
+        self.assertIsInstance(result, threat_actor.ThreatActor)
         self.assertEqual(result.type, "threat-actor")
 
     def test_attack_pattern(self) -> None:
-        result = AttackPattern(
+        result = attack_pattern.AttackPattern(
             name="Abuse Elevation Control Mechanism",
             aliases=["T1548"],
             kill_chain_phases=["mitre-attack:Privilege Escalation"],
@@ -57,19 +57,19 @@ def test_attack_pattern(self) -> None:
         self.assertIn("T1548", result.aliases)
 
     def test_entity_dupe_name_type(self) -> None:
-        oldm = Malware(name="APT123").save()
-        ta = ThreatActor.find(name="APT123")
-        m = Malware.find(name="APT123")
+        oldm = malware.Malware(name="APT123").save()
+        ta = threat_actor.ThreatActor.find(name="APT123")
+        m = malware.Malware.find(name="APT123")
         self.assertEqual(ta.id, self.ta1.id)
         self.assertEqual(m.id, oldm.id)
-        self.assertIsInstance(m, Malware)
-        self.assertIsInstance(ta, ThreatActor)
+        self.assertIsInstance(m, malware.Malware)
+        self.assertIsInstance(ta, threat_actor.ThreatActor)
 
     def test_list_entities(self) -> None:
         all_entities = list(Entity.list())
-        threat_actor_entities = list(ThreatActor.list())
-        tool_entities = list(Tool.list())
-        malware_entities = list(Malware.list())
+        threat_actor_entities = list(threat_actor.ThreatActor.list())
+        tool_entities = list(tool.Tool.list())
+        malware_entities = list(malware.Malware.list())
 
         self.assertEqual(len(all_entities), 4)
 
@@ -122,7 +122,7 @@ def test_filter_entities_time(self):
         self.assertNotIn(self.malware1, entities)
 
     def test_entity_with_tags(self):
-        entity = ThreatActor(name="APT0").save()
+        entity = threat_actor.ThreatActor(name="APT0").save()
         entity.tag(["tag1", "tag2"])
         observable = hostname.Hostname(value="doman.com").save()
 
@@ -145,13 +145,13 @@ def test_entity_with_tags(self):
 
     def test_duplicate_name(self):
         """Tests that saving an entity with an existing name will return the existing entity."""
-        ta = ThreatActor(name="APT123").save()
+        ta = threat_actor.ThreatActor(name="APT123").save()
         self.assertEqual(ta.id, self.ta1.id)
 
     def test_entity_duplicate_name(self):
         """Tests that two entities of different types can have the same name."""
-        psexec_tool = Tool(name="psexec").save()
-        psexec_ap = AttackPattern(name="psexec").save()
+        psexec_tool = tool.Tool(name="psexec").save()
+        psexec_ap = attack_pattern.AttackPattern(name="psexec").save()
         self.assertNotEqual(psexec_tool.id, psexec_ap.id)
         self.assertEqual(psexec_tool.type, "tool")
         self.assertEqual(psexec_ap.type, "attack-pattern")
@@ -159,12 +159,12 @@ def test_entity_duplicate_name(self):
     def test_no_empty_name(self):
         """Tests that an entity with an empty name cannot be saved."""
         with self.assertRaises(ValueError):
-            ThreatActor(name="").save()
+            threat_actor.ThreatActor(name="").save()
 
     def test_bad_cve_name(self):
-        vulnerability = Vulnerability(name="1337-4242").save()
-        self.assertEqual(Vulnerability.is_valid(vulnerability), False)
+        vuln = vulnerability.Vulnerability(name="1337-4242").save()
+        self.assertEqual(vulnerability.Vulnerability.is_valid(vuln), False)
 
     def test_correct_cve_name(self):
-        vulnerability = Vulnerability(name="CVE-1337-4242").save()
-        self.assertEqual(Vulnerability.is_valid(vulnerability), True)
+        vuln = vulnerability.Vulnerability(name="CVE-1337-4242").save()
+        self.assertEqual(vulnerability.Vulnerability.is_valid(vuln), True)
diff --git a/tests/schemas/fixture.py b/tests/schemas/fixture.py
index 4ba63a045..8267393d1 100644
--- a/tests/schemas/fixture.py
+++ b/tests/schemas/fixture.py
@@ -1,7 +1,7 @@
 import unittest
 
 from core import database_arango
-from core.schemas.entity import Investigation, Malware, ThreatActor
+from core.schemas.entities import investigation, malware, threat_actor
 from core.schemas.indicator import DiamondModel, Query, Regex
 from core.schemas.observables import (
     bic,
@@ -57,7 +57,7 @@ def test_something(self):
         ibantest.link_to(bictest, "bic", "BIC")
         ibantest.tag(["example"])
 
-        ta = ThreatActor(name="HackerActor").save()
+        ta = threat_actor.ThreatActor(name="HackerActor").save()
         ta.tag(["Hack!ré T@ëst"])
         ta.link_to(hacker, "uses", "Uses domain")
 
@@ -68,7 +68,7 @@ def test_something(self):
             diamond=DiamondModel.capability,
         ).save()
         regex.link_to(hacker, "indicates", "Domain dropped by this regex")
-        xmrig = Malware(name="xmrig").save()
+        xmrig = malware.Malware(name="xmrig").save()
         xmrig.tag(["xmrig"])
         regex.link_to(xmrig, "indicates", "Usual name for dropped binary")
 
@@ -81,7 +81,7 @@ def test_something(self):
             target_systems=["timesketch", "plaso"],
             relevant_tags=["ssh", "login"],
         ).save()
-        i = Investigation(
+        i = investigation.Investigation(
             name="coin mining case",
             reference="http://timesketch-server/sketch/12345",
             relevant_tags=["coin", "mining"],
diff --git a/tests/schemas/graph.py b/tests/schemas/graph.py
index e55b50211..6682925a9 100644
--- a/tests/schemas/graph.py
+++ b/tests/schemas/graph.py
@@ -3,7 +3,8 @@
 from fastapi.testclient import TestClient
 
 from core import database_arango
-from core.schemas.entity import Campaign, Entity, Malware
+from core.schemas.entities import campaign, malware
+from core.schemas.entity import Entity
 from core.schemas.graph import GraphFilter, Relationship
 from core.schemas.observables import hostname, ipv4, user_agent
 from core.web import webapp
@@ -19,9 +20,9 @@ def setUp(self) -> None:
         self.observable2 = ipv4.IPv4(value="127.0.0.1").save()
         self.observable3 = ipv4.IPv4(value="8.8.8.8").save()
         self.observable4 = user_agent.UserAgent(value="Mozilla/5.0").save()
-        self.entity1 = Malware(name="plugx").save()
-        self.entity2 = Campaign(name="campaign1").save()
-        self.entity3 = Campaign(name="campaign2").save()
+        self.entity1 = malware.Malware(name="plugx").save()
+        self.entity2 = campaign.Campaign(name="campaign1").save()
+        self.entity3 = campaign.Campaign(name="campaign2").save()
 
     def tearDown(self) -> None:
         database_arango.db.clear()

From 9a095041e6d30da0ed9dcbd6cede1c6f348df05a Mon Sep 17 00:00:00 2001
From: Fred Baguelin 
Date: Tue, 17 Sep 2024 08:40:09 +0200
Subject: [PATCH 17/39] Ruff format

---
 core/common/utils.py                   | 6 +++---
 core/schemas/__init__.py               | 5 ++---
 core/schemas/entities/vulnerability.py | 5 ++++-
 core/schemas/entity.py                 | 8 ++++++--
 core/schemas/observable.py             | 6 +++++-
 tests/apiv2/graph.py                   | 4 +++-
 tests/schemas/entity.py                | 4 +++-
 7 files changed, 26 insertions(+), 12 deletions(-)

diff --git a/core/common/utils.py b/core/common/utils.py
index a0010274c..28d81c447 100644
--- a/core/common/utils.py
+++ b/core/common/utils.py
@@ -14,9 +14,9 @@
 
 if hasattr(yeti_config, "tldextract"):
     if yeti_config.tldextract.extra_suffixes:
-        tld_extract_dict["extra_suffixes"] = (
-            yeti_config.tldextract.extra_suffixes.split(",")
-        )
+        tld_extract_dict[
+            "extra_suffixes"
+        ] = yeti_config.tldextract.extra_suffixes.split(",")
     if yeti_config.tldextract.suffix_list_urls:
         tld_extract_dict["suffix_list_urls"] = yeti_config.tldextract.suffix_list_urls
 
diff --git a/core/schemas/__init__.py b/core/schemas/__init__.py
index a8093a8fc..5ab9bc2f5 100644
--- a/core/schemas/__init__.py
+++ b/core/schemas/__init__.py
@@ -9,6 +9,7 @@
 
 logger = logging.getLogger(__name__)
 
+
 def load_observables():
     logger.info("Registering observable types")
     for observable_file in Path(__file__).parent.glob("observables/**/*.py"):
@@ -43,9 +44,7 @@ def load_entities():
             module_name = f"core.schemas.entities.private.{entity_file.stem}"
         enum_value = entity_file.stem.replace("_", "-")
         if entity_file.stem not in entity.EntityType.__members__:
-            aenum.extend_enum(
-                entity.EntityType, entity_file.stem, enum_value
-            )
+            aenum.extend_enum(entity.EntityType, entity_file.stem, enum_value)
         module = importlib.import_module(module_name)
         for _, obj in inspect.getmembers(module, inspect.isclass):
             if issubclass(obj, entity.Entity):
diff --git a/core/schemas/entities/vulnerability.py b/core/schemas/entities/vulnerability.py
index 58fbe26e8..af694d72d 100644
--- a/core/schemas/entities/vulnerability.py
+++ b/core/schemas/entities/vulnerability.py
@@ -6,7 +6,10 @@
 
 from core.schemas import entity
 
-vulnerability_matcher = re.compile(r"(?P
\W?)(?PCVE-\d{4}-\d{4,7})(?P\W?)")
+vulnerability_matcher = re.compile(
+    r"(?P
\W?)(?PCVE-\d{4}-\d{4,7})(?P\W?)"
+)
+
 
 class SeverityType(str, Enum):
     none = "none"
diff --git a/core/schemas/entity.py b/core/schemas/entity.py
index 6b3f39982..23331397a 100644
--- a/core/schemas/entity.py
+++ b/core/schemas/entity.py
@@ -10,10 +10,14 @@
 
 
 # forward declarations
-class EntityType(str, Enum): ...
+class EntityType(str, Enum):
+    ...
+
+
 EntityTypes = ()
 TYPE_MAPPING = {}
 
+
 class Entity(YetiTagModel, database_arango.ArangoYetiConnector):
     _exclude_overwrite: list[str] = ["related_observables_count"]
     _collection_name: ClassVar[str] = "entities"
@@ -77,4 +81,4 @@ def add_context(
         return self.save()
 
 
-TYPE_MAPPING.update({"entities": Entity, "entity": Entity})
\ No newline at end of file
+TYPE_MAPPING.update({"entities": Entity, "entity": Entity})
diff --git a/core/schemas/observable.py b/core/schemas/observable.py
index 9cd64fa6e..ded73db3f 100644
--- a/core/schemas/observable.py
+++ b/core/schemas/observable.py
@@ -18,9 +18,13 @@
 
 
 # forward declarations
-class ObservableType(str, Enum): ...
+class ObservableType(str, Enum):
+    ...
+
+
 TYPE_MAPPING = {}
 
+
 class Observable(YetiTagModel, database_arango.ArangoYetiConnector):
     _collection_name: ClassVar[str] = "observables"
     _type_filter: ClassVar[str | None] = None
diff --git a/tests/apiv2/graph.py b/tests/apiv2/graph.py
index 639adafb0..c72d5c15f 100644
--- a/tests/apiv2/graph.py
+++ b/tests/apiv2/graph.py
@@ -299,7 +299,9 @@ def test_neighbor_filter_in(self):
         data = response.json()
         self.assertEqual(response.status_code, 200, data)
         self.assertEqual(len(data["vertices"]), 1)
-        self.assertEqual(data["vertices"][malware_entity.extended_id]["name"], "malware1")
+        self.assertEqual(
+            data["vertices"][malware_entity.extended_id]["name"], "malware1"
+        )
 
     def test_add_link(self):
         response = client.post(
diff --git a/tests/schemas/entity.py b/tests/schemas/entity.py
index b6d38208c..fabbd8181 100644
--- a/tests/schemas/entity.py
+++ b/tests/schemas/entity.py
@@ -19,7 +19,9 @@ def setUp(self) -> None:
         database_arango.db.connect(database="yeti_test")
         database_arango.db.clear()
         self.ta1 = threat_actor.ThreatActor(name="APT123", aliases=["CrazyFrog"]).save()
-        self.vuln1 = vulnerability.Vulnerability(name="CVE-2018-1337", title="elite exploit").save()
+        self.vuln1 = vulnerability.Vulnerability(
+            name="CVE-2018-1337", title="elite exploit"
+        ).save()
         self.malware1 = malware.Malware(
             name="zeus", created=datetime.datetime(2020, 1, 1)
         ).save()

From 93a501d327d0197ccfdaacd52dd106383a67ffc0 Mon Sep 17 00:00:00 2001
From: Fred Baguelin 
Date: Tue, 17 Sep 2024 08:58:02 +0200
Subject: [PATCH 18/39] Add __init__.py in entities/observables private folder

---
 core/schemas/entities/private/.gitignore     | 3 ++-
 core/schemas/entities/private/__init__.py    | 0
 core/schemas/observables/private/.gitignore  | 3 ++-
 core/schemas/observables/private/__init__.py | 0
 4 files changed, 4 insertions(+), 2 deletions(-)
 create mode 100644 core/schemas/entities/private/__init__.py
 create mode 100644 core/schemas/observables/private/__init__.py

diff --git a/core/schemas/entities/private/.gitignore b/core/schemas/entities/private/.gitignore
index e5af87e9b..a91bcd1d7 100644
--- a/core/schemas/entities/private/.gitignore
+++ b/core/schemas/entities/private/.gitignore
@@ -1,3 +1,4 @@
 *
 !.gitignore
-!README.md
\ No newline at end of file
+!README.md
+!__init__.py
\ No newline at end of file
diff --git a/core/schemas/entities/private/__init__.py b/core/schemas/entities/private/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/core/schemas/observables/private/.gitignore b/core/schemas/observables/private/.gitignore
index e5af87e9b..a91bcd1d7 100644
--- a/core/schemas/observables/private/.gitignore
+++ b/core/schemas/observables/private/.gitignore
@@ -1,3 +1,4 @@
 *
 !.gitignore
-!README.md
\ No newline at end of file
+!README.md
+!__init__.py
\ No newline at end of file
diff --git a/core/schemas/observables/private/__init__.py b/core/schemas/observables/private/__init__.py
new file mode 100644
index 000000000..e69de29bb

From 002c3b14a3637dd39fcdf0d4c78cdcdd3487d13f Mon Sep 17 00:00:00 2001
From: Fred Baguelin 
Date: Tue, 17 Sep 2024 12:53:03 +0200
Subject: [PATCH 19/39] Add indicators and update type literal for all objects

---
 core/schemas/__init__.py                    |  95 ++++--
 core/schemas/dfiq.py                        |  11 +-
 core/schemas/entities/attack_pattern.py     |   4 +-
 core/schemas/entities/campaign.py           |   4 +-
 core/schemas/entities/company.py            |   4 +-
 core/schemas/entities/course_of_action.py   |   4 +-
 core/schemas/entities/identity.py           |   4 +-
 core/schemas/entities/intrusion_set.py      |   4 +-
 core/schemas/entities/investigation.py      |   4 +-
 core/schemas/entities/malware.py            |   4 +-
 core/schemas/entities/note.py               |   4 +-
 core/schemas/entities/phone.py              |   4 +-
 core/schemas/entities/private/README.md     |   4 +-
 core/schemas/entities/threat_actor.py       |   4 +-
 core/schemas/entities/tool.py               |   4 +-
 core/schemas/entities/vulnerability.py      |   4 +-
 core/schemas/entity.py                      |   5 +-
 core/schemas/indicator.py                   | 311 +-------------------
 core/schemas/indicators/__init__.py         |   0
 core/schemas/indicators/forensicartifact.py | 177 +++++++++++
 core/schemas/indicators/private/.gitignore  |   4 +
 core/schemas/indicators/private/README.md   |   2 +
 core/schemas/indicators/private/__init__.py |   0
 core/schemas/indicators/query.py            |  16 +
 core/schemas/indicators/regex.py            |  33 +++
 core/schemas/indicators/sigma.py            |  16 +
 core/schemas/indicators/suricata.py         |  38 +++
 core/schemas/indicators/yara.py             |  16 +
 core/schemas/observable.py                  |   6 +-
 core/schemas/observables/asn.py             |   4 +-
 core/schemas/observables/bic.py             |   3 +-
 core/schemas/observables/certificate.py     |   3 +-
 core/schemas/observables/cidr.py            |   4 +-
 core/schemas/observables/command_line.py    |   4 +-
 core/schemas/observables/docker_image.py    |   4 +-
 core/schemas/observables/email.py           |   4 +-
 core/schemas/observables/file.py            |   4 +-
 core/schemas/observables/generic.py         |   2 +-
 core/schemas/observables/hostname.py        |   4 +-
 core/schemas/observables/iban.py            |   4 +-
 core/schemas/observables/imphash.py         |   4 +-
 core/schemas/observables/ipv4.py            |   4 +-
 core/schemas/observables/ipv6.py            |   4 +-
 core/schemas/observables/ja3.py             |   4 +-
 core/schemas/observables/jarm.py            |   4 +-
 core/schemas/observables/mac_address.py     |   4 +-
 core/schemas/observables/md5.py             |   4 +-
 core/schemas/observables/mutex.py           |   4 +-
 core/schemas/observables/named_pipe.py      |   4 +-
 core/schemas/observables/package.py         |   4 +-
 core/schemas/observables/path.py            |   3 +-
 core/schemas/observables/registry_key.py    |   2 +-
 core/schemas/observables/sha1.py            |   4 +-
 core/schemas/observables/sha256.py          |   4 +-
 core/schemas/observables/ssdeep.py          |   4 +-
 core/schemas/observables/tlsh.py            |   4 +-
 core/schemas/observables/url.py             |   4 +-
 core/schemas/observables/user_account.py    |   2 +-
 core/schemas/observables/user_agent.py      |   4 +-
 core/schemas/observables/wallet.py          |   4 +-
 core/web/apiv2/entities.py                  |   4 +-
 core/web/apiv2/graph.py                     |   2 +-
 core/web/apiv2/indicators.py                |   4 +-
 core/web/apiv2/observables.py               |  23 +-
 tests/apiv2/graph.py                        |   9 +-
 tests/apiv2/indicators.py                   |   7 +-
 tests/schemas/fixture.py                    |  11 +-
 tests/schemas/indicator.py                  |  43 ++-
 68 files changed, 548 insertions(+), 456 deletions(-)
 create mode 100644 core/schemas/indicators/__init__.py
 create mode 100644 core/schemas/indicators/forensicartifact.py
 create mode 100644 core/schemas/indicators/private/.gitignore
 create mode 100644 core/schemas/indicators/private/README.md
 create mode 100644 core/schemas/indicators/private/__init__.py
 create mode 100644 core/schemas/indicators/query.py
 create mode 100644 core/schemas/indicators/regex.py
 create mode 100644 core/schemas/indicators/sigma.py
 create mode 100644 core/schemas/indicators/suricata.py
 create mode 100644 core/schemas/indicators/yara.py

diff --git a/core/schemas/__init__.py b/core/schemas/__init__.py
index 5ab9bc2f5..905167e7f 100644
--- a/core/schemas/__init__.py
+++ b/core/schemas/__init__.py
@@ -5,35 +5,13 @@
 
 import aenum
 
-from core.schemas import entity, observable
+from core.schemas import entity, indicator, observable
 
 logger = logging.getLogger(__name__)
 
-
-def load_observables():
-    logger.info("Registering observable types")
-    for observable_file in Path(__file__).parent.glob("observables/**/*.py"):
-        if observable_file.stem == "__init__":
-            continue
-        logger.info(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}"
-        if observable_file.stem not in observable.ObservableType.__members__:
-            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
-    if "guess" not in observable.ObservableType.__members__:
-        aenum.extend_enum(observable.ObservableType, "guess", "guess")
-
-
 def load_entities():
-    logger.info("Registering entity types")
+    logger.info("Registering entities")
+    modules = dict()
     for entity_file in Path(__file__).parent.glob("entities/**/*.py"):
         if entity_file.stem == "__init__":
             continue
@@ -45,6 +23,9 @@ def load_entities():
         enum_value = entity_file.stem.replace("_", "-")
         if entity_file.stem not in entity.EntityType.__members__:
             aenum.extend_enum(entity.EntityType, entity_file.stem, enum_value)
+        modules[module_name] = enum_value
+    entity.TYPE_MAPPING = {"entity": entity.Entity, "entities": entity.Entity}
+    for module_name, enum_value in modules.items():
         module = importlib.import_module(module_name)
         for _, obj in inspect.getmembers(module, inspect.isclass):
             if issubclass(obj, entity.Entity):
@@ -59,5 +40,69 @@ def load_entities():
             entity.EntityTypes |= cls
 
 
+def load_indicators():
+    logger.info("Registering indicators")
+    modules = dict()
+    for indicator_file in Path(__file__).parent.glob("indicators/**/*.py"):
+        if indicator_file.stem == "__init__":
+            continue
+        logger.info(f"Registering indicator type {indicator_file.stem}")
+        if indicator_file.parent.stem == "indicators":
+            module_name = f"core.schemas.indicators.{indicator_file.stem}"
+        elif indicator_file.parent.stem == "private":
+            module_name = f"core.schemas.indicators.private.{indicator_file.stem}"
+        enum_value = indicator_file.stem
+        if indicator_file.stem not in indicator.IndicatorType.__members__:
+            aenum.extend_enum(indicator.IndicatorType, indicator_file.stem, enum_value)
+        modules[module_name] = enum_value
+    indicator.TYPE_MAPPING = {"indicator": indicator.Indicator, "indicators": indicator.Indicator}
+    for module_name, enum_value in modules.items():
+        module = importlib.import_module(module_name)
+        for _, obj in inspect.getmembers(module, inspect.isclass):
+            if issubclass(obj, indicator.Indicator):
+                indicator.TYPE_MAPPING[enum_value] = obj
+    for key in indicator.TYPE_MAPPING:
+        if key in ["indicator", "indicators"]:
+            continue
+        cls = indicator.TYPE_MAPPING[key]
+        if not indicator.IndicatorTypes:
+            indicator.IndicatorTypes = cls
+        else:
+            indicator.IndicatorTypes |= cls
+
+def load_observables():
+    logger.info("Registering observables")
+    modules = dict()
+    for observable_file in Path(__file__).parent.glob("observables/**/*.py"):
+        if observable_file.stem == "__init__":
+            continue
+        logger.info(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}"
+        if observable_file.stem not in observable.ObservableType.__members__:
+            aenum.extend_enum(
+                observable.ObservableType, observable_file.stem, observable_file.stem
+            )
+        modules[module_name] = observable_file.stem
+    if "guess" not in observable.ObservableType.__members__:
+        aenum.extend_enum(observable.ObservableType, "guess", "guess")
+    observable.TYPE_MAPPING = {"observable": observable.Observable, "observables": observable.Observable}
+    for module_name, enum_value in modules.items():
+        module = importlib.import_module(module_name)
+        for _, obj in inspect.getmembers(module, inspect.isclass):
+            if issubclass(obj, observable.Observable):
+                observable.TYPE_MAPPING[enum_value] = obj
+    for key in observable.TYPE_MAPPING:
+        if key in ["observable", "observables"]:
+            continue
+        cls = observable.TYPE_MAPPING[key]
+        if not observable.ObservableTypes:
+            observable.ObservableTypes = cls
+        else:
+            observable.ObservableTypes |= cls
+
 load_observables()
 load_entities()
+load_indicators()
\ No newline at end of file
diff --git a/core/schemas/dfiq.py b/core/schemas/dfiq.py
index 282c5ea10..2001aa645 100644
--- a/core/schemas/dfiq.py
+++ b/core/schemas/dfiq.py
@@ -14,6 +14,7 @@
 from core.config.config import yeti_config
 from core.helpers import now
 from core.schemas import indicator
+from core.schemas.indicators import forensicartifact, query
 from core.schemas.model import YetiModel
 
 LATEST_SUPPORTED_DFIQ_VERSION = "1.1.0"
@@ -98,7 +99,7 @@ def extract_indicators(question: "DFIQQuestion") -> None:
                 continue
 
             if step.type in ("ForensicArtifact", "artifact"):
-                artifact = indicator.ForensicArtifact.find(name=step.value)
+                artifact = forensicartifact.ForensicArtifact.find(name=step.value)
                 if not artifact:
                     logging.warning(
                         "Missing artifact %s in %s", step.value, question.dfiq_id
@@ -108,9 +109,9 @@ def extract_indicators(question: "DFIQQuestion") -> None:
                 continue
 
             elif step.type and step.value and "query" in step.type:
-                query = indicator.Query.find(pattern=step.value)
-                if not query:
-                    query = indicator.Query(
+                query_indicator = query.Query.find(pattern=step.value)
+                if not query_indicator:
+                    query_indicator = query.Query(
                         name=f"{step.name} ({step.type})",
                         description=step.description or "",
                         pattern=step.value,
@@ -119,7 +120,7 @@ def extract_indicators(question: "DFIQQuestion") -> None:
                         location=step.type,
                         diamond=indicator.DiamondModel.victim,
                     ).save()
-                question.link_to(query, "query", "Uses query")
+                question.link_to(query_indicator, "query", "Uses query")
 
             else:
                 logging.warning(
diff --git a/core/schemas/entities/attack_pattern.py b/core/schemas/entities/attack_pattern.py
index 6074257ba..dc5fb41d7 100644
--- a/core/schemas/entities/attack_pattern.py
+++ b/core/schemas/entities/attack_pattern.py
@@ -1,10 +1,10 @@
-from typing import ClassVar
+from typing import ClassVar, Literal
 
 from core.schemas import entity
 
 
 class AttackPattern(entity.Entity):
     _type_filter: ClassVar[str] = entity.EntityType.attack_pattern
-    type: entity.EntityType = entity.EntityType.attack_pattern
+    type: Literal[entity.EntityType.attack_pattern] = entity.EntityType.attack_pattern
     aliases: list[str] = []
     kill_chain_phases: list[str] = []
diff --git a/core/schemas/entities/campaign.py b/core/schemas/entities/campaign.py
index 16012ea9a..8101aa7eb 100644
--- a/core/schemas/entities/campaign.py
+++ b/core/schemas/entities/campaign.py
@@ -1,5 +1,5 @@
 import datetime
-from typing import ClassVar
+from typing import ClassVar, Literal
 
 from pydantic import Field
 
@@ -9,7 +9,7 @@
 
 class Campaign(entity.Entity):
     _type_filter: ClassVar[str] = entity.EntityType.campaign
-    type: entity.EntityType = entity.EntityType.campaign
+    type: Literal[entity.EntityType.campaign] = entity.EntityType.campaign
 
     aliases: list[str] = []
     first_seen: datetime.datetime = Field(default_factory=now)
diff --git a/core/schemas/entities/company.py b/core/schemas/entities/company.py
index 2e3f60c68..ed4c245f8 100644
--- a/core/schemas/entities/company.py
+++ b/core/schemas/entities/company.py
@@ -1,8 +1,8 @@
-from typing import ClassVar
+from typing import ClassVar, Literal
 
 from core.schemas import entity
 
 
 class Company(entity.Entity):
-    type: entity.EntityType = entity.EntityType.company
+    type: Literal[entity.EntityType.company] = entity.EntityType.company
     _type_filter: ClassVar[str] = entity.EntityType.company
diff --git a/core/schemas/entities/course_of_action.py b/core/schemas/entities/course_of_action.py
index 395ee1073..09e154f02 100644
--- a/core/schemas/entities/course_of_action.py
+++ b/core/schemas/entities/course_of_action.py
@@ -1,8 +1,8 @@
-from typing import ClassVar
+from typing import ClassVar, Literal
 
 from core.schemas import entity
 
 
 class CourseOfAction(entity.Entity):
     _type_filter: ClassVar[str] = entity.EntityType.course_of_action
-    type: entity.EntityType = entity.EntityType.course_of_action
+    type: Literal[entity.EntityType.course_of_action] = entity.EntityType.course_of_action
diff --git a/core/schemas/entities/identity.py b/core/schemas/entities/identity.py
index d3f5b20c8..6359c80b6 100644
--- a/core/schemas/entities/identity.py
+++ b/core/schemas/entities/identity.py
@@ -1,11 +1,11 @@
-from typing import ClassVar
+from typing import ClassVar, Literal
 
 from core.schemas import entity
 
 
 class Identity(entity.Entity):
     _type_filter: ClassVar[str] = entity.EntityType.identity
-    type: entity.EntityType = entity.EntityType.identity
+    type: Literal[entity.EntityType.identity] = entity.EntityType.identity
 
     identity_class: str = ""
     sectors: list[str] = []
diff --git a/core/schemas/entities/intrusion_set.py b/core/schemas/entities/intrusion_set.py
index 35dfa4501..2c4c06943 100644
--- a/core/schemas/entities/intrusion_set.py
+++ b/core/schemas/entities/intrusion_set.py
@@ -1,5 +1,5 @@
 import datetime
-from typing import ClassVar
+from typing import ClassVar, Literal
 
 from pydantic import Field
 
@@ -9,7 +9,7 @@
 
 class IntrusionSet(entity.Entity):
     _type_filter: ClassVar[str] = entity.EntityType.intrusion_set
-    type: entity.EntityType = entity.EntityType.intrusion_set
+    type: Literal[entity.EntityType.intrusion_set] = entity.EntityType.intrusion_set
 
     aliases: list[str] = []
     first_seen: datetime.datetime = Field(default_factory=now)
diff --git a/core/schemas/entities/investigation.py b/core/schemas/entities/investigation.py
index b993f6a9b..57fe5ba95 100644
--- a/core/schemas/entities/investigation.py
+++ b/core/schemas/entities/investigation.py
@@ -1,10 +1,10 @@
-from typing import ClassVar
+from typing import ClassVar, Literal
 
 from core.schemas import entity
 
 
 class Investigation(entity.Entity):
     _type_filter: ClassVar[str] = entity.EntityType.investigation
-    type: entity.EntityType = entity.EntityType.investigation
+    type: Literal[entity.EntityType.investigation] = entity.EntityType.investigation
 
     reference: str = ""
diff --git a/core/schemas/entities/malware.py b/core/schemas/entities/malware.py
index dbdcdaafa..d4d390181 100644
--- a/core/schemas/entities/malware.py
+++ b/core/schemas/entities/malware.py
@@ -1,11 +1,11 @@
-from typing import ClassVar
+from typing import ClassVar, Literal
 
 from core.schemas import entity
 
 
 class Malware(entity.Entity):
     _type_filter: ClassVar[str] = entity.EntityType.malware
-    type: entity.EntityType = entity.EntityType.malware
+    type: Literal[entity.EntityType.malware] = entity.EntityType.malware
 
     kill_chain_phases: list[str] = []
     aliases: list[str] = []
diff --git a/core/schemas/entities/note.py b/core/schemas/entities/note.py
index 706e42e85..d14bb3756 100644
--- a/core/schemas/entities/note.py
+++ b/core/schemas/entities/note.py
@@ -1,8 +1,8 @@
-from typing import ClassVar
+from typing import ClassVar, Literal
 
 from core.schemas import entity
 
 
 class Note(entity.Entity):
-    type: entity.EntityType = entity.EntityType.note
+    type: Literal[entity.EntityType.note] = entity.EntityType.note
     _type_filter: ClassVar[str] = entity.EntityType.note
diff --git a/core/schemas/entities/phone.py b/core/schemas/entities/phone.py
index 809eb3745..10e486a80 100644
--- a/core/schemas/entities/phone.py
+++ b/core/schemas/entities/phone.py
@@ -1,8 +1,8 @@
-from typing import ClassVar
+from typing import ClassVar, Literal
 
 from core.schemas import entity
 
 
 class Phone(entity.Entity):
-    type: entity.EntityType = entity.EntityType.phone
+    type: Literal[entity.EntityType.phone] = entity.EntityType.phone
     _type_filter: ClassVar[str] = entity.EntityType.phone
diff --git a/core/schemas/entities/private/README.md b/core/schemas/entities/private/README.md
index 4f9a2488b..c34e233c4 100644
--- a/core/schemas/entities/private/README.md
+++ b/core/schemas/entities/private/README.md
@@ -1,2 +1,2 @@
-### Private entities
-This directory is where you should place your private entities. It could be named anything else, but this one has a `.gitignore` so you don't mess things up. ;-)
+### Private indicators
+This directory is where you should place your private indicators. It could be named anything else, but this one has a `.gitignore` so you don't mess things up. ;-)
diff --git a/core/schemas/entities/threat_actor.py b/core/schemas/entities/threat_actor.py
index 2f9841d00..330955be6 100644
--- a/core/schemas/entities/threat_actor.py
+++ b/core/schemas/entities/threat_actor.py
@@ -1,5 +1,5 @@
 import datetime
-from typing import ClassVar
+from typing import ClassVar, Literal
 
 from pydantic import Field
 
@@ -9,7 +9,7 @@
 
 class ThreatActor(entity.Entity):
     _type_filter: ClassVar[str] = entity.EntityType.threat_actor
-    type: entity.EntityType = entity.EntityType.threat_actor
+    type: Literal[entity.EntityType.threat_actor] = entity.EntityType.threat_actor
 
     threat_actor_types: list[str] = []
     aliases: list[str] = []
diff --git a/core/schemas/entities/tool.py b/core/schemas/entities/tool.py
index c184561f7..6d3822aa7 100644
--- a/core/schemas/entities/tool.py
+++ b/core/schemas/entities/tool.py
@@ -1,11 +1,11 @@
-from typing import ClassVar
+from typing import ClassVar, Literal
 
 from core.schemas import entity
 
 
 class Tool(entity.Entity):
     _type_filter: ClassVar[str] = entity.EntityType.tool
-    type: entity.EntityType = entity.EntityType.tool
+    type: Literal[entity.EntityType.tool] = entity.EntityType.tool
 
     aliases: list[str] = []
     kill_chain_phases: list[str] = []
diff --git a/core/schemas/entities/vulnerability.py b/core/schemas/entities/vulnerability.py
index af694d72d..bb633d68d 100644
--- a/core/schemas/entities/vulnerability.py
+++ b/core/schemas/entities/vulnerability.py
@@ -1,6 +1,6 @@
 import re
 from enum import Enum
-from typing import ClassVar
+from typing import ClassVar, Literal
 
 from pydantic import Field
 
@@ -32,7 +32,7 @@ class Vulnerability(entity.Entity):
     """
 
     _type_filter: ClassVar[str] = entity.EntityType.vulnerability
-    type: entity.EntityType = entity.EntityType.vulnerability
+    type: Literal[entity.EntityType.vulnerability] = entity.EntityType.vulnerability
 
     title: str = ""
     base_score: float = Field(ge=0.0, le=10.0, default=0.0)
diff --git a/core/schemas/entity.py b/core/schemas/entity.py
index 23331397a..62d798401 100644
--- a/core/schemas/entity.py
+++ b/core/schemas/entity.py
@@ -78,7 +78,4 @@ def add_context(
         else:
             context["source"] = source
             self.context.append(context)
-        return self.save()
-
-
-TYPE_MAPPING.update({"entities": Entity, "entity": Entity})
+        return self.save()
\ No newline at end of file
diff --git a/core/schemas/indicator.py b/core/schemas/indicator.py
index 4b3b8bc1c..8e0870936 100644
--- a/core/schemas/indicator.py
+++ b/core/schemas/indicator.py
@@ -1,15 +1,9 @@
 import datetime
-import io
 import logging
-import re
 from enum import Enum
-from typing import Annotated, ClassVar, List, Literal, Type, Union
+from typing import ClassVar, Literal
 
-import yaml
-from artifacts import definitions, reader, writer
-from artifacts import errors as artifacts_errors
-from idstools import rule
-from pydantic import BaseModel, Field, PrivateAttr, computed_field, field_validator
+from pydantic import BaseModel, Field, computed_field
 
 from core import database_arango
 from core.helpers import now
@@ -24,14 +18,13 @@ def future():
 
 DEFAULT_INDICATOR_VALIDITY_DAYS = 30
 
-
+# forward declarations
 class IndicatorType(str, Enum):
-    regex = "regex"
-    yara = "yara"
-    sigma = "sigma"
-    query = "query"
-    suricata = "suricata"
-    forensicartifact = "forensicartifact"
+    ...
+
+IndicatorTypes = ()
+TYPE_MAPPING = {}
+
 
 
 class IndicatorMatch(BaseModel):
@@ -94,290 +87,4 @@ def search(cls, observables: list[str]) -> list[tuple[str, "Indicator"]]:
                 except NotImplementedError as error:
                     logging.error(
                         f"Indicator type {indicator.type} has not implemented match(): {error}"
-                    )
-
-
-class Regex(Indicator):
-    _type_filter: ClassVar[str] = IndicatorType.regex
-    _compiled_pattern: re.Pattern | None = PrivateAttr(None)
-    type: Literal["regex"] = IndicatorType.regex
-
-    @property
-    def compiled_pattern(self):
-        if not self._compiled_pattern:
-            self._compiled_pattern = re.compile(self.pattern)
-        return self._compiled_pattern
-
-    @field_validator("pattern")
-    @classmethod
-    def validate_regex(cls, value) -> str:
-        try:
-            re.compile(value)
-        except re.error as error:
-            raise ValueError(f"Invalid regex pattern: {error}")
-        return value
-
-    def match(self, value: str) -> IndicatorMatch | None:
-        result = self.compiled_pattern.search(value)
-        if result:
-            return IndicatorMatch(name=self.name, match=result.group())
-        return None
-
-
-class Query(Indicator):
-    """Represents a query that can be sent to another system."""
-
-    _type_filter: ClassVar[str] = IndicatorType.query
-    type: Literal["query"] = IndicatorType.query
-
-    query_type: str
-    target_systems: list[str] = []
-
-    def match(self, value: str) -> IndicatorMatch | None:
-        return None
-
-
-class Yara(Indicator):
-    """Represents a Yara rule.
-
-    Parsing and matching is yet TODO.
-    """
-
-    _type_filter: ClassVar[str] = IndicatorType.yara
-    type: Literal["yara"] = IndicatorType.yara
-
-    def match(self, value: str) -> IndicatorMatch | None:
-        raise NotImplementedError
-
-
-class Suricata(Indicator):
-    """Represents a Suricata rule.
-
-    Parsing and matching is yet TODO.
-    """
-
-    _type_filter: ClassVar[str] = IndicatorType.suricata
-    type: Literal["suricata"] = IndicatorType.suricata
-    sid: int = 0
-    metadata: List[str] = []
-    references: List[str] = []
-
-    def match(self, value: str) -> IndicatorMatch | None:
-        raise NotImplementedError
-
-    @field_validator("pattern")
-    @classmethod
-    def validate_rules(cls, value) -> str:
-        try:
-            rule.parse(value)
-        except Exception as e:
-            raise ValueError(f"invalid {cls.pattern} {e}")
-        return value
-
-    def parse(self) -> rule.Rule | None:
-        try:
-            return rule.parse(self.pattern)
-        except Exception as e:
-            logging.error(f" Error parsing {self.pattern} {e}")
-
-
-class Sigma(Indicator):
-    """Represents a Sigma rule.
-
-    Parsing and matching is yet TODO.
-    """
-
-    _type_filter: ClassVar[str] = IndicatorType.sigma
-    type: Literal["sigma"] = IndicatorType.sigma
-
-    def match(self, value: str) -> IndicatorMatch | None:
-        raise NotImplementedError
-
-
-class ForensicArtifact(Indicator):
-    """Represents a Forensic Artifact
-
-    As defined in https://github.com/ForensicArtifacts/artifacts
-    """
-
-    _type_filter: ClassVar[str] = IndicatorType.forensicartifact
-    type: Literal[IndicatorType.forensicartifact] = IndicatorType.forensicartifact
-
-    sources: list[dict] = []
-    aliases: list[str] = []
-    supported_os: list[str] = []
-
-    def match(self, value: str) -> IndicatorMatch | None:
-        raise NotImplementedError
-
-    @field_validator("pattern")
-    @classmethod
-    def validate_artifact(cls, value) -> str:
-        artifact_reader = reader.YamlArtifactsReader()
-        try:
-            list(artifact_reader.ReadFileObject(io.StringIO(value)))
-        except artifacts_errors.FormatError as error:
-            raise ValueError(f"Invalid ForensicArtifact YAML: {error}")
-        return value
-
-    @classmethod
-    def from_yaml_string(
-        cls, yaml_string: str, update_parents: bool = False
-    ) -> list["ForensicArtifact"]:
-        artifact_reader = reader.YamlArtifactsReader()
-        artifact_writer = writer.YamlArtifactsWriter()
-
-        artifacts_dict = {}
-
-        for definition in artifact_reader.ReadFileObject(io.StringIO(yaml_string)):
-            definition_dict = definition.AsDict()
-            definition_dict["description"] = definition_dict.pop("doc")
-            if definition.urls:
-                definition_dict["description"] += "\n\nURLs:\n"
-                definition_dict["description"] += " ".join(
-                    [f"* {url}\n" for url in definition.urls]
-                )
-            definition_dict["pattern"] = artifact_writer.FormatArtifacts([definition])
-            definition_dict["location"] = "host"
-            definition_dict["diamond"] = DiamondModel.victim
-            definition_dict["relevant_tags"] = [definition_dict["name"]]
-            forensic_indicator = cls(**definition_dict).save()
-            artifacts_dict[definition.name] = forensic_indicator
-
-        if update_parents:
-            for artifact in artifacts_dict.values():
-                artifact.update_parents(artifacts_dict)
-
-        return list(artifacts_dict.values())
-
-    def update_yaml(self):
-        artifact_reader = reader.YamlArtifactsReader()
-        definition_dict = next(
-            artifact_reader.ReadFileObject(io.StringIO(self.pattern))
-        ).AsDict()
-        definition_dict["doc"] = self.description.split("\n\nURLs:")[0]
-        definition_dict["name"] = self.name
-        definition_dict["supported_os"] = self.supported_os
-        self.pattern = yaml.safe_dump(definition_dict)
-
-    def update_parents(self, artifacts_dict: dict[str, "ForensicArtifact"]) -> None:
-        for source in self.sources:
-            if not source["type"] == definitions.TYPE_INDICATOR_ARTIFACT_GROUP:
-                continue
-            for child_name in source["attributes"]["names"]:
-                child = artifacts_dict.get(child_name)
-                if not child:
-                    logging.error(f"Missing child {child_name} for {self.name}")
-                    continue
-
-                add_tags = set(self.relevant_tags + [self.name])
-                child.relevant_tags = list(add_tags | set(child.relevant_tags))
-                child.save()
-                child.link_to(
-                    self,
-                    "included in",
-                    f"Included in ForensicArtifact definition for {self.name}",
-                )
-
-    def save_indicators(self, create_links: bool = False):
-        indicators = []
-        for source in self.sources:
-            if source["type"] == definitions.TYPE_INDICATOR_FILE:
-                for path in source["attributes"]["paths"]:
-                    # TODO: consider using https://github.com/log2timeline/dfvfs/blob/main/dfvfs/lib/glob2regex.py
-                    pattern = ARTIFACT_INTERPOLATION_RE.sub("*", path)
-                    pattern = re.escape(pattern).replace("\\*", ".*")
-                    # Account for different path separators
-                    pattern = re.sub(r"\\\\", r"[\\|/]", pattern)
-                    indicator = Regex.find(name=path)
-                    if not indicator:
-                        try:
-                            indicator = Regex(
-                                name=path,
-                                pattern=pattern,
-                                location="filesystem",
-                                diamond=DiamondModel.victim,
-                                relevant_tags=self.relevant_tags,
-                            ).save()
-                            indicators.append(indicator)
-                        except Exception as error:
-                            logging.error(
-                                f"Failed to create indicator for {path} (was: {source['attributes']['paths']}): {error}"
-                            )
-                            continue
-
-                    else:
-                        indicator.relevant_tags = list(
-                            set(indicator.relevant_tags + self.relevant_tags)
-                        )
-                        indicator.save()
-        if source["type"] == definitions.TYPE_INDICATOR_WINDOWS_REGISTRY_KEY:
-            for key in source["attributes"]["keys"]:
-                pattern = re.sub(r"\\\*$", "", key)
-                pattern = ARTIFACT_INTERPOLATION_RE.sub("*", pattern)
-                pattern = re.escape(pattern)
-                pattern = pattern.replace(
-                    "HKEY_USERS\\\\\\*",
-                    r"(HKEY_USERS\\*|HKEY_CURRENT_USER)",
-                )
-                pattern = pattern.replace("*", r".*").replace("?", r".")
-                if "CurrentControlSet" in pattern:
-                    pattern = pattern.replace(
-                        "CurrentControlSet", "(CurrentControlSet|ControlSet[0-9]+)"
-                    )
-                    pattern = pattern.replace("HKEY_LOCAL_MACHINE\\\\System\\\\", "")
-
-                indicator = Regex.find(name=key)
-
-                if not indicator:
-                    try:
-                        indicator = Regex(
-                            name=key,
-                            pattern=pattern,
-                            location="registry",
-                            diamond=DiamondModel.victim,
-                            relevant_tags=self.relevant_tags,
-                        ).save()
-                        indicators.append(indicator)
-                    except Exception as error:
-                        logging.error(
-                            f"Failed to create indicator for {key} (was: {source['attributes']['keys']}): {error}"
-                        )
-                        continue
-                else:
-                    indicator.relevant_tags = list(
-                        set(indicator.relevant_tags + self.relevant_tags)
-                    )
-                    indicator.save()
-        if create_links:
-            for indicator in indicators:
-                indicator.link_to(self, "indicates", f"Indicates {indicator.name}")
-        return indicators
-
-
-ARTIFACT_INTERPOLATION_RE = re.compile(r"%%[a-z._]+%%")
-ARTIFACT_INTERPOLATION_RE_HKEY_USERS = re.compile(r"HKEY_USERS\\%%users.sid%%")
-
-TYPE_MAPPING = {
-    "regex": Regex,
-    "yara": Yara,
-    "sigma": Sigma,
-    "suricata": Suricata,
-    "query": Query,
-    "forensicartifact": ForensicArtifact,
-    "indicator": Indicator,
-    "indicators": Indicator,
-}
-
-IndicatorTypes = Annotated[
-    Union[Regex, Yara, Suricata, Sigma, Query, ForensicArtifact],
-    Field(discriminator="type"),
-]
-IndicatorClasses = (
-    Type[Regex]
-    | Type[Yara]
-    | Type[Suricata]
-    | Type[Sigma]
-    | Type[Query]
-    | Type[ForensicArtifact]
-)
+                    )
\ No newline at end of file
diff --git a/core/schemas/indicators/__init__.py b/core/schemas/indicators/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/core/schemas/indicators/forensicartifact.py b/core/schemas/indicators/forensicartifact.py
new file mode 100644
index 000000000..c21930315
--- /dev/null
+++ b/core/schemas/indicators/forensicartifact.py
@@ -0,0 +1,177 @@
+import io
+import logging
+import re
+from typing import ClassVar, Literal
+
+import yaml
+from artifacts import definitions, reader, writer
+from artifacts import errors as artifacts_errors
+from pydantic import field_validator
+
+from core.schemas import indicator
+from core.schemas.indicators import regex
+
+
+class ForensicArtifact(indicator.Indicator):
+    """Represents a Forensic Artifact
+
+    As defined in https://github.com/ForensicArtifacts/artifacts
+    """
+
+    _type_filter: ClassVar[str] = indicator.IndicatorType.forensicartifact
+    type: Literal[indicator.IndicatorType.forensicartifact] = indicator.IndicatorType.forensicartifact
+
+    sources: list[dict] = []
+    aliases: list[str] = []
+    supported_os: list[str] = []
+
+    def match(self, value: str) -> indicator.IndicatorMatch | None:
+        raise NotImplementedError
+
+    @field_validator("pattern")
+    @classmethod
+    def validate_artifact(cls, value) -> str:
+        artifact_reader = reader.YamlArtifactsReader()
+        try:
+            list(artifact_reader.ReadFileObject(io.StringIO(value)))
+        except artifacts_errors.FormatError as error:
+            raise ValueError(f"Invalid ForensicArtifact YAML: {error}")
+        return value
+
+    @classmethod
+    def from_yaml_string(
+        cls, yaml_string: str, update_parents: bool = False
+    ) -> list["ForensicArtifact"]:
+        artifact_reader = reader.YamlArtifactsReader()
+        artifact_writer = writer.YamlArtifactsWriter()
+
+        artifacts_dict = {}
+
+        for definition in artifact_reader.ReadFileObject(io.StringIO(yaml_string)):
+            definition_dict = definition.AsDict()
+            definition_dict["description"] = definition_dict.pop("doc")
+            if definition.urls:
+                definition_dict["description"] += "\n\nURLs:\n"
+                definition_dict["description"] += " ".join(
+                    [f"* {url}\n" for url in definition.urls]
+                )
+            definition_dict["pattern"] = artifact_writer.FormatArtifacts([definition])
+            definition_dict["location"] = "host"
+            definition_dict["diamond"] = indicator.DiamondModel.victim
+            definition_dict["relevant_tags"] = [definition_dict["name"]]
+            forensic_indicator = cls(**definition_dict).save()
+            artifacts_dict[definition.name] = forensic_indicator
+
+        if update_parents:
+            for artifact in artifacts_dict.values():
+                artifact.update_parents(artifacts_dict)
+
+        return list(artifacts_dict.values())
+
+    def update_yaml(self):
+        artifact_reader = reader.YamlArtifactsReader()
+        definition_dict = next(
+            artifact_reader.ReadFileObject(io.StringIO(self.pattern))
+        ).AsDict()
+        definition_dict["doc"] = self.description.split("\n\nURLs:")[0]
+        definition_dict["name"] = self.name
+        definition_dict["supported_os"] = self.supported_os
+        self.pattern = yaml.safe_dump(definition_dict)
+
+    def update_parents(self, artifacts_dict: dict[str, "ForensicArtifact"]) -> None:
+        for source in self.sources:
+            if not source["type"] == definitions.TYPE_INDICATOR_ARTIFACT_GROUP:
+                continue
+            for child_name in source["attributes"]["names"]:
+                child = artifacts_dict.get(child_name)
+                if not child:
+                    logging.error(f"Missing child {child_name} for {self.name}")
+                    continue
+
+                add_tags = set(self.relevant_tags + [self.name])
+                child.relevant_tags = list(add_tags | set(child.relevant_tags))
+                child.save()
+                child.link_to(
+                    self,
+                    "included in",
+                    f"Included in ForensicArtifact definition for {self.name}",
+                )
+
+    def save_indicators(self, create_links: bool = False):
+        indicators = []
+        for source in self.sources:
+            if source["type"] == definitions.TYPE_INDICATOR_FILE:
+                for path in source["attributes"]["paths"]:
+                    # TODO: consider using https://github.com/log2timeline/dfvfs/blob/main/dfvfs/lib/glob2regex.py
+                    pattern = ARTIFACT_INTERPOLATION_RE.sub("*", path)
+                    pattern = re.escape(pattern).replace("\\*", ".*")
+                    # Account for different path separators
+                    pattern = re.sub(r"\\\\", r"[\\|/]", pattern)
+                    regex_indicator = regex.Regex.find(name=path)
+                    if not regex_indicator:
+                        try:
+                            regex_indicator = regex.Regex(
+                                name=path,
+                                pattern=pattern,
+                                location="filesystem",
+                                diamond=indicator.DiamondModel.victim,
+                                relevant_tags=self.relevant_tags,
+                            ).save()
+                            indicators.append(regex_indicator)
+                        except Exception as error:
+                            logging.error(
+                                f"Failed to create indicator for {path} (was: {source['attributes']['paths']}): {error}"
+                            )
+                            continue
+
+                    else:
+                        regex_indicator.relevant_tags = list(
+                            set(regex_indicator.relevant_tags + self.relevant_tags)
+                        )
+                        regex_indicator.save()
+        if source["type"] == definitions.TYPE_INDICATOR_WINDOWS_REGISTRY_KEY:
+            for key in source["attributes"]["keys"]:
+                pattern = re.sub(r"\\\*$", "", key)
+                pattern = ARTIFACT_INTERPOLATION_RE.sub("*", pattern)
+                pattern = re.escape(pattern)
+                pattern = pattern.replace(
+                    "HKEY_USERS\\\\\\*",
+                    r"(HKEY_USERS\\*|HKEY_CURRENT_USER)",
+                )
+                pattern = pattern.replace("*", r".*").replace("?", r".")
+                if "CurrentControlSet" in pattern:
+                    pattern = pattern.replace(
+                        "CurrentControlSet", "(CurrentControlSet|ControlSet[0-9]+)"
+                    )
+                    pattern = pattern.replace("HKEY_LOCAL_MACHINE\\\\System\\\\", "")
+
+                regex_indicator = regex.Regex.find(name=key)
+
+                if not regex_indicator:
+                    try:
+                        regex_indicator = regex.Regex(
+                            name=key,
+                            pattern=pattern,
+                            location="registry",
+                            diamond=indicator.DiamondModel.victim,
+                            relevant_tags=self.relevant_tags,
+                        ).save()
+                        indicators.append(regex_indicator)
+                    except Exception as error:
+                        logging.error(
+                            f"Failed to create indicator for {key} (was: {source['attributes']['keys']}): {error}"
+                        )
+                        continue
+                else:
+                    regex_indicator.relevant_tags = list(
+                        set(regex_indicator.relevant_tags + self.relevant_tags)
+                    )
+                    regex_indicator.save()
+        if create_links:
+            for indicator_obj in indicators:
+                indicator_obj.link_to(self, "indicates", f"Indicates {indicator_obj.name}")
+        return indicators
+
+
+ARTIFACT_INTERPOLATION_RE = re.compile(r"%%[a-z._]+%%")
+ARTIFACT_INTERPOLATION_RE_HKEY_USERS = re.compile(r"HKEY_USERS\\%%users.sid%%")
diff --git a/core/schemas/indicators/private/.gitignore b/core/schemas/indicators/private/.gitignore
new file mode 100644
index 000000000..a91bcd1d7
--- /dev/null
+++ b/core/schemas/indicators/private/.gitignore
@@ -0,0 +1,4 @@
+*
+!.gitignore
+!README.md
+!__init__.py
\ No newline at end of file
diff --git a/core/schemas/indicators/private/README.md b/core/schemas/indicators/private/README.md
new file mode 100644
index 000000000..4f9a2488b
--- /dev/null
+++ b/core/schemas/indicators/private/README.md
@@ -0,0 +1,2 @@
+### Private entities
+This directory is where you should place your private entities. It could be named anything else, but this one has a `.gitignore` so you don't mess things up. ;-)
diff --git a/core/schemas/indicators/private/__init__.py b/core/schemas/indicators/private/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/core/schemas/indicators/query.py b/core/schemas/indicators/query.py
new file mode 100644
index 000000000..89acb3d14
--- /dev/null
+++ b/core/schemas/indicators/query.py
@@ -0,0 +1,16 @@
+from typing import ClassVar, Literal
+
+from core.schemas import indicator
+
+
+class Query(indicator.Indicator):
+    """Represents a query that can be sent to another system."""
+
+    _type_filter: ClassVar[str] = indicator.IndicatorType.query
+    type: Literal[indicator.IndicatorType.query] = indicator.IndicatorType.query
+
+    query_type: str
+    target_systems: list[str] = []
+
+    def match(self, value: str) -> indicator.IndicatorMatch | None:
+        return None
diff --git a/core/schemas/indicators/regex.py b/core/schemas/indicators/regex.py
new file mode 100644
index 000000000..67d58d945
--- /dev/null
+++ b/core/schemas/indicators/regex.py
@@ -0,0 +1,33 @@
+import re
+from typing import ClassVar, Literal
+
+from pydantic import PrivateAttr, field_validator
+
+from core.schemas import indicator
+
+
+class Regex(indicator.Indicator):
+    _type_filter: ClassVar[str] = indicator.IndicatorType.regex
+    _compiled_pattern: re.Pattern | None = PrivateAttr(None)
+    type: Literal[indicator.IndicatorType.regex] = indicator.IndicatorType.regex
+
+    @property
+    def compiled_pattern(self):
+        if not self._compiled_pattern:
+            self._compiled_pattern = re.compile(self.pattern)
+        return self._compiled_pattern
+
+    @field_validator("pattern")
+    @classmethod
+    def validate_regex(cls, value) -> str:
+        try:
+            re.compile(value)
+        except re.error as error:
+            raise ValueError(f"Invalid regex pattern: {error}")
+        return value
+
+    def match(self, value: str) -> indicator.IndicatorMatch | None:
+        result = self.compiled_pattern.search(value)
+        if result:
+            return indicator.IndicatorMatch(name=self.name, match=result.group())
+        return None
\ No newline at end of file
diff --git a/core/schemas/indicators/sigma.py b/core/schemas/indicators/sigma.py
new file mode 100644
index 000000000..4777820ff
--- /dev/null
+++ b/core/schemas/indicators/sigma.py
@@ -0,0 +1,16 @@
+from typing import ClassVar, Literal
+
+from core.schemas import indicator
+
+
+class Sigma(indicator.Indicator):
+    """Represents a Sigma rule.
+
+    Parsing and matching is yet TODO.
+    """
+
+    _type_filter: ClassVar[str] = indicator.IndicatorType.sigma
+    type: Literal[indicator.IndicatorType.sigma] = indicator.IndicatorType.sigma
+
+    def match(self, value: str) -> indicator.IndicatorMatch | None:
+        raise NotImplementedError
diff --git a/core/schemas/indicators/suricata.py b/core/schemas/indicators/suricata.py
new file mode 100644
index 000000000..6d62f6590
--- /dev/null
+++ b/core/schemas/indicators/suricata.py
@@ -0,0 +1,38 @@
+import logging
+from typing import ClassVar, List, Literal
+
+from idstools import rule
+from pydantic import field_validator
+
+from core.schemas import indicator
+
+
+class Suricata(indicator.Indicator):
+    """Represents a Suricata rule.
+
+    Parsing and matching is yet TODO.
+    """
+
+    _type_filter: ClassVar[str] = indicator.IndicatorType.suricata
+    type: Literal[indicator.IndicatorType.suricata] = indicator.IndicatorType.suricata
+    sid: int = 0
+    metadata: List[str] = []
+    references: List[str] = []
+
+    def match(self, value: str) -> indicator.IndicatorMatch | None:
+        raise NotImplementedError
+
+    @field_validator("pattern")
+    @classmethod
+    def validate_rules(cls, value) -> str:
+        try:
+            rule.parse(value)
+        except Exception as e:
+            raise ValueError(f"invalid {cls.pattern} {e}")
+        return value
+
+    def parse(self) -> rule.Rule | None:
+        try:
+            return rule.parse(self.pattern)
+        except Exception as e:
+            logging.error(f" Error parsing {self.pattern} {e}")
diff --git a/core/schemas/indicators/yara.py b/core/schemas/indicators/yara.py
new file mode 100644
index 000000000..7f1d1b1a2
--- /dev/null
+++ b/core/schemas/indicators/yara.py
@@ -0,0 +1,16 @@
+from typing import ClassVar, Literal
+
+from core.schemas import indicator
+
+
+class Yara(indicator.Indicator):
+    """Represents a Yara rule.
+
+    Parsing and matching is yet TODO.
+    """
+
+    _type_filter: ClassVar[str] = indicator.IndicatorType.yara
+    type: Literal[indicator.IndicatorType.yara] = indicator.IndicatorType.yara
+
+    def match(self, value: str) -> indicator.IndicatorMatch | None:
+        raise NotImplementedError
diff --git a/core/schemas/observable.py b/core/schemas/observable.py
index ded73db3f..277b13699 100644
--- a/core/schemas/observable.py
+++ b/core/schemas/observable.py
@@ -21,7 +21,7 @@
 class ObservableType(str, Enum):
     ...
 
-
+ObservableTypes = ()
 TYPE_MAPPING = {}
 
 
@@ -127,10 +127,6 @@ def delete_context(
                 break
         return self.save()
 
-
-TYPE_MAPPING.update({"observable": Observable, "observables": Observable})
-
-
 def find_type(value: str) -> ObservableType | None:
     for obs_type, obj in TYPE_MAPPING.items():
         if obj.is_valid(value):
diff --git a/core/schemas/observables/asn.py b/core/schemas/observables/asn.py
index 8866e010c..f4c3cf240 100644
--- a/core/schemas/observables/asn.py
+++ b/core/schemas/observables/asn.py
@@ -1,7 +1,9 @@
+from typing import Literal
+
 from core.schemas import observable
 
 
 class ASN(observable.Observable):
-    type: observable.ObservableType = observable.ObservableType.asn
+    type: Literal[observable.ObservableType.asn] = observable.ObservableType.asn
     country: str | None = None
     description: str | None = None
diff --git a/core/schemas/observables/bic.py b/core/schemas/observables/bic.py
index 68b7c8ea3..8e631c8ad 100644
--- a/core/schemas/observables/bic.py
+++ b/core/schemas/observables/bic.py
@@ -1,4 +1,5 @@
 import re
+from typing import Literal
 
 from core.schemas import observable
 
@@ -6,7 +7,7 @@
 
 
 class BIC(observable.Observable):
-    type: observable.ObservableType = observable.ObservableType.bic
+    type: Literal[observable.ObservableType.bic] = observable.ObservableType.bic
 
     @staticmethod
     def is_valid(value: str) -> bool:
diff --git a/core/schemas/observables/certificate.py b/core/schemas/observables/certificate.py
index a9433988e..7dfcf3623 100644
--- a/core/schemas/observables/certificate.py
+++ b/core/schemas/observables/certificate.py
@@ -1,5 +1,6 @@
 import datetime
 import hashlib
+from typing import Literal
 
 from pydantic import Field
 
@@ -21,7 +22,7 @@ class Certificate(observable.Observable):
         fingerprint: the certificate fingerprint.
     """
 
-    type: observable.ObservableType = observable.ObservableType.certificate
+    type: Literal[observable.ObservableType.certificate] = observable.ObservableType.certificate
     last_seen: datetime.datetime = Field(default_factory=now)
     first_seen: datetime.datetime = Field(default_factory=now)
     issuer: str | None = None
diff --git a/core/schemas/observables/cidr.py b/core/schemas/observables/cidr.py
index 97433ff8e..981f47a96 100644
--- a/core/schemas/observables/cidr.py
+++ b/core/schemas/observables/cidr.py
@@ -1,5 +1,7 @@
+from typing import Literal
+
 from core.schemas import observable
 
 
 class CIDR(observable.Observable):
-    type: observable.ObservableType = observable.ObservableType.cidr
+    type: Literal[observable.ObservableType.cidr] = observable.ObservableType.cidr
diff --git a/core/schemas/observables/command_line.py b/core/schemas/observables/command_line.py
index c2119ab7f..4a9e86dbc 100644
--- a/core/schemas/observables/command_line.py
+++ b/core/schemas/observables/command_line.py
@@ -1,5 +1,7 @@
+from typing import Literal
+
 from core.schemas import observable
 
 
 class CommandLine(observable.Observable):
-    type: observable.ObservableType = observable.ObservableType.command_line
+    type: Literal[observable.ObservableType.command_line] = observable.ObservableType.command_line
diff --git a/core/schemas/observables/docker_image.py b/core/schemas/observables/docker_image.py
index 3d166e4eb..ca044f1d1 100644
--- a/core/schemas/observables/docker_image.py
+++ b/core/schemas/observables/docker_image.py
@@ -1,5 +1,7 @@
+from typing import Literal
+
 from core.schemas import observable
 
 
 class DockerImage(observable.Observable):
-    type: observable.ObservableType = observable.ObservableType.docker_image
+    type: Literal[observable.ObservableType.docker_image] = observable.ObservableType.docker_image
diff --git a/core/schemas/observables/email.py b/core/schemas/observables/email.py
index ba619ec3f..ae0c41d13 100644
--- a/core/schemas/observables/email.py
+++ b/core/schemas/observables/email.py
@@ -1,10 +1,12 @@
+from typing import Literal
+
 import validators
 
 from core.schemas import observable
 
 
 class Email(observable.Observable):
-    type: observable.ObservableType = observable.ObservableType.email
+    type: Literal[observable.ObservableType.email] = observable.ObservableType.email
 
     @staticmethod
     def is_valid(value: str) -> bool:
diff --git a/core/schemas/observables/file.py b/core/schemas/observables/file.py
index de2d8e043..82625b7fa 100644
--- a/core/schemas/observables/file.py
+++ b/core/schemas/observables/file.py
@@ -1,3 +1,5 @@
+from typing import Literal
+
 from core.schemas import observable
 
 
@@ -8,7 +10,7 @@ class File(observable.Observable):
     Value should to be in the form FILE:.
     """
 
-    type: observable.ObservableType = observable.ObservableType.file
+    type: Literal[observable.ObservableType.file] = observable.ObservableType.file
     name: str | None = None
     size: int | None = None
     sha256: str | None = None
diff --git a/core/schemas/observables/generic.py b/core/schemas/observables/generic.py
index a8e9bbe68..cbdae38cd 100644
--- a/core/schemas/observables/generic.py
+++ b/core/schemas/observables/generic.py
@@ -6,4 +6,4 @@
 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
+    type: Literal[observable.ObservableType.generic] = observable.ObservableType.generic
diff --git a/core/schemas/observables/hostname.py b/core/schemas/observables/hostname.py
index db0b892a6..812e03bd9 100644
--- a/core/schemas/observables/hostname.py
+++ b/core/schemas/observables/hostname.py
@@ -1,10 +1,12 @@
+from typing import Literal
+
 import validators
 
 from core.schemas import observable
 
 
 class Hostname(observable.Observable):
-    type: observable.ObservableType = observable.ObservableType.hostname
+    type: Literal[observable.ObservableType.hostname] = observable.ObservableType.hostname
 
     @staticmethod
     def is_valid(value: str) -> bool:
diff --git a/core/schemas/observables/iban.py b/core/schemas/observables/iban.py
index 9641080a2..01cf3c20e 100644
--- a/core/schemas/observables/iban.py
+++ b/core/schemas/observables/iban.py
@@ -1,10 +1,12 @@
+from typing import Literal
+
 import validators
 
 from core.schemas import observable
 
 
 class IBAN(observable.Observable):
-    type: observable.ObservableType = observable.ObservableType.iban
+    type: Literal[observable.ObservableType.iban] = observable.ObservableType.iban
 
     @staticmethod
     def is_valid(value: str) -> bool:
diff --git a/core/schemas/observables/imphash.py b/core/schemas/observables/imphash.py
index 316e4ce29..ee302549f 100644
--- a/core/schemas/observables/imphash.py
+++ b/core/schemas/observables/imphash.py
@@ -1,5 +1,7 @@
+from typing import Literal
+
 from core.schemas import observable
 
 
 class Imphash(observable.Observable):
-    type: observable.ObservableType = observable.ObservableType.imphash
+    type: Literal[observable.ObservableType.imphash] = observable.ObservableType.imphash
diff --git a/core/schemas/observables/ipv4.py b/core/schemas/observables/ipv4.py
index 9eab620c3..ec3a23f21 100644
--- a/core/schemas/observables/ipv4.py
+++ b/core/schemas/observables/ipv4.py
@@ -1,10 +1,12 @@
+from typing import Literal
+
 import validators
 
 from core.schemas import observable
 
 
 class IPv4(observable.Observable):
-    type: observable.ObservableType = observable.ObservableType.ipv4
+    type: Literal[observable.ObservableType.ipv4] = observable.ObservableType.ipv4
 
     @staticmethod
     def is_valid(value: str) -> bool:
diff --git a/core/schemas/observables/ipv6.py b/core/schemas/observables/ipv6.py
index f1fae0be9..0619eacac 100644
--- a/core/schemas/observables/ipv6.py
+++ b/core/schemas/observables/ipv6.py
@@ -1,10 +1,12 @@
+from typing import Literal
+
 import validators
 
 from core.schemas import observable
 
 
 class IPv6(observable.Observable):
-    type: observable.ObservableType = observable.ObservableType.ipv6
+    type: Literal[observable.ObservableType.ipv6] = observable.ObservableType.ipv6
 
     @staticmethod
     def is_valid(value: str) -> bool:
diff --git a/core/schemas/observables/ja3.py b/core/schemas/observables/ja3.py
index 5b93f8755..774b36f59 100644
--- a/core/schemas/observables/ja3.py
+++ b/core/schemas/observables/ja3.py
@@ -1,5 +1,7 @@
+from typing import Literal
+
 from core.schemas import observable
 
 
 class JA3(observable.Observable):
-    type: observable.ObservableType = observable.ObservableType.ja3
+    type: Literal[observable.ObservableType.ja3] = observable.ObservableType.ja3
diff --git a/core/schemas/observables/jarm.py b/core/schemas/observables/jarm.py
index bcf21a87c..f7a6694ea 100644
--- a/core/schemas/observables/jarm.py
+++ b/core/schemas/observables/jarm.py
@@ -1,5 +1,7 @@
+from typing import Literal
+
 from core.schemas import observable
 
 
 class JARM(observable.Observable):
-    type: observable.ObservableType = observable.ObservableType.jarm
+    type: Literal[observable.ObservableType.jarm] = observable.ObservableType.jarm
diff --git a/core/schemas/observables/mac_address.py b/core/schemas/observables/mac_address.py
index 5165d0260..d900da416 100644
--- a/core/schemas/observables/mac_address.py
+++ b/core/schemas/observables/mac_address.py
@@ -1,5 +1,7 @@
+from typing import Literal
+
 from core.schemas import observable
 
 
 class MacAddress(observable.Observable):
-    type: observable.ObservableType = observable.ObservableType.mac_address
+    type: Literal[observable.ObservableType.mac_address] = observable.ObservableType.mac_address
diff --git a/core/schemas/observables/md5.py b/core/schemas/observables/md5.py
index 9e783c55e..c48101b15 100644
--- a/core/schemas/observables/md5.py
+++ b/core/schemas/observables/md5.py
@@ -1,10 +1,12 @@
+from typing import Literal
+
 import validators
 
 from core.schemas import observable
 
 
 class MD5(observable.Observable):
-    type: observable.ObservableType = observable.ObservableType.md5
+    type: Literal[observable.ObservableType.md5] = observable.ObservableType.md5
 
     @staticmethod
     def is_valid(value: str) -> bool:
diff --git a/core/schemas/observables/mutex.py b/core/schemas/observables/mutex.py
index 5124aed83..52cec411c 100644
--- a/core/schemas/observables/mutex.py
+++ b/core/schemas/observables/mutex.py
@@ -1,5 +1,7 @@
+from typing import Literal
+
 from core.schemas import observable
 
 
 class Mutex(observable.Observable):
-    type: observable.ObservableType = observable.ObservableType.mutex
+    type: Literal[observable.ObservableType.mutex] = observable.ObservableType.mutex
diff --git a/core/schemas/observables/named_pipe.py b/core/schemas/observables/named_pipe.py
index 314e6e2ea..b126fe479 100644
--- a/core/schemas/observables/named_pipe.py
+++ b/core/schemas/observables/named_pipe.py
@@ -1,5 +1,7 @@
+from typing import Literal
+
 from core.schemas import observable
 
 
 class NamedPipe(observable.Observable):
-    type: observable.ObservableType = observable.ObservableType.named_pipe
+    type: Literal[observable.ObservableType.named_pipe] = observable.ObservableType.named_pipe
diff --git a/core/schemas/observables/package.py b/core/schemas/observables/package.py
index 23cff722f..a997acc02 100644
--- a/core/schemas/observables/package.py
+++ b/core/schemas/observables/package.py
@@ -1,7 +1,9 @@
+from typing import Literal
+
 from core.schemas import observable
 
 
 class Package(observable.Observable):
-    type: observable.ObservableType = observable.ObservableType.package
+    type: Literal[observable.ObservableType.package] = observable.ObservableType.package
     version: str = None
     regitry_type: str = None
diff --git a/core/schemas/observables/path.py b/core/schemas/observables/path.py
index 5d90025e3..b9035f01a 100644
--- a/core/schemas/observables/path.py
+++ b/core/schemas/observables/path.py
@@ -1,4 +1,5 @@
 import re
+from typing import Literal
 
 from core.schemas import observable
 
@@ -13,7 +14,7 @@ def path_validator(value):
 
 
 class Path(observable.Observable):
-    type: observable.ObservableType = observable.ObservableType.path
+    type: Literal[observable.ObservableType.path] = observable.ObservableType.path
 
     @staticmethod
     def is_valid(value: str) -> bool:
diff --git a/core/schemas/observables/registry_key.py b/core/schemas/observables/registry_key.py
index c9d7b5956..ccc07baca 100644
--- a/core/schemas/observables/registry_key.py
+++ b/core/schemas/observables/registry_key.py
@@ -26,7 +26,7 @@ class RegistryKey(observable.Observable):
         path_file: The filesystem path to the file that contains the registry key value.
     """
 
-    type: observable.ObservableType = observable.ObservableType.registry_key
+    type: Literal[observable.ObservableType.registry_key] = observable.ObservableType.registry_key
     key: str
     data: bytes
     hive: RegistryHive
diff --git a/core/schemas/observables/sha1.py b/core/schemas/observables/sha1.py
index db60b5f09..cc44ee7ac 100644
--- a/core/schemas/observables/sha1.py
+++ b/core/schemas/observables/sha1.py
@@ -1,10 +1,12 @@
+from typing import Literal
+
 import validators
 
 from core.schemas import observable
 
 
 class SHA1(observable.Observable):
-    type: observable.ObservableType = observable.ObservableType.sha1
+    type: Literal[observable.ObservableType.sha1] = observable.ObservableType.sha1
 
     @staticmethod
     def is_valid(value: str) -> bool:
diff --git a/core/schemas/observables/sha256.py b/core/schemas/observables/sha256.py
index 51910b03c..929ab5f47 100644
--- a/core/schemas/observables/sha256.py
+++ b/core/schemas/observables/sha256.py
@@ -1,10 +1,12 @@
+from typing import Literal
+
 import validators
 
 from core.schemas import observable
 
 
 class SHA256(observable.Observable):
-    type: observable.ObservableType = observable.ObservableType.sha256
+    type: Literal[observable.ObservableType.sha256] = observable.ObservableType.sha256
 
     @staticmethod
     def is_valid(value: str) -> bool:
diff --git a/core/schemas/observables/ssdeep.py b/core/schemas/observables/ssdeep.py
index 35b68de5c..74c2e0dfd 100644
--- a/core/schemas/observables/ssdeep.py
+++ b/core/schemas/observables/ssdeep.py
@@ -1,5 +1,7 @@
+from typing import Literal
+
 from core.schemas import observable
 
 
 class Ssdeep(observable.Observable):
-    type: observable.ObservableType = observable.ObservableType.ssdeep
+    type: Literal[observable.ObservableType.ssdeep] = observable.ObservableType.ssdeep
diff --git a/core/schemas/observables/tlsh.py b/core/schemas/observables/tlsh.py
index edb97be94..2cebcc6d5 100644
--- a/core/schemas/observables/tlsh.py
+++ b/core/schemas/observables/tlsh.py
@@ -1,5 +1,7 @@
+from typing import Literal
+
 from core.schemas import observable
 
 
 class TLSH(observable.Observable):
-    type: observable.ObservableType = observable.ObservableType.tlsh
+    type: Literal[observable.ObservableType.tlsh] = observable.ObservableType.tlsh
diff --git a/core/schemas/observables/url.py b/core/schemas/observables/url.py
index 25c7ceb4f..3ba1f40b9 100644
--- a/core/schemas/observables/url.py
+++ b/core/schemas/observables/url.py
@@ -1,10 +1,12 @@
+from typing import Literal
+
 import validators
 
 from core.schemas import observable
 
 
 class Url(observable.Observable):
-    type: observable.ObservableType = observable.ObservableType.url
+    type: Literal[observable.ObservableType.url] = observable.ObservableType.url
 
     @staticmethod
     def is_valid(value: str) -> bool:
diff --git a/core/schemas/observables/user_account.py b/core/schemas/observables/user_account.py
index 9c77009e5..e36656c0c 100644
--- a/core/schemas/observables/user_account.py
+++ b/core/schemas/observables/user_account.py
@@ -14,7 +14,7 @@ class UserAccount(observable.Observable):
     Value should to be in the form :.
     """
 
-    type: observable.ObservableType = observable.ObservableType.user_account
+    type: Literal[observable.ObservableType.user_account] = observable.ObservableType.user_account
     user_id: str | None = None
     credential: str | None = None
     account_login: str | None = None
diff --git a/core/schemas/observables/user_agent.py b/core/schemas/observables/user_agent.py
index 40c0d8cdc..2fcb83d36 100644
--- a/core/schemas/observables/user_agent.py
+++ b/core/schemas/observables/user_agent.py
@@ -1,5 +1,7 @@
+from typing import Literal
+
 from core.schemas import observable
 
 
 class UserAgent(observable.Observable):
-    type: observable.ObservableType = observable.ObservableType.user_agent
+    type: Literal[observable.ObservableType.user_agent] = observable.ObservableType.user_agent
diff --git a/core/schemas/observables/wallet.py b/core/schemas/observables/wallet.py
index 4670f5ca3..cd77eaf73 100644
--- a/core/schemas/observables/wallet.py
+++ b/core/schemas/observables/wallet.py
@@ -1,3 +1,5 @@
+from typing import Literal
+
 from core.schemas import observable
 
 
@@ -8,6 +10,6 @@ class Wallet(observable.Observable):
     Value should be in the form :
.
     """
 
-    type: observable.ObservableType = observable.ObservableType.wallet
+    type: Literal[observable.ObservableType.wallet] = observable.ObservableType.wallet
     coin: str | None = None
     address: str | None = None
diff --git a/core/web/apiv2/entities.py b/core/web/apiv2/entities.py
index 42e6c60e6..0597b22cc 100644
--- a/core/web/apiv2/entities.py
+++ b/core/web/apiv2/entities.py
@@ -10,14 +10,14 @@
 class NewEntityRequest(BaseModel):
     model_config = ConfigDict(extra="forbid")
 
-    entity: EntityTypes
+    entity: EntityTypes = Field(discriminator="type")
     tags: conlist(str, max_length=MAX_TAGS_REQUEST) = []
 
 
 class PatchEntityRequest(BaseModel):
     model_config = ConfigDict(extra="forbid")
 
-    entity: EntityTypes
+    entity: EntityTypes = Field(discriminator="type")
 
 
 class EntitySearchRequest(BaseModel):
diff --git a/core/web/apiv2/graph.py b/core/web/apiv2/graph.py
index 38ab2026b..e009eccbc 100644
--- a/core/web/apiv2/graph.py
+++ b/core/web/apiv2/graph.py
@@ -111,7 +111,7 @@ class GraphSearchResponse(BaseModel):
         | entity.EntityTypes
         | indicator.IndicatorTypes
         | tag.Tag
-        | dfiq.DFIQTypes,
+        | dfiq.DFIQTypes
     ]
     paths: list[list[graph.Relationship | graph.TagRelationship]]
     total: int
diff --git a/core/web/apiv2/indicators.py b/core/web/apiv2/indicators.py
index 709802c4f..3c6de8397 100644
--- a/core/web/apiv2/indicators.py
+++ b/core/web/apiv2/indicators.py
@@ -3,11 +3,11 @@
 
 from core.schemas import graph
 from core.schemas.indicator import (
-    ForensicArtifact,
     Indicator,
     IndicatorType,
     IndicatorTypes,
 )
+from core.schemas.indicators import forensicartifact
 from core.schemas.tag import MAX_TAGS_REQUEST
 
 
@@ -79,7 +79,7 @@ async def patch(request: PatchIndicatorRequest, indicator_id) -> IndicatorTypes:
 
     if db_indicator.type == IndicatorType.forensicartifact:
         if db_indicator.pattern != request.indicator.pattern:
-            return ForensicArtifact.from_yaml_string(request.indicator.pattern)[0]
+            return forensicartifact.ForensicArtifact.from_yaml_string(request.indicator.pattern)[0]
 
     update_data = request.indicator.model_dump(exclude_unset=True)
     updated_indicator = db_indicator.model_copy(update=update_data)
diff --git a/core/web/apiv2/observables.py b/core/web/apiv2/observables.py
index 88704d11b..956d7878b 100644
--- a/core/web/apiv2/observables.py
+++ b/core/web/apiv2/observables.py
@@ -4,20 +4,14 @@
 from pydantic import BaseModel, ConfigDict, Field, conlist, field_validator
 
 from core.schemas import graph
-from core.schemas.observable import TYPE_MAPPING, Observable, ObservableType
+from core.schemas.observable import (
+    TYPE_MAPPING,
+    Observable,
+    ObservableType,
+    ObservableTypes,
+)
 from core.schemas.tag import MAX_TAG_LENGTH, MAX_TAGS_REQUEST
 
-ObservableTypes = ()
-
-for key in TYPE_MAPPING:
-    if key in ["observable", "observables"]:
-        continue
-    cls = TYPE_MAPPING[key]
-    if not ObservableTypes:
-        ObservableTypes = cls
-    else:
-        ObservableTypes |= cls
-
 
 class TagRequestMixin(BaseModel):
     tags: conlist(str, max_length=MAX_TAGS_REQUEST) = []
@@ -44,14 +38,13 @@ class NewObservableRequest(TagRequestMixin):
 class NewExtendedObservableRequest(TagRequestMixin):
     model_config = ConfigDict(extra="forbid")
 
-    observable: ObservableTypes = Field(discriminant="type")
+    observable: ObservableTypes = Field(discriminator="type")
 
 
 class PatchObservableRequest(BaseModel):
     model_config = ConfigDict(extra="forbid")
 
-    observable: ObservableTypes = Field(discriminant="type")
-
+    observable: ObservableTypes = Field(discriminator="type")
 
 class NewBulkObservableAddRequest(BaseModel):
     model_config = ConfigDict(extra="forbid")
diff --git a/tests/apiv2/graph.py b/tests/apiv2/graph.py
index c72d5c15f..f8f561c49 100644
--- a/tests/apiv2/graph.py
+++ b/tests/apiv2/graph.py
@@ -7,7 +7,8 @@
 from core import database_arango
 from core.schemas.entities import attack_pattern, malware, threat_actor
 from core.schemas.graph import Relationship
-from core.schemas.indicator import DiamondModel, ForensicArtifact, Query, Regex
+from core.schemas.indicator import DiamondModel
+from core.schemas.indicators import forensicartifact, query, regex
 from core.schemas.observables import hostname, ipv4, url
 from core.schemas.user import UserSensitive
 from core.web import webapp
@@ -28,7 +29,7 @@ def setUp(self) -> None:
         self.observable1 = hostname.Hostname(value="tomchop.me").save()
         self.observable2 = ipv4.IPv4(value="127.0.0.1").save()
         self.entity1 = threat_actor.ThreatActor(name="actor0").save()
-        self.indicator1 = Query(
+        self.indicator1 = query.Query(
             name="query1",
             query_type="opensearch",
             target_systems=["system1"],
@@ -378,7 +379,7 @@ def setUp(self) -> None:
         self.observable2 = hostname.Hostname(value="test2.com").save()
         self.observable3 = url.Url(value="http://test1.com/admin").save()
         self.entity1 = threat_actor.ThreatActor(name="tester").save()
-        self.indicator1 = Regex(
+        self.indicator1 = regex.Regex(
             name="test c2",
             pattern="test[0-9].com",
             location="network",
@@ -580,7 +581,7 @@ def setUp(self) -> None:
 
         self.persistence = attack_pattern.AttackPattern(name="persistence").save()
         self.persistence.tag(["triage"])
-        self.persistence_artifact = ForensicArtifact.from_yaml_string(
+        self.persistence_artifact = forensicartifact.ForensicArtifact.from_yaml_string(
             """doc: Crontab files.
 name: LinuxCronTabs
 sources:
diff --git a/tests/apiv2/indicators.py b/tests/apiv2/indicators.py
index a021d4d1b..5ddfc8d58 100644
--- a/tests/apiv2/indicators.py
+++ b/tests/apiv2/indicators.py
@@ -7,6 +7,7 @@
 
 from core import database_arango
 from core.schemas import indicator
+from core.schemas.indicators import query, regex
 from core.schemas.user import UserSensitive
 from core.web import webapp
 
@@ -23,14 +24,14 @@ def setUp(self) -> None:
             "/api/v2/auth/api-token", headers={"x-yeti-apikey": user.api_key}
         ).json()
         client.headers = {"Authorization": "Bearer " + token_data["access_token"]}
-        self.indicator1 = indicator.Regex(
+        self.indicator1 = regex.Regex(
             name="hex",
             pattern="[0-9a-f]",
             location="filesystem",
             diamond=indicator.DiamondModel.capability,
         ).save()
         self.indicator1.tag(["hextag"])
-        self.indicator2 = indicator.Regex(
+        self.indicator2 = regex.Regex(
             name="localhost",
             pattern="127.0.0.1",
             location="network",
@@ -101,7 +102,7 @@ def test_search_indicators_subfields(self):
         self.assertEqual(data["indicators"][0]["type"], "regex")
 
     def test_search_indicators_by_alias(self):
-        indicator.Query(
+        query.Query(
             name="query1",
             pattern="SELECT * FROM table",
             location="database",
diff --git a/tests/schemas/fixture.py b/tests/schemas/fixture.py
index 8267393d1..b6e068061 100644
--- a/tests/schemas/fixture.py
+++ b/tests/schemas/fixture.py
@@ -2,7 +2,8 @@
 
 from core import database_arango
 from core.schemas.entities import investigation, malware, threat_actor
-from core.schemas.indicator import DiamondModel, Query, Regex
+from core.schemas.indicator import DiamondModel
+from core.schemas.indicators import query, regex
 from core.schemas.observables import (
     bic,
     generic,
@@ -61,18 +62,18 @@ def test_something(self):
         ta.tag(["Hack!ré T@ëst"])
         ta.link_to(hacker, "uses", "Uses domain")
 
-        regex = Regex(
+        regex_indicator = regex.Regex(
             name="hex",
             pattern="/tmp/[0-9a-f]",
             location="bodyfile",
             diamond=DiamondModel.capability,
         ).save()
-        regex.link_to(hacker, "indicates", "Domain dropped by this regex")
+        regex_indicator.link_to(hacker, "indicates", "Domain dropped by this regex")
         xmrig = malware.Malware(name="xmrig").save()
         xmrig.tag(["xmrig"])
-        regex.link_to(xmrig, "indicates", "Usual name for dropped binary")
+        regex_indicator.link_to(xmrig, "indicates", "Usual name for dropped binary")
 
-        Query(
+        query.Query(
             name="ssh succesful logins",
             location="syslogs",
             diamond=DiamondModel.capability,
diff --git a/tests/schemas/indicator.py b/tests/schemas/indicator.py
index ee1505ee0..8ac9b2429 100644
--- a/tests/schemas/indicator.py
+++ b/tests/schemas/indicator.py
@@ -1,13 +1,8 @@
 import unittest
 
 from core import database_arango
-from core.schemas.indicator import (
-    DiamondModel,
-    ForensicArtifact,
-    Indicator,
-    Query,
-    Regex,
-)
+from core.schemas.indicator import DiamondModel, Indicator
+from core.schemas.indicators import forensicartifact, query, regex
 
 
 class IndicatorTest(unittest.TestCase):
@@ -19,7 +14,7 @@ def tearDown(self) -> None:
         database_arango.db.clear()
 
     def test_create_indicator(self) -> None:
-        result = Regex(
+        result = regex.Regex(
             name="regex1",
             pattern="asd",
             location="any",
@@ -31,7 +26,7 @@ def test_create_indicator(self) -> None:
         self.assertEqual(result.type, "regex")
 
     def test_filter_entities_different_types(self) -> None:
-        regex = Regex(
+        regex_indicator = regex.Regex(
             name="regex1",
             pattern="asd",
             location="any",
@@ -39,53 +34,53 @@ def test_filter_entities_different_types(self) -> None:
         ).save()
 
         all_entities = list(Indicator.list())
-        regex_entities = list(Regex.list())
+        regex_indicators = list(regex.Regex.list())
 
         self.assertEqual(len(all_entities), 1)
-        self.assertEqual(len(regex_entities), 1)
-        self.assertEqual(regex_entities[0].model_dump_json(), regex.model_dump_json())
+        self.assertEqual(len(regex_indicators), 1)
+        self.assertEqual(regex_indicators[0].model_dump_json(), regex_indicator.model_dump_json())
 
     def test_create_indicator_same_name_diff_types(self) -> None:
-        regex = Regex(
+        regex1 = regex.Regex(
             name="persistence1",
             pattern="asd",
             location="any",
             diamond=DiamondModel.capability,
         ).save()
-        regex2 = Query(
+        regex2 = query.Query(
             name="persistence1",
             pattern="asd",
             location="any",
             query_type="query",
             diamond=DiamondModel.capability,
         ).save()
-        self.assertNotEqual(regex.id, regex2.id)
-        r = Regex.find(name="persistence1")
-        q = Query.find(name="persistence1")
+        self.assertNotEqual(regex1.id, regex2.id)
+        r = regex.Regex.find(name="persistence1")
+        q = query.Query.find(name="persistence1")
         self.assertNotEqual(r.id, q.id)
 
     def test_regex_match(self) -> None:
-        regex = Regex(
+        regex_indicator = regex.Regex(
             name="regex1",
             pattern="Ba+dString",
             location="any",
             diamond=DiamondModel.capability,
         ).save()
 
-        result = regex.match("ThisIsAReallyBaaaadStringIsntIt")
+        result = regex_indicator.match("ThisIsAReallyBaaaadStringIsntIt")
         assert result is not None
         self.assertIsNotNone(result)
         self.assertEqual(result.name, "regex1")
         self.assertEqual(result.match, "BaaaadString")
 
     def test_regex_nomatch(self) -> None:
-        regex = Regex(
+        regex_indicator = regex.Regex(
             name="regex1",
             pattern="Blah",
             location="any",
             diamond=DiamondModel.capability,
         ).save()
-        result = regex.match("ThisIsAReallyBaaaadStringIsntIt")
+        result = regex_indicator.match("ThisIsAReallyBaaaadStringIsntIt")
         self.assertIsNone(result)
 
     def test_forensics_artifacts_indicator_extraction_file(self) -> None:
@@ -108,7 +103,7 @@ def test_forensics_artifacts_indicator_extraction_file(self) -> None:
         - Darwin
         - Linux"""
 
-        artifacts = ForensicArtifact.from_yaml_string(pattern)
+        artifacts = forensicartifact.ForensicArtifact.from_yaml_string(pattern)
         db_artifact = artifacts[0]
         self.assertIsNotNone(db_artifact.id)
         self.assertIsNotNone(db_artifact.created)
@@ -173,7 +168,7 @@ def test_forensics_artifacts_indicator_extraction_registry(self) -> None:
 supported_os:
 - Windows"""
 
-        artifacts = ForensicArtifact.from_yaml_string(pattern)
+        artifacts = forensicartifact.ForensicArtifact.from_yaml_string(pattern)
         db_artifact = artifacts[0]
         self.assertIsNotNone(db_artifact.id)
         self.assertIsNotNone(db_artifact.created)
@@ -247,7 +242,7 @@ def test_forensic_artifacts_parent_extraction(self):
     - blah3
 """
 
-        artifacts = ForensicArtifact.from_yaml_string(pattern, update_parents=True)
+        artifacts = forensicartifact.ForensicArtifact.from_yaml_string(pattern, update_parents=True)
         self.assertEqual(len(artifacts), 3)
 
         vertices, _, total = artifacts[0].neighbors()

From a2ea66c497d1afbcbba954231a9545086d1b2162 Mon Sep 17 00:00:00 2001
From: Fred Baguelin 
Date: Tue, 17 Sep 2024 12:54:02 +0200
Subject: [PATCH 20/39] Ruff format

---
 core/schemas/__init__.py                    | 15 ++++++++++++---
 core/schemas/entities/course_of_action.py   |  4 +++-
 core/schemas/entity.py                      |  2 +-
 core/schemas/indicator.py                   |  5 +++--
 core/schemas/indicators/forensicartifact.py |  8 ++++++--
 core/schemas/indicators/regex.py            |  2 +-
 core/schemas/observable.py                  |  2 ++
 core/schemas/observables/certificate.py     |  4 +++-
 core/schemas/observables/command_line.py    |  4 +++-
 core/schemas/observables/docker_image.py    |  4 +++-
 core/schemas/observables/hostname.py        |  4 +++-
 core/schemas/observables/mac_address.py     |  4 +++-
 core/schemas/observables/named_pipe.py      |  4 +++-
 core/schemas/observables/registry_key.py    |  4 +++-
 core/schemas/observables/user_account.py    |  4 +++-
 core/schemas/observables/user_agent.py      |  4 +++-
 core/web/apiv2/graph.py                     |  2 +-
 core/web/apiv2/indicators.py                |  4 +++-
 core/web/apiv2/observables.py               |  1 +
 tests/schemas/indicator.py                  |  8 ++++++--
 20 files changed, 66 insertions(+), 23 deletions(-)

diff --git a/core/schemas/__init__.py b/core/schemas/__init__.py
index 905167e7f..a3e5da30f 100644
--- a/core/schemas/__init__.py
+++ b/core/schemas/__init__.py
@@ -9,6 +9,7 @@
 
 logger = logging.getLogger(__name__)
 
+
 def load_entities():
     logger.info("Registering entities")
     modules = dict()
@@ -55,7 +56,10 @@ def load_indicators():
         if indicator_file.stem not in indicator.IndicatorType.__members__:
             aenum.extend_enum(indicator.IndicatorType, indicator_file.stem, enum_value)
         modules[module_name] = enum_value
-    indicator.TYPE_MAPPING = {"indicator": indicator.Indicator, "indicators": indicator.Indicator}
+    indicator.TYPE_MAPPING = {
+        "indicator": indicator.Indicator,
+        "indicators": indicator.Indicator,
+    }
     for module_name, enum_value in modules.items():
         module = importlib.import_module(module_name)
         for _, obj in inspect.getmembers(module, inspect.isclass):
@@ -70,6 +74,7 @@ def load_indicators():
         else:
             indicator.IndicatorTypes |= cls
 
+
 def load_observables():
     logger.info("Registering observables")
     modules = dict()
@@ -88,7 +93,10 @@ def load_observables():
         modules[module_name] = observable_file.stem
     if "guess" not in observable.ObservableType.__members__:
         aenum.extend_enum(observable.ObservableType, "guess", "guess")
-    observable.TYPE_MAPPING = {"observable": observable.Observable, "observables": observable.Observable}
+    observable.TYPE_MAPPING = {
+        "observable": observable.Observable,
+        "observables": observable.Observable,
+    }
     for module_name, enum_value in modules.items():
         module = importlib.import_module(module_name)
         for _, obj in inspect.getmembers(module, inspect.isclass):
@@ -103,6 +111,7 @@ def load_observables():
         else:
             observable.ObservableTypes |= cls
 
+
 load_observables()
 load_entities()
-load_indicators()
\ No newline at end of file
+load_indicators()
diff --git a/core/schemas/entities/course_of_action.py b/core/schemas/entities/course_of_action.py
index 09e154f02..3eec594e8 100644
--- a/core/schemas/entities/course_of_action.py
+++ b/core/schemas/entities/course_of_action.py
@@ -5,4 +5,6 @@
 
 class CourseOfAction(entity.Entity):
     _type_filter: ClassVar[str] = entity.EntityType.course_of_action
-    type: Literal[entity.EntityType.course_of_action] = entity.EntityType.course_of_action
+    type: Literal[
+        entity.EntityType.course_of_action
+    ] = entity.EntityType.course_of_action
diff --git a/core/schemas/entity.py b/core/schemas/entity.py
index 62d798401..a70bf8536 100644
--- a/core/schemas/entity.py
+++ b/core/schemas/entity.py
@@ -78,4 +78,4 @@ def add_context(
         else:
             context["source"] = source
             self.context.append(context)
-        return self.save()
\ No newline at end of file
+        return self.save()
diff --git a/core/schemas/indicator.py b/core/schemas/indicator.py
index 8e0870936..ad1a54582 100644
--- a/core/schemas/indicator.py
+++ b/core/schemas/indicator.py
@@ -18,15 +18,16 @@ def future():
 
 DEFAULT_INDICATOR_VALIDITY_DAYS = 30
 
+
 # forward declarations
 class IndicatorType(str, Enum):
     ...
 
+
 IndicatorTypes = ()
 TYPE_MAPPING = {}
 
 
-
 class IndicatorMatch(BaseModel):
     name: str
     match: str
@@ -87,4 +88,4 @@ def search(cls, observables: list[str]) -> list[tuple[str, "Indicator"]]:
                 except NotImplementedError as error:
                     logging.error(
                         f"Indicator type {indicator.type} has not implemented match(): {error}"
-                    )
\ No newline at end of file
+                    )
diff --git a/core/schemas/indicators/forensicartifact.py b/core/schemas/indicators/forensicartifact.py
index c21930315..676f6c478 100644
--- a/core/schemas/indicators/forensicartifact.py
+++ b/core/schemas/indicators/forensicartifact.py
@@ -19,7 +19,9 @@ class ForensicArtifact(indicator.Indicator):
     """
 
     _type_filter: ClassVar[str] = indicator.IndicatorType.forensicartifact
-    type: Literal[indicator.IndicatorType.forensicartifact] = indicator.IndicatorType.forensicartifact
+    type: Literal[
+        indicator.IndicatorType.forensicartifact
+    ] = indicator.IndicatorType.forensicartifact
 
     sources: list[dict] = []
     aliases: list[str] = []
@@ -169,7 +171,9 @@ def save_indicators(self, create_links: bool = False):
                     regex_indicator.save()
         if create_links:
             for indicator_obj in indicators:
-                indicator_obj.link_to(self, "indicates", f"Indicates {indicator_obj.name}")
+                indicator_obj.link_to(
+                    self, "indicates", f"Indicates {indicator_obj.name}"
+                )
         return indicators
 
 
diff --git a/core/schemas/indicators/regex.py b/core/schemas/indicators/regex.py
index 67d58d945..68a74b444 100644
--- a/core/schemas/indicators/regex.py
+++ b/core/schemas/indicators/regex.py
@@ -30,4 +30,4 @@ def match(self, value: str) -> indicator.IndicatorMatch | None:
         result = self.compiled_pattern.search(value)
         if result:
             return indicator.IndicatorMatch(name=self.name, match=result.group())
-        return None
\ No newline at end of file
+        return None
diff --git a/core/schemas/observable.py b/core/schemas/observable.py
index 277b13699..acdbc1323 100644
--- a/core/schemas/observable.py
+++ b/core/schemas/observable.py
@@ -21,6 +21,7 @@
 class ObservableType(str, Enum):
     ...
 
+
 ObservableTypes = ()
 TYPE_MAPPING = {}
 
@@ -127,6 +128,7 @@ def delete_context(
                 break
         return self.save()
 
+
 def find_type(value: str) -> ObservableType | None:
     for obs_type, obj in TYPE_MAPPING.items():
         if obj.is_valid(value):
diff --git a/core/schemas/observables/certificate.py b/core/schemas/observables/certificate.py
index 7dfcf3623..eeb545b46 100644
--- a/core/schemas/observables/certificate.py
+++ b/core/schemas/observables/certificate.py
@@ -22,7 +22,9 @@ class Certificate(observable.Observable):
         fingerprint: the certificate fingerprint.
     """
 
-    type: Literal[observable.ObservableType.certificate] = observable.ObservableType.certificate
+    type: Literal[
+        observable.ObservableType.certificate
+    ] = observable.ObservableType.certificate
     last_seen: datetime.datetime = Field(default_factory=now)
     first_seen: datetime.datetime = Field(default_factory=now)
     issuer: str | None = None
diff --git a/core/schemas/observables/command_line.py b/core/schemas/observables/command_line.py
index 4a9e86dbc..375522e64 100644
--- a/core/schemas/observables/command_line.py
+++ b/core/schemas/observables/command_line.py
@@ -4,4 +4,6 @@
 
 
 class CommandLine(observable.Observable):
-    type: Literal[observable.ObservableType.command_line] = observable.ObservableType.command_line
+    type: Literal[
+        observable.ObservableType.command_line
+    ] = observable.ObservableType.command_line
diff --git a/core/schemas/observables/docker_image.py b/core/schemas/observables/docker_image.py
index ca044f1d1..dc6c78850 100644
--- a/core/schemas/observables/docker_image.py
+++ b/core/schemas/observables/docker_image.py
@@ -4,4 +4,6 @@
 
 
 class DockerImage(observable.Observable):
-    type: Literal[observable.ObservableType.docker_image] = observable.ObservableType.docker_image
+    type: Literal[
+        observable.ObservableType.docker_image
+    ] = observable.ObservableType.docker_image
diff --git a/core/schemas/observables/hostname.py b/core/schemas/observables/hostname.py
index 812e03bd9..148ec7c4e 100644
--- a/core/schemas/observables/hostname.py
+++ b/core/schemas/observables/hostname.py
@@ -6,7 +6,9 @@
 
 
 class Hostname(observable.Observable):
-    type: Literal[observable.ObservableType.hostname] = observable.ObservableType.hostname
+    type: Literal[
+        observable.ObservableType.hostname
+    ] = observable.ObservableType.hostname
 
     @staticmethod
     def is_valid(value: str) -> bool:
diff --git a/core/schemas/observables/mac_address.py b/core/schemas/observables/mac_address.py
index d900da416..628c9fa97 100644
--- a/core/schemas/observables/mac_address.py
+++ b/core/schemas/observables/mac_address.py
@@ -4,4 +4,6 @@
 
 
 class MacAddress(observable.Observable):
-    type: Literal[observable.ObservableType.mac_address] = observable.ObservableType.mac_address
+    type: Literal[
+        observable.ObservableType.mac_address
+    ] = observable.ObservableType.mac_address
diff --git a/core/schemas/observables/named_pipe.py b/core/schemas/observables/named_pipe.py
index b126fe479..0c660bb96 100644
--- a/core/schemas/observables/named_pipe.py
+++ b/core/schemas/observables/named_pipe.py
@@ -4,4 +4,6 @@
 
 
 class NamedPipe(observable.Observable):
-    type: Literal[observable.ObservableType.named_pipe] = observable.ObservableType.named_pipe
+    type: Literal[
+        observable.ObservableType.named_pipe
+    ] = observable.ObservableType.named_pipe
diff --git a/core/schemas/observables/registry_key.py b/core/schemas/observables/registry_key.py
index ccc07baca..a7ed37560 100644
--- a/core/schemas/observables/registry_key.py
+++ b/core/schemas/observables/registry_key.py
@@ -26,7 +26,9 @@ class RegistryKey(observable.Observable):
         path_file: The filesystem path to the file that contains the registry key value.
     """
 
-    type: Literal[observable.ObservableType.registry_key] = observable.ObservableType.registry_key
+    type: Literal[
+        observable.ObservableType.registry_key
+    ] = observable.ObservableType.registry_key
     key: str
     data: bytes
     hive: RegistryHive
diff --git a/core/schemas/observables/user_account.py b/core/schemas/observables/user_account.py
index e36656c0c..3e26a298c 100644
--- a/core/schemas/observables/user_account.py
+++ b/core/schemas/observables/user_account.py
@@ -14,7 +14,9 @@ class UserAccount(observable.Observable):
     Value should to be in the form :.
     """
 
-    type: Literal[observable.ObservableType.user_account] = observable.ObservableType.user_account
+    type: Literal[
+        observable.ObservableType.user_account
+    ] = observable.ObservableType.user_account
     user_id: str | None = None
     credential: str | None = None
     account_login: str | None = None
diff --git a/core/schemas/observables/user_agent.py b/core/schemas/observables/user_agent.py
index 2fcb83d36..ca80b3859 100644
--- a/core/schemas/observables/user_agent.py
+++ b/core/schemas/observables/user_agent.py
@@ -4,4 +4,6 @@
 
 
 class UserAgent(observable.Observable):
-    type: Literal[observable.ObservableType.user_agent] = observable.ObservableType.user_agent
+    type: Literal[
+        observable.ObservableType.user_agent
+    ] = observable.ObservableType.user_agent
diff --git a/core/web/apiv2/graph.py b/core/web/apiv2/graph.py
index e009eccbc..38ab2026b 100644
--- a/core/web/apiv2/graph.py
+++ b/core/web/apiv2/graph.py
@@ -111,7 +111,7 @@ class GraphSearchResponse(BaseModel):
         | entity.EntityTypes
         | indicator.IndicatorTypes
         | tag.Tag
-        | dfiq.DFIQTypes
+        | dfiq.DFIQTypes,
     ]
     paths: list[list[graph.Relationship | graph.TagRelationship]]
     total: int
diff --git a/core/web/apiv2/indicators.py b/core/web/apiv2/indicators.py
index 3c6de8397..3638401c9 100644
--- a/core/web/apiv2/indicators.py
+++ b/core/web/apiv2/indicators.py
@@ -79,7 +79,9 @@ async def patch(request: PatchIndicatorRequest, indicator_id) -> IndicatorTypes:
 
     if db_indicator.type == IndicatorType.forensicartifact:
         if db_indicator.pattern != request.indicator.pattern:
-            return forensicartifact.ForensicArtifact.from_yaml_string(request.indicator.pattern)[0]
+            return forensicartifact.ForensicArtifact.from_yaml_string(
+                request.indicator.pattern
+            )[0]
 
     update_data = request.indicator.model_dump(exclude_unset=True)
     updated_indicator = db_indicator.model_copy(update=update_data)
diff --git a/core/web/apiv2/observables.py b/core/web/apiv2/observables.py
index 956d7878b..4c1ae9de1 100644
--- a/core/web/apiv2/observables.py
+++ b/core/web/apiv2/observables.py
@@ -46,6 +46,7 @@ class PatchObservableRequest(BaseModel):
 
     observable: ObservableTypes = Field(discriminator="type")
 
+
 class NewBulkObservableAddRequest(BaseModel):
     model_config = ConfigDict(extra="forbid")
 
diff --git a/tests/schemas/indicator.py b/tests/schemas/indicator.py
index 8ac9b2429..d256116ee 100644
--- a/tests/schemas/indicator.py
+++ b/tests/schemas/indicator.py
@@ -38,7 +38,9 @@ def test_filter_entities_different_types(self) -> None:
 
         self.assertEqual(len(all_entities), 1)
         self.assertEqual(len(regex_indicators), 1)
-        self.assertEqual(regex_indicators[0].model_dump_json(), regex_indicator.model_dump_json())
+        self.assertEqual(
+            regex_indicators[0].model_dump_json(), regex_indicator.model_dump_json()
+        )
 
     def test_create_indicator_same_name_diff_types(self) -> None:
         regex1 = regex.Regex(
@@ -242,7 +244,9 @@ def test_forensic_artifacts_parent_extraction(self):
     - blah3
 """
 
-        artifacts = forensicartifact.ForensicArtifact.from_yaml_string(pattern, update_parents=True)
+        artifacts = forensicartifact.ForensicArtifact.from_yaml_string(
+            pattern, update_parents=True
+        )
         self.assertEqual(len(artifacts), 3)
 
         vertices, _, total = artifacts[0].neighbors()

From e209466f02796cc9abbf263442548840484c3c39 Mon Sep 17 00:00:00 2001
From: Fred Baguelin 
Date: Tue, 17 Sep 2024 13:11:54 +0200
Subject: [PATCH 21/39] Update analytics to import relevant entities

---
 plugins/analytics/public/macaddress_io.py | 2 +-
 plugins/analytics/public/network_whois.py | 2 +-
 plugins/analytics/public/passive_total.py | 3 ++-
 plugins/analytics/public/shodan_api.py    | 2 +-
 4 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/plugins/analytics/public/macaddress_io.py b/plugins/analytics/public/macaddress_io.py
index 81137c2c2..a3cac4780 100644
--- a/plugins/analytics/public/macaddress_io.py
+++ b/plugins/analytics/public/macaddress_io.py
@@ -9,7 +9,7 @@
 from core import taskmanager
 from core.config.config import yeti_config
 from core.schemas import task
-from core.schemas.entity import Company
+from core.schemas.entities.company import Company
 from core.schemas.observable import ObservableType
 from core.schemas.observables.mac_address import MacAddress
 
diff --git a/plugins/analytics/public/network_whois.py b/plugins/analytics/public/network_whois.py
index 265f82dc9..cd1927ed0 100644
--- a/plugins/analytics/public/network_whois.py
+++ b/plugins/analytics/public/network_whois.py
@@ -2,7 +2,7 @@
 
 from core import taskmanager
 from core.schemas import task
-from core.schemas.entity import Company
+from core.schemas.entities.company import Company
 from core.schemas.observable import ObservableType
 from core.schemas.observables import email, ipv4
 
diff --git a/plugins/analytics/public/passive_total.py b/plugins/analytics/public/passive_total.py
index 68903c7bd..243bdc6e1 100644
--- a/plugins/analytics/public/passive_total.py
+++ b/plugins/analytics/public/passive_total.py
@@ -7,7 +7,8 @@
 from core import taskmanager
 from core.config.config import yeti_config
 from core.schemas import task
-from core.schemas.entity import Company, Phone
+from core.schemas.entities.company import Company
+from core.schemas.entities.phone import Phone
 from core.schemas.observable import Observable, ObservableType
 from core.schemas.observables import email, hostname, sha256
 
diff --git a/plugins/analytics/public/shodan_api.py b/plugins/analytics/public/shodan_api.py
index 56e6ef8a9..618c496b9 100644
--- a/plugins/analytics/public/shodan_api.py
+++ b/plugins/analytics/public/shodan_api.py
@@ -5,7 +5,7 @@
 from core import taskmanager
 from core.config.config import yeti_config
 from core.schemas import task
-from core.schemas.entity import Company
+from core.schemas.entities.company import Company
 from core.schemas.observable import Observable, ObservableType
 from core.schemas.observables import asn, hostname, ipv4
 

From bf834865fc2d6c3bb952eac9648e765bbe7a47ae Mon Sep 17 00:00:00 2001
From: Fred Baguelin 
Date: Tue, 17 Sep 2024 13:55:23 +0200
Subject: [PATCH 22/39] Update analytics to correctly import schemas

---
 plugins/analytics/public/censys.py   | 5 +++--
 plugins/analytics/public/github.py   | 5 +++--
 plugins/analytics/public/malshare.py | 2 +-
 plugins/analytics/public/shodan.py   | 5 +++--
 4 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/plugins/analytics/public/censys.py b/plugins/analytics/public/censys.py
index 5c51a03a3..b1daaabb2 100644
--- a/plugins/analytics/public/censys.py
+++ b/plugins/analytics/public/censys.py
@@ -5,7 +5,8 @@
 
 from core import taskmanager
 from core.config.config import yeti_config
-from core.schemas import indicator, task
+from core.schemas import task
+from core.schemas.indicators.query import Query
 from core.schemas.observable import Observable
 
 
@@ -31,7 +32,7 @@ def run(self):
             api_secret=api_secret,
         )
 
-        censys_queries, _ = indicator.Query.filter({"query_type": "censys"})
+        censys_queries, _ = Query.filter({"query_type": "censys"})
 
         for query in censys_queries:
             ip_addresses = query_censys(hosts_api, query.pattern)
diff --git a/plugins/analytics/public/github.py b/plugins/analytics/public/github.py
index d2ee66abd..af3c21564 100644
--- a/plugins/analytics/public/github.py
+++ b/plugins/analytics/public/github.py
@@ -7,7 +7,8 @@
 
 from core import taskmanager
 from core.config.config import yeti_config
-from core.schemas import indicator, observable, task
+from core.schemas import observable, task
+from core.schemas.indicators.query import Query
 from core.schemas.observable import ObservableType
 
 
@@ -163,7 +164,7 @@ def run(self):
         auth = Auth.Token(github_token)
         self._github_api = Github(auth=auth)
 
-        github_query_indicators, _ = indicator.Query.filter({"query_type": "github"})
+        github_query_indicators, _ = Query.filter({"query_type": "github"})
         logging.info(
             f"[+] Found {len(github_query_indicators)} Github queries: {github_query_indicators}"
         )
diff --git a/plugins/analytics/public/malshare.py b/plugins/analytics/public/malshare.py
index ff81349e7..4cb540ba0 100644
--- a/plugins/analytics/public/malshare.py
+++ b/plugins/analytics/public/malshare.py
@@ -95,7 +95,7 @@ def each(self, observable: Observable):
             new_hash.add_context("malshare.com", context)
 
         if json_result["SSDEEP"]:
-            ssdeep_data = ssdeep.SsdeepHash(value=json_result["SSDEEP"]).save()
+            ssdeep_data = ssdeep.Ssdeep(value=json_result["SSDEEP"]).save()
             ssdeep_data.add_context("malshare.com", context)
             ssdeep_data.link_to(observable, "ssdeep", "malshare_query")
 
diff --git a/plugins/analytics/public/shodan.py b/plugins/analytics/public/shodan.py
index e74b9f533..c5e8c6c53 100644
--- a/plugins/analytics/public/shodan.py
+++ b/plugins/analytics/public/shodan.py
@@ -5,7 +5,8 @@
 
 from core import taskmanager
 from core.config.config import yeti_config
-from core.schemas import indicator, task
+from core.schemas import task
+from core.schemas.indicators.query import Query
 from core.schemas.observable import Observable
 
 
@@ -28,7 +29,7 @@ def run(self):
 
         shodan_api = Shodan(api_key)
 
-        shodan_queries, _ = indicator.Query.filter({"query_type": "shodan"})
+        shodan_queries, _ = Query.filter({"query_type": "shodan"})
 
         for query in shodan_queries:
             ip_addresses = query_shodan(shodan_api, query.pattern, result_limit)

From 120d299317deac205333db0f7f441dfd6569c325 Mon Sep 17 00:00:00 2001
From: Fred Baguelin 
Date: Tue, 17 Sep 2024 13:59:37 +0200
Subject: [PATCH 23/39] Update feeds to correctly import schemas

---
 plugins/feeds/public/abusech_malwarebazaar.py |  2 +-
 plugins/feeds/public/artifacts.py             |  5 ++--
 plugins/feeds/public/attack.py                | 29 ++++++++++++-------
 plugins/feeds/public/cisa_kev.py              | 17 ++++++-----
 plugins/feeds/public/et_open.py               | 18 +++++++-----
 plugins/feeds/public/lolbas.py                |  9 ++++--
 plugins/feeds/public/malpedia.py              | 22 +++++++-------
 plugins/feeds/public/otx_alienvault.py        |  3 +-
 plugins/feeds/public/timesketch.py            |  2 +-
 .../public/wiz_cloud_threat_landscape.py      |  5 +++-
 plugins/feeds/public/yaraify.py               |  9 +++---
 11 files changed, 73 insertions(+), 48 deletions(-)

diff --git a/plugins/feeds/public/abusech_malwarebazaar.py b/plugins/feeds/public/abusech_malwarebazaar.py
index 947ddb9e2..a9d8543ec 100644
--- a/plugins/feeds/public/abusech_malwarebazaar.py
+++ b/plugins/feeds/public/abusech_malwarebazaar.py
@@ -131,7 +131,7 @@ def analyze(self, block):
             imphash_data.tag(tags)
             malware_file.link_to(imphash_data, "imphash", self.name)
 
-        ssdeep_data = ssdeep.SsdeepHash(value=context["ssdeep"]).save()
+        ssdeep_data = ssdeep.Ssdeep(value=context["ssdeep"]).save()
         ssdeep_data.tag(tags)
         malware_file.link_to(ssdeep_data, "ssdeep", self.name)
 
diff --git a/plugins/feeds/public/artifacts.py b/plugins/feeds/public/artifacts.py
index c730e864c..c38f25aa1 100644
--- a/plugins/feeds/public/artifacts.py
+++ b/plugins/feeds/public/artifacts.py
@@ -9,7 +9,8 @@
 from artifacts.scripts import validator
 
 from core import taskmanager
-from core.schemas import indicator, task
+from core.schemas import task
+from core.schemas.indicators.forensicartifact import ForensicArtifact
 
 
 class ForensicArtifacts(task.FeedTask):
@@ -47,7 +48,7 @@ def run(self):
             with open(file, "r") as f:
                 yaml_string = f.read()
 
-            forensic_indicators = indicator.ForensicArtifact.from_yaml_string(
+            forensic_indicators = ForensicArtifact.from_yaml_string(
                 yaml_string, update_parents=False
             )
             for fi in forensic_indicators:
diff --git a/plugins/feeds/public/attack.py b/plugins/feeds/public/attack.py
index 56e502fdd..cc221b36d 100644
--- a/plugins/feeds/public/attack.py
+++ b/plugins/feeds/public/attack.py
@@ -7,7 +7,16 @@
 from zipfile import ZipFile
 
 from core import taskmanager
-from core.schemas import entity, task
+from core.schemas import task
+from core.schemas.entities.attack_pattern import AttackPattern
+from core.schemas.entities.campaign import Campaign
+from core.schemas.entities.course_of_action import CourseOfAction
+from core.schemas.entities.identity import Identity
+from core.schemas.entities.intrusion_set import IntrusionSet
+from core.schemas.entities.malware import Malware
+from core.schemas.entities.threat_actor import ThreatActor
+from core.schemas.entities.tool import Tool
+from core.schemas.entities.vulnerability import Vulnerability
 
 
 def _format_context_from_obj(obj):
@@ -35,7 +44,7 @@ def _format_context_from_obj(obj):
 
 
 def _process_intrusion_set(obj):
-    intrusion_set = entity.IntrusionSet(
+    intrusion_set = IntrusionSet(
         name=obj["name"],
         aliases=obj.get("aliases", []) + obj.get("x_mitre_aliases", []),
         created=obj["created"],
@@ -47,7 +56,7 @@ def _process_intrusion_set(obj):
 
 
 def _process_malware(obj):
-    malware = entity.Malware(
+    malware = Malware(
         name=obj["name"],
         created=obj["created"],
         modified=obj["modified"],
@@ -68,7 +77,7 @@ def _process_attack_pattern(obj):
     for phase in obj["kill_chain_phases"]:
         kill_chain_phases.add(f'{phase["kill_chain_name"]}:{phase["phase_name"]}')
 
-    attack_pattern = entity.AttackPattern(
+    attack_pattern = AttackPattern(
         name=obj["name"],
         created=obj["created"],
         modified=obj["modified"],
@@ -81,7 +90,7 @@ def _process_attack_pattern(obj):
 
 
 def _process_course_of_action(obj):
-    course_of_action = entity.CourseOfAction(
+    course_of_action = CourseOfAction(
         name=obj["name"],
         created=obj["created"],
         modified=obj["modified"],
@@ -92,7 +101,7 @@ def _process_course_of_action(obj):
 
 
 def _process_identity(obj):
-    identity = entity.Identity(
+    identity = Identity(
         name=obj["name"],
         created=obj["created"],
         modified=obj["modified"],
@@ -106,7 +115,7 @@ def _process_identity(obj):
 
 
 def _process_threat_actor(obj):
-    threat_actor = entity.ThreatActor(
+    threat_actor = ThreatActor(
         name=obj["name"],
         created=obj["created"],
         modified=obj["modified"],
@@ -121,7 +130,7 @@ def _process_threat_actor(obj):
 
 
 def _process_campaign(obj):
-    campaign = entity.Campaign(
+    campaign = Campaign(
         name=obj["name"],
         created=obj["created"],
         modified=obj["modified"],
@@ -135,7 +144,7 @@ def _process_campaign(obj):
 
 
 def _process_tool(obj):
-    tool = entity.Tool(
+    tool = Tool(
         name=obj["name"],
         created=obj["created"],
         modified=obj["modified"],
@@ -152,7 +161,7 @@ def _process_tool(obj):
 
 
 def _process_vulnerability(obj):
-    vulnerabilty = entity.Vulnerability(
+    vulnerabilty = Vulnerability(
         name=obj["name"],
         created=obj["created"],
         modified=obj["modified"],
diff --git a/plugins/feeds/public/cisa_kev.py b/plugins/feeds/public/cisa_kev.py
index 0db462a15..e034ffed3 100644
--- a/plugins/feeds/public/cisa_kev.py
+++ b/plugins/feeds/public/cisa_kev.py
@@ -3,7 +3,8 @@
 from typing import ClassVar
 
 from core import taskmanager
-from core.schemas import entity, task
+from core.schemas import task
+from core.schemas.entities.vulnerability import Vulnerability
 
 
 def _cves_as_dict(data):
@@ -41,12 +42,12 @@ class CisaKEV(task.FeedTask):
         "source": "https://www.cisa.gov/known-exploited-vulnerabilities-catalog",
     }
 
-    CISA_SOURCE: ClassVar["str"] = (
-        "https://www.cisa.gov/sites/default/files/feeds/known_exploited_vulnerabilities.json"
-    )
-    NVD_SOURCE: ClassVar["str"] = (
-        "https://services.nvd.nist.gov/rest/json/cves/2.0?hasKev"
-    )
+    CISA_SOURCE: ClassVar[
+        "str"
+    ] = "https://www.cisa.gov/sites/default/files/feeds/known_exploited_vulnerabilities.json"
+    NVD_SOURCE: ClassVar[
+        "str"
+    ] = "https://services.nvd.nist.gov/rest/json/cves/2.0?hasKev"
 
     def run(self):
         response = self._make_request(self.CISA_SOURCE, sort=False)
@@ -151,7 +152,7 @@ def analyze_entry(self, entry: dict, cve_details: dict):
 
         name = f"{cve_id}"
         title = entry.get("vulnerabilityName", "")
-        vulnerability = entity.Vulnerability(
+        vulnerability = Vulnerability(
             name=name,
             title=title,
             description=description,
diff --git a/plugins/feeds/public/et_open.py b/plugins/feeds/public/et_open.py
index 0c3d12f2f..6812107d5 100644
--- a/plugins/feeds/public/et_open.py
+++ b/plugins/feeds/public/et_open.py
@@ -8,6 +8,10 @@
 from core import taskmanager
 from core.config.config import yeti_config
 from core.schemas import entity, indicator, task
+from core.schemas.entities.attack_pattern import AttackPattern
+from core.schemas.entities.malware import Malware
+from core.schemas.entities.vulnerability import Vulnerability
+from core.schemas.indicators.suricata import Suricata
 
 
 class ETOpen(task.FeedTask):
@@ -37,7 +41,7 @@ def run(self):
     def analyze(self, rule_suricata: rule.Rule):
         if not self._filter_rule(rule_suricata.metadata):
             return
-        ind_suricata_rule = indicator.Suricata(
+        ind_suricata_rule = Suricata(
             name=rule_suricata["msg"],
             pattern=rule_suricata["raw"],
             metadata=rule_suricata.metadata,
@@ -61,20 +65,20 @@ def analyze(self, rule_suricata: rule.Rule):
         if tags:
             ind_suricata_rule.tag(tags)
 
-    def _extract_cve(self, meta: str) -> entity.Vulnerability:
+    def _extract_cve(self, meta: str) -> Vulnerability:
         _, cve = meta.split(" ")
         if "_" in cve:
             cve = cve.replace("_", "-")
-        ind_cve = entity.Vulnerability.find(name=cve)
+        ind_cve = Vulnerability.find(name=cve)
         if not ind_cve:
-            ind_cve = entity.Vulnerability(name=cve).save()
+            ind_cve = Vulnerability(name=cve).save()
         return ind_cve
 
     def _extract_malware_family(self, meta: str):
         _, malware_family = meta.split(" ")
-        ind_malware_family = entity.Malware.find(name=malware_family)
+        ind_malware_family = Malware.find(name=malware_family)
         if not ind_malware_family:
-            ind_malware_family = entity.Malware(name=malware_family).save()
+            ind_malware_family = Malware(name=malware_family).save()
         return ind_malware_family
 
     def _extract_tags(self, metadata: list[str]) -> list[str]:
@@ -85,7 +89,7 @@ def _extract_tags(self, metadata: list[str]) -> list[str]:
                 tags.append(tag)
         return tags
 
-    def _extract_mitre_attack(self, meta: str) -> entity.AttackPattern | None:
+    def _extract_mitre_attack(self, meta: str) -> AttackPattern | None:
         _, mitre_id = meta.split(" ")
         ind_mitre_attack, nb_ent = entity.Entity.filter(
             query_args={"type": entity.EntityType.attack_pattern},
diff --git a/plugins/feeds/public/lolbas.py b/plugins/feeds/public/lolbas.py
index 25a650e41..ab37c0b77 100644
--- a/plugins/feeds/public/lolbas.py
+++ b/plugins/feeds/public/lolbas.py
@@ -6,6 +6,9 @@
 
 from core import taskmanager
 from core.schemas import entity, indicator, task
+from core.schemas.entities.attack_pattern import AttackPattern
+from core.schemas.entities.tool import Tool
+from core.schemas.indicators.sigma import Sigma
 from core.schemas.observables import path
 
 
@@ -24,7 +27,7 @@ def run(self):
         if not response:
             return
         lolbas_json = response.json()
-        self._lolbas_attackpattern = entity.AttackPattern(name="LOLBAS usage").save()
+        self._lolbas_attackpattern = AttackPattern(name="LOLBAS usage").save()
         if not self._lolbas_attackpattern.description:
             self._lolbas_attackpattern.description = (
                 "Usage of living-off-the-land binaries and scripts"
@@ -72,7 +75,7 @@ def analyze_entry(self, entry: dict):
                         "Error processing sigma rule for %s: %s", entry["Name"], error
                     )
 
-    def process_sigma_rule(self, tool: entity.Tool, detection: dict) -> None:
+    def process_sigma_rule(self, tool: Tool, detection: dict) -> None:
         """Processes a Sigma rule as specified in the lolbas json."""
         url = detection["Sigma"]
         if not url:
@@ -92,7 +95,7 @@ def process_sigma_rule(self, tool: entity.Tool, detection: dict) -> None:
         date = sigma_data["date"]
         date = datetime.strptime(date.strip(), "%Y/%m/%d")
         # create sigma indicator
-        sigma = indicator.Sigma(
+        sigma = Sigma(
             name=title,
             description=description,
             created=date,
diff --git a/plugins/feeds/public/malpedia.py b/plugins/feeds/public/malpedia.py
index 32799938b..54c2ed6fc 100644
--- a/plugins/feeds/public/malpedia.py
+++ b/plugins/feeds/public/malpedia.py
@@ -3,7 +3,9 @@
 from typing import ClassVar
 
 from core import taskmanager
-from core.schemas import entity, task
+from core.schemas import task
+from core.schemas.entities.intrusion_set import IntrusionSet
+from core.schemas.entities.malware import Malware
 
 
 class MalpediaMalware(task.FeedTask):
@@ -14,9 +16,9 @@ class MalpediaMalware(task.FeedTask):
         "source": "https://malpedia.caad.fkie.fraunhofer.de/",
     }
 
-    _SOURCE: ClassVar["str"] = (
-        "https://malpedia.caad.fkie.fraunhofer.de/api/get/families"
-    )
+    _SOURCE: ClassVar[
+        "str"
+    ] = "https://malpedia.caad.fkie.fraunhofer.de/api/get/families"
 
     def run(self):
         response = self._make_request(self._SOURCE)
@@ -32,9 +34,9 @@ def analyze_entry(self, malware_name: str, entry: dict):
         if not entry.get("common_name"):
             return
 
-        m = entity.Malware.find(name=entry["common_name"])
+        m = Malware.find(name=entry["common_name"])
         if not m:
-            m = entity.Malware(name=entry["common_name"])
+            m = Malware(name=entry["common_name"])
 
         m.aliases = entry.get("aliases", [])
         refs = entry.get("urls", [])
@@ -48,9 +50,9 @@ def analyze_entry(self, malware_name: str, entry: dict):
         m.add_context(context["source"], context)
         attributions = entry.get("attribution", [])
         for attribution in attributions:
-            intrusion_set = entity.IntrusionSet.find(name=attribution)
+            intrusion_set = IntrusionSet.find(name=attribution)
             if not intrusion_set:
-                intrusion_set = entity.IntrusionSet(name=attribution).save()
+                intrusion_set = IntrusionSet(name=attribution).save()
             intrusion_set.link_to(m, "uses", "Malpedia")
 
         tags = []
@@ -80,9 +82,9 @@ def run(self):
             self.analyze_entry(actor_name, entry)
 
     def analyze_entry(self, actor_name: str, entry: dict):
-        intrusion_set = entity.IntrusionSet.find(name=entry["value"])
+        intrusion_set = IntrusionSet.find(name=entry["value"])
         if not intrusion_set:
-            intrusion_set = entity.IntrusionSet(name=entry["value"])
+            intrusion_set = IntrusionSet(name=entry["value"])
 
         refs = entry.get("meta", {}).get("refs", [])
         context = {
diff --git a/plugins/feeds/public/otx_alienvault.py b/plugins/feeds/public/otx_alienvault.py
index 602cca5cf..2fb8910cc 100644
--- a/plugins/feeds/public/otx_alienvault.py
+++ b/plugins/feeds/public/otx_alienvault.py
@@ -10,6 +10,7 @@
 from core import taskmanager
 from core.config.config import yeti_config
 from core.schemas import entity, indicator, observable, task
+from core.schemas.entities.investigation import Investigation
 
 
 class OTXAlienvault(task.FeedTask):
@@ -67,7 +68,7 @@ def analyze(self, item):
         context["references"] = "\r\n".join(item["references"])
         context["description"] = item["description"]
         context["link"] = "https://otx.alienvault.com/pulse/%s" % item["id"]
-        investigation = entity.Investigation(
+        investigation = Investigation(
             name=item["name"],
             description=item["description"],
             reference=f"https://otx.alienvault.com/pulse/{item['id']}",
diff --git a/plugins/feeds/public/timesketch.py b/plugins/feeds/public/timesketch.py
index 628ffe9b7..1009bb488 100644
--- a/plugins/feeds/public/timesketch.py
+++ b/plugins/feeds/public/timesketch.py
@@ -6,7 +6,7 @@
 from core import taskmanager
 from core.config.config import yeti_config
 from core.schemas import observable, task
-from core.schemas.entity import Investigation
+from core.schemas.entities.investigation import Investigation
 from core.schemas.observables import hostname, ipv4, md5, path, sha1, sha256, url
 
 TIMESKETCH_TYPE_MAPPING = {
diff --git a/plugins/feeds/public/wiz_cloud_threat_landscape.py b/plugins/feeds/public/wiz_cloud_threat_landscape.py
index b7cd0a982..62378aa85 100644
--- a/plugins/feeds/public/wiz_cloud_threat_landscape.py
+++ b/plugins/feeds/public/wiz_cloud_threat_landscape.py
@@ -7,7 +7,10 @@
 
 from core import taskmanager
 from core.schemas import task
-from core.schemas.entity import AttackPattern, Campaign, IntrusionSet, Tool
+from core.schemas.entities.attack_pattern import AttackPattern
+from core.schemas.entities.campaign import Campaign
+from core.schemas.entities.intrusion_set import IntrusionSet
+from core.schemas.entities.tool import Tool
 
 VALUE_PROPERTIES = [
     "tags",
diff --git a/plugins/feeds/public/yaraify.py b/plugins/feeds/public/yaraify.py
index 51221a340..591b0acc8 100644
--- a/plugins/feeds/public/yaraify.py
+++ b/plugins/feeds/public/yaraify.py
@@ -8,6 +8,7 @@
 
 from core import taskmanager
 from core.schemas import indicator, task
+from core.schemas.indicators.yara import Yara
 
 
 class YARAify(task.FeedTask):
@@ -18,9 +19,9 @@ class YARAify(task.FeedTask):
         "source": "",
     }
 
-    _SOURCE_ALL_RULES: ClassVar["str"] = (
-        "https://yaraify.abuse.ch/yarahub/yaraify-rules.zip"
-    )
+    _SOURCE_ALL_RULES: ClassVar[
+        "str"
+    ] = "https://yaraify.abuse.ch/yarahub/yaraify-rules.zip"
 
     def run(self):
         response = self._make_request(self._SOURCE_ALL_RULES)
@@ -40,7 +41,7 @@ def analyze_entry(self, entry: str):
             logging.error(f"Error compiling yara rule: {e}")
             return
         for r in yara_rules:
-            ind_obj = indicator.Yara(
+            ind_obj = Yara(
                 name=f"{r.identifier}",
                 pattern=entry,
                 diamond=indicator.DiamondModel.capability,

From 59260a4b3993a4748d55486ef5f7a61d96d8af40 Mon Sep 17 00:00:00 2001
From: tomchop 
Date: Wed, 18 Sep 2024 03:34:57 +0000
Subject: [PATCH 24/39] global naming change

---
 core/schemas/observables/bic.py  | 4 ++--
 core/schemas/observables/path.py | 6 +++---
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/core/schemas/observables/bic.py b/core/schemas/observables/bic.py
index 8e631c8ad..2aff5c745 100644
--- a/core/schemas/observables/bic.py
+++ b/core/schemas/observables/bic.py
@@ -3,7 +3,7 @@
 
 from core.schemas import observable
 
-bic_matcher = re.compile("^[A-Z]{6}[A-Z0-9]{2}[A-Z0-9]{3}?")
+BIC_MATCHER_REGEX = re.compile("^[A-Z]{6}[A-Z0-9]{2}[A-Z0-9]{3}?")
 
 
 class BIC(observable.Observable):
@@ -11,4 +11,4 @@ class BIC(observable.Observable):
 
     @staticmethod
     def is_valid(value: str) -> bool:
-        return bic_matcher.match(value)
+        return BIC_MATCHER_REGEX.match(value)
diff --git a/core/schemas/observables/path.py b/core/schemas/observables/path.py
index b9035f01a..dba75c59c 100644
--- a/core/schemas/observables/path.py
+++ b/core/schemas/observables/path.py
@@ -3,14 +3,14 @@
 
 from core.schemas import observable
 
-linux_path_matcher = re.compile(r"^(\/[^\/\0]+)+$")
-windows_path_matcher = re.compile(
+LINUX_PATH_REGEX = re.compile(r"^(\/[^\/\0]+)+$")
+WINDOWS_PATH_REGEX = re.compile(
     r"^(?:[a-zA-Z]\:|\\\\[\w\.]+\\[\w.$]+)\\(?:[\w]+\\)*\w([\w.])+"
 )
 
 
 def path_validator(value):
-    return linux_path_matcher.match(value) or windows_path_matcher.match(value)
+    return LINUX_PATH_REGEX.match(value) or WINDOWS_PATH_REGEX.match(value)
 
 
 class Path(observable.Observable):

From d99a42e4d2803f982c1f1e9e12ae7be3c40257e7 Mon Sep 17 00:00:00 2001
From: tomchop 
Date: Wed, 18 Sep 2024 03:35:03 +0000
Subject: [PATCH 25/39] Formatting

---
 core/schemas/entities/private/README.md | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/core/schemas/entities/private/README.md b/core/schemas/entities/private/README.md
index c34e233c4..0f25cdf78 100644
--- a/core/schemas/entities/private/README.md
+++ b/core/schemas/entities/private/README.md
@@ -1,2 +1,5 @@
 ### Private indicators
-This directory is where you should place your private indicators. It could be named anything else, but this one has a `.gitignore` so you don't mess things up. ;-)
+
+This directory is where you should place your private indicators. It could be
+named anything else, but this one has a `.gitignore` so you don't mess things
+up. ;-)

From bc8c4018e0c59b1c78fe24fdddb5bafa7493c955 Mon Sep 17 00:00:00 2001
From: tomchop 
Date: Wed, 18 Sep 2024 03:35:31 +0000
Subject: [PATCH 26/39] Ruff formatting

---
 core/common/utils.py                        |  6 +++---
 core/schemas/entities/course_of_action.py   |  4 ++--
 core/schemas/entity.py                      |  3 +--
 core/schemas/indicator.py                   |  3 +--
 core/schemas/indicators/forensicartifact.py |  4 ++--
 core/schemas/observable.py                  |  3 +--
 core/schemas/observables/certificate.py     |  4 ++--
 core/schemas/observables/command_line.py    |  4 ++--
 core/schemas/observables/docker_image.py    |  4 ++--
 core/schemas/observables/hostname.py        |  4 ++--
 core/schemas/observables/mac_address.py     |  4 ++--
 core/schemas/observables/named_pipe.py      |  4 ++--
 core/schemas/observables/registry_key.py    |  4 ++--
 core/schemas/observables/user_account.py    |  4 ++--
 core/schemas/observables/user_agent.py      |  4 ++--
 plugins/feeds/public/cisa_kev.py            | 12 ++++++------
 plugins/feeds/public/malpedia.py            |  6 +++---
 plugins/feeds/public/yaraify.py             |  6 +++---
 tests/apiv2/templates.py                    | 12 ++++++------
 19 files changed, 46 insertions(+), 49 deletions(-)

diff --git a/core/common/utils.py b/core/common/utils.py
index 28d81c447..a0010274c 100644
--- a/core/common/utils.py
+++ b/core/common/utils.py
@@ -14,9 +14,9 @@
 
 if hasattr(yeti_config, "tldextract"):
     if yeti_config.tldextract.extra_suffixes:
-        tld_extract_dict[
-            "extra_suffixes"
-        ] = yeti_config.tldextract.extra_suffixes.split(",")
+        tld_extract_dict["extra_suffixes"] = (
+            yeti_config.tldextract.extra_suffixes.split(",")
+        )
     if yeti_config.tldextract.suffix_list_urls:
         tld_extract_dict["suffix_list_urls"] = yeti_config.tldextract.suffix_list_urls
 
diff --git a/core/schemas/entities/course_of_action.py b/core/schemas/entities/course_of_action.py
index 3eec594e8..8a5d2722e 100644
--- a/core/schemas/entities/course_of_action.py
+++ b/core/schemas/entities/course_of_action.py
@@ -5,6 +5,6 @@
 
 class CourseOfAction(entity.Entity):
     _type_filter: ClassVar[str] = entity.EntityType.course_of_action
-    type: Literal[
+    type: Literal[entity.EntityType.course_of_action] = (
         entity.EntityType.course_of_action
-    ] = entity.EntityType.course_of_action
+    )
diff --git a/core/schemas/entity.py b/core/schemas/entity.py
index a70bf8536..c4add2ecf 100644
--- a/core/schemas/entity.py
+++ b/core/schemas/entity.py
@@ -10,8 +10,7 @@
 
 
 # forward declarations
-class EntityType(str, Enum):
-    ...
+class EntityType(str, Enum): ...
 
 
 EntityTypes = ()
diff --git a/core/schemas/indicator.py b/core/schemas/indicator.py
index ad1a54582..c7434f74b 100644
--- a/core/schemas/indicator.py
+++ b/core/schemas/indicator.py
@@ -20,8 +20,7 @@ def future():
 
 
 # forward declarations
-class IndicatorType(str, Enum):
-    ...
+class IndicatorType(str, Enum): ...
 
 
 IndicatorTypes = ()
diff --git a/core/schemas/indicators/forensicartifact.py b/core/schemas/indicators/forensicartifact.py
index 676f6c478..793f16df1 100644
--- a/core/schemas/indicators/forensicartifact.py
+++ b/core/schemas/indicators/forensicartifact.py
@@ -19,9 +19,9 @@ class ForensicArtifact(indicator.Indicator):
     """
 
     _type_filter: ClassVar[str] = indicator.IndicatorType.forensicartifact
-    type: Literal[
+    type: Literal[indicator.IndicatorType.forensicartifact] = (
         indicator.IndicatorType.forensicartifact
-    ] = indicator.IndicatorType.forensicartifact
+    )
 
     sources: list[dict] = []
     aliases: list[str] = []
diff --git a/core/schemas/observable.py b/core/schemas/observable.py
index acdbc1323..7f97b4e57 100644
--- a/core/schemas/observable.py
+++ b/core/schemas/observable.py
@@ -18,8 +18,7 @@
 
 
 # forward declarations
-class ObservableType(str, Enum):
-    ...
+class ObservableType(str, Enum): ...
 
 
 ObservableTypes = ()
diff --git a/core/schemas/observables/certificate.py b/core/schemas/observables/certificate.py
index eeb545b46..4356f60f3 100644
--- a/core/schemas/observables/certificate.py
+++ b/core/schemas/observables/certificate.py
@@ -22,9 +22,9 @@ class Certificate(observable.Observable):
         fingerprint: the certificate fingerprint.
     """
 
-    type: Literal[
+    type: Literal[observable.ObservableType.certificate] = (
         observable.ObservableType.certificate
-    ] = observable.ObservableType.certificate
+    )
     last_seen: datetime.datetime = Field(default_factory=now)
     first_seen: datetime.datetime = Field(default_factory=now)
     issuer: str | None = None
diff --git a/core/schemas/observables/command_line.py b/core/schemas/observables/command_line.py
index 375522e64..2273b9cd4 100644
--- a/core/schemas/observables/command_line.py
+++ b/core/schemas/observables/command_line.py
@@ -4,6 +4,6 @@
 
 
 class CommandLine(observable.Observable):
-    type: Literal[
+    type: Literal[observable.ObservableType.command_line] = (
         observable.ObservableType.command_line
-    ] = observable.ObservableType.command_line
+    )
diff --git a/core/schemas/observables/docker_image.py b/core/schemas/observables/docker_image.py
index dc6c78850..a117a00c5 100644
--- a/core/schemas/observables/docker_image.py
+++ b/core/schemas/observables/docker_image.py
@@ -4,6 +4,6 @@
 
 
 class DockerImage(observable.Observable):
-    type: Literal[
+    type: Literal[observable.ObservableType.docker_image] = (
         observable.ObservableType.docker_image
-    ] = observable.ObservableType.docker_image
+    )
diff --git a/core/schemas/observables/hostname.py b/core/schemas/observables/hostname.py
index 148ec7c4e..f3d0154d2 100644
--- a/core/schemas/observables/hostname.py
+++ b/core/schemas/observables/hostname.py
@@ -6,9 +6,9 @@
 
 
 class Hostname(observable.Observable):
-    type: Literal[
+    type: Literal[observable.ObservableType.hostname] = (
         observable.ObservableType.hostname
-    ] = observable.ObservableType.hostname
+    )
 
     @staticmethod
     def is_valid(value: str) -> bool:
diff --git a/core/schemas/observables/mac_address.py b/core/schemas/observables/mac_address.py
index 628c9fa97..3b5c0f578 100644
--- a/core/schemas/observables/mac_address.py
+++ b/core/schemas/observables/mac_address.py
@@ -4,6 +4,6 @@
 
 
 class MacAddress(observable.Observable):
-    type: Literal[
+    type: Literal[observable.ObservableType.mac_address] = (
         observable.ObservableType.mac_address
-    ] = observable.ObservableType.mac_address
+    )
diff --git a/core/schemas/observables/named_pipe.py b/core/schemas/observables/named_pipe.py
index 0c660bb96..8bc49e971 100644
--- a/core/schemas/observables/named_pipe.py
+++ b/core/schemas/observables/named_pipe.py
@@ -4,6 +4,6 @@
 
 
 class NamedPipe(observable.Observable):
-    type: Literal[
+    type: Literal[observable.ObservableType.named_pipe] = (
         observable.ObservableType.named_pipe
-    ] = observable.ObservableType.named_pipe
+    )
diff --git a/core/schemas/observables/registry_key.py b/core/schemas/observables/registry_key.py
index a7ed37560..07d9f2db4 100644
--- a/core/schemas/observables/registry_key.py
+++ b/core/schemas/observables/registry_key.py
@@ -26,9 +26,9 @@ class RegistryKey(observable.Observable):
         path_file: The filesystem path to the file that contains the registry key value.
     """
 
-    type: Literal[
+    type: Literal[observable.ObservableType.registry_key] = (
         observable.ObservableType.registry_key
-    ] = observable.ObservableType.registry_key
+    )
     key: str
     data: bytes
     hive: RegistryHive
diff --git a/core/schemas/observables/user_account.py b/core/schemas/observables/user_account.py
index 3e26a298c..76e8490d1 100644
--- a/core/schemas/observables/user_account.py
+++ b/core/schemas/observables/user_account.py
@@ -14,9 +14,9 @@ class UserAccount(observable.Observable):
     Value should to be in the form :.
     """
 
-    type: Literal[
+    type: Literal[observable.ObservableType.user_account] = (
         observable.ObservableType.user_account
-    ] = observable.ObservableType.user_account
+    )
     user_id: str | None = None
     credential: str | None = None
     account_login: str | None = None
diff --git a/core/schemas/observables/user_agent.py b/core/schemas/observables/user_agent.py
index ca80b3859..f6327c7e1 100644
--- a/core/schemas/observables/user_agent.py
+++ b/core/schemas/observables/user_agent.py
@@ -4,6 +4,6 @@
 
 
 class UserAgent(observable.Observable):
-    type: Literal[
+    type: Literal[observable.ObservableType.user_agent] = (
         observable.ObservableType.user_agent
-    ] = observable.ObservableType.user_agent
+    )
diff --git a/plugins/feeds/public/cisa_kev.py b/plugins/feeds/public/cisa_kev.py
index e034ffed3..ccb188e53 100644
--- a/plugins/feeds/public/cisa_kev.py
+++ b/plugins/feeds/public/cisa_kev.py
@@ -42,12 +42,12 @@ class CisaKEV(task.FeedTask):
         "source": "https://www.cisa.gov/known-exploited-vulnerabilities-catalog",
     }
 
-    CISA_SOURCE: ClassVar[
-        "str"
-    ] = "https://www.cisa.gov/sites/default/files/feeds/known_exploited_vulnerabilities.json"
-    NVD_SOURCE: ClassVar[
-        "str"
-    ] = "https://services.nvd.nist.gov/rest/json/cves/2.0?hasKev"
+    CISA_SOURCE: ClassVar["str"] = (
+        "https://www.cisa.gov/sites/default/files/feeds/known_exploited_vulnerabilities.json"
+    )
+    NVD_SOURCE: ClassVar["str"] = (
+        "https://services.nvd.nist.gov/rest/json/cves/2.0?hasKev"
+    )
 
     def run(self):
         response = self._make_request(self.CISA_SOURCE, sort=False)
diff --git a/plugins/feeds/public/malpedia.py b/plugins/feeds/public/malpedia.py
index 54c2ed6fc..50c6af6e1 100644
--- a/plugins/feeds/public/malpedia.py
+++ b/plugins/feeds/public/malpedia.py
@@ -16,9 +16,9 @@ class MalpediaMalware(task.FeedTask):
         "source": "https://malpedia.caad.fkie.fraunhofer.de/",
     }
 
-    _SOURCE: ClassVar[
-        "str"
-    ] = "https://malpedia.caad.fkie.fraunhofer.de/api/get/families"
+    _SOURCE: ClassVar["str"] = (
+        "https://malpedia.caad.fkie.fraunhofer.de/api/get/families"
+    )
 
     def run(self):
         response = self._make_request(self._SOURCE)
diff --git a/plugins/feeds/public/yaraify.py b/plugins/feeds/public/yaraify.py
index 591b0acc8..748530de0 100644
--- a/plugins/feeds/public/yaraify.py
+++ b/plugins/feeds/public/yaraify.py
@@ -19,9 +19,9 @@ class YARAify(task.FeedTask):
         "source": "",
     }
 
-    _SOURCE_ALL_RULES: ClassVar[
-        "str"
-    ] = "https://yaraify.abuse.ch/yarahub/yaraify-rules.zip"
+    _SOURCE_ALL_RULES: ClassVar["str"] = (
+        "https://yaraify.abuse.ch/yarahub/yaraify-rules.zip"
+    )
 
     def run(self):
         response = self._make_request(self._SOURCE_ALL_RULES)
diff --git a/tests/apiv2/templates.py b/tests/apiv2/templates.py
index e1db94b4c..d5e9e2c4a 100644
--- a/tests/apiv2/templates.py
+++ b/tests/apiv2/templates.py
@@ -90,9 +90,9 @@ def test_render_template_by_id(self):
             },
         )
         data = response.text
-        response.headers[
-            "Content-Disposition"
-        ] = "attachment; filename=FakeTemplate.txt"
+        response.headers["Content-Disposition"] = (
+            "attachment; filename=FakeTemplate.txt"
+        )
         self.assertEqual(response.status_code, 200, data)
         self.assertEqual(data, "\n1.1.1.1\n2.2.2.2\n3.3.3.3\n\n\n")
 
@@ -106,8 +106,8 @@ def test_render_template_by_search(self):
             json={"template_id": self.template.id, "search_query": "yeti"},
         )
         data = response.text
-        response.headers[
-            "Content-Disposition"
-        ] = "attachment; filename=FakeTemplate.txt"
+        response.headers["Content-Disposition"] = (
+            "attachment; filename=FakeTemplate.txt"
+        )
         self.assertEqual(response.status_code, 200, data)
         self.assertEqual(data, "\nyeti1.com\nyeti2.com\nyeti3.com\n\n\n")

From 580acfdfc17ba58b8382a2b5d6817f0a0ae932c8 Mon Sep 17 00:00:00 2001
From: Fred Baguelin 
Date: Wed, 18 Sep 2024 12:11:06 +0200
Subject: [PATCH 27/39] Register entities, observables and indicators object
 accordingly to schemas

---
 core/schemas/__init__.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/core/schemas/__init__.py b/core/schemas/__init__.py
index a3e5da30f..daa14abb5 100644
--- a/core/schemas/__init__.py
+++ b/core/schemas/__init__.py
@@ -31,6 +31,7 @@ def load_entities():
         for _, obj in inspect.getmembers(module, inspect.isclass):
             if issubclass(obj, entity.Entity):
                 entity.TYPE_MAPPING[enum_value] = obj
+                setattr(entity, obj.__name__, obj)
     for key in entity.TYPE_MAPPING:
         if key in ["entity", "entities"]:
             continue
@@ -65,6 +66,7 @@ def load_indicators():
         for _, obj in inspect.getmembers(module, inspect.isclass):
             if issubclass(obj, indicator.Indicator):
                 indicator.TYPE_MAPPING[enum_value] = obj
+                setattr(indicator, obj.__name__, obj)
     for key in indicator.TYPE_MAPPING:
         if key in ["indicator", "indicators"]:
             continue
@@ -102,6 +104,7 @@ def load_observables():
         for _, obj in inspect.getmembers(module, inspect.isclass):
             if issubclass(obj, observable.Observable):
                 observable.TYPE_MAPPING[enum_value] = obj
+                setattr(observable, obj.__name__, obj)
     for key in observable.TYPE_MAPPING:
         if key in ["observable", "observables"]:
             continue

From 8afe49964ba5c8aaffcb7321e61ef5cea4dd3dfb Mon Sep 17 00:00:00 2001
From: Fred Baguelin 
Date: Wed, 18 Sep 2024 12:18:25 +0200
Subject: [PATCH 28/39] Revert tests to main version

---
 tests/apiv2/entities.py    |  9 +++---
 tests/apiv2/graph.py       | 25 +++++++--------
 tests/apiv2/indicators.py  |  7 ++---
 tests/schemas/entity.py    | 64 ++++++++++++++++++--------------------
 tests/schemas/fixture.py   | 23 +++++++-------
 tests/schemas/graph.py     |  9 +++---
 tests/schemas/indicator.py | 47 ++++++++++++++--------------
 7 files changed, 89 insertions(+), 95 deletions(-)

diff --git a/tests/apiv2/entities.py b/tests/apiv2/entities.py
index 458879e15..541550a20 100644
--- a/tests/apiv2/entities.py
+++ b/tests/apiv2/entities.py
@@ -7,7 +7,6 @@
 
 from core import database_arango
 from core.schemas import entity
-from core.schemas.entities import attack_pattern, malware, threat_actor
 from core.schemas.user import UserSensitive
 from core.web import webapp
 
@@ -24,11 +23,11 @@ def setUp(self) -> None:
             "/api/v2/auth/api-token", headers={"x-yeti-apikey": user.api_key}
         ).json()
         client.headers = {"Authorization": "Bearer " + token_data["access_token"]}
-        self.entity1 = threat_actor.ThreatActor(
+        self.entity1 = entity.ThreatActor(
             name="ta1", aliases=["badactor"], created=datetime.datetime(2020, 1, 1)
         ).save()
         self.entity1.tag(["ta1"])
-        self.entity2 = threat_actor.ThreatActor(name="bears").save()
+        self.entity2 = entity.ThreatActor(name="bears").save()
 
     def tearDown(self) -> None:
         database_arango.db.clear()
@@ -92,8 +91,8 @@ def test_search_entities_subfields(self):
         self.assertEqual(data["entities"][0]["type"], "threat-actor")
 
     def test_search_entities_multiple_types(self):
-        attack_pattern.AttackPattern(name="ttp1").save()
-        malware.Malware(name="malware1").save()  # this won't show up
+        entity.AttackPattern(name="ttp1").save()
+        entity.Malware(name="malware1").save()  # this won't show up
         response = client.post(
             "/api/v2/entities/search",
             json={"query": {"type__in": ["threat-actor", "attack-pattern"]}},
diff --git a/tests/apiv2/graph.py b/tests/apiv2/graph.py
index f8f561c49..53ee43875 100644
--- a/tests/apiv2/graph.py
+++ b/tests/apiv2/graph.py
@@ -5,10 +5,9 @@
 from fastapi.testclient import TestClient
 
 from core import database_arango
-from core.schemas.entities import attack_pattern, malware, threat_actor
+from core.schemas.entity import AttackPattern, Malware, ThreatActor
 from core.schemas.graph import Relationship
-from core.schemas.indicator import DiamondModel
-from core.schemas.indicators import forensicartifact, query, regex
+from core.schemas.indicator import DiamondModel, ForensicArtifact, Query, Regex
 from core.schemas.observables import hostname, ipv4, url
 from core.schemas.user import UserSensitive
 from core.web import webapp
@@ -28,8 +27,8 @@ def setUp(self) -> None:
         client.headers = {"Authorization": "Bearer " + token_data["access_token"]}
         self.observable1 = hostname.Hostname(value="tomchop.me").save()
         self.observable2 = ipv4.IPv4(value="127.0.0.1").save()
-        self.entity1 = threat_actor.ThreatActor(name="actor0").save()
-        self.indicator1 = query.Query(
+        self.entity1 = ThreatActor(name="actor0").save()
+        self.indicator1 = Query(
             name="query1",
             query_type="opensearch",
             target_systems=["system1"],
@@ -284,8 +283,8 @@ def test_neighbors_filter(self):
 
     def test_neighbor_filter_in(self):
         self.entity1.link_to(self.observable1, "uses", "asd")
-        malware_entity = malware.Malware(name="malware1", aliases=["blah"]).save()
-        self.entity1.link_to(malware_entity, "uses", "asd")
+        malware = Malware(name="malware1", aliases=["blah"]).save()
+        self.entity1.link_to(malware, "uses", "asd")
         response = client.post(
             "/api/v2/graph/search",
             json={
@@ -300,9 +299,7 @@ def test_neighbor_filter_in(self):
         data = response.json()
         self.assertEqual(response.status_code, 200, data)
         self.assertEqual(len(data["vertices"]), 1)
-        self.assertEqual(
-            data["vertices"][malware_entity.extended_id]["name"], "malware1"
-        )
+        self.assertEqual(data["vertices"][malware.extended_id]["name"], "malware1")
 
     def test_add_link(self):
         response = client.post(
@@ -378,8 +375,8 @@ def setUp(self) -> None:
         self.observable1.tag(["tag1", "tag2"])
         self.observable2 = hostname.Hostname(value="test2.com").save()
         self.observable3 = url.Url(value="http://test1.com/admin").save()
-        self.entity1 = threat_actor.ThreatActor(name="tester").save()
-        self.indicator1 = regex.Regex(
+        self.entity1 = ThreatActor(name="tester").save()
+        self.indicator1 = Regex(
             name="test c2",
             pattern="test[0-9].com",
             location="network",
@@ -579,9 +576,9 @@ def setUp(self) -> None:
         ).json()
         client.headers = {"Authorization": "Bearer " + token_data["access_token"]}
 
-        self.persistence = attack_pattern.AttackPattern(name="persistence").save()
+        self.persistence = AttackPattern(name="persistence").save()
         self.persistence.tag(["triage"])
-        self.persistence_artifact = forensicartifact.ForensicArtifact.from_yaml_string(
+        self.persistence_artifact = ForensicArtifact.from_yaml_string(
             """doc: Crontab files.
 name: LinuxCronTabs
 sources:
diff --git a/tests/apiv2/indicators.py b/tests/apiv2/indicators.py
index 5ddfc8d58..a021d4d1b 100644
--- a/tests/apiv2/indicators.py
+++ b/tests/apiv2/indicators.py
@@ -7,7 +7,6 @@
 
 from core import database_arango
 from core.schemas import indicator
-from core.schemas.indicators import query, regex
 from core.schemas.user import UserSensitive
 from core.web import webapp
 
@@ -24,14 +23,14 @@ def setUp(self) -> None:
             "/api/v2/auth/api-token", headers={"x-yeti-apikey": user.api_key}
         ).json()
         client.headers = {"Authorization": "Bearer " + token_data["access_token"]}
-        self.indicator1 = regex.Regex(
+        self.indicator1 = indicator.Regex(
             name="hex",
             pattern="[0-9a-f]",
             location="filesystem",
             diamond=indicator.DiamondModel.capability,
         ).save()
         self.indicator1.tag(["hextag"])
-        self.indicator2 = regex.Regex(
+        self.indicator2 = indicator.Regex(
             name="localhost",
             pattern="127.0.0.1",
             location="network",
@@ -102,7 +101,7 @@ def test_search_indicators_subfields(self):
         self.assertEqual(data["indicators"][0]["type"], "regex")
 
     def test_search_indicators_by_alias(self):
-        query.Query(
+        indicator.Query(
             name="query1",
             pattern="SELECT * FROM table",
             location="database",
diff --git a/tests/schemas/entity.py b/tests/schemas/entity.py
index fabbd8181..ea33cf99c 100644
--- a/tests/schemas/entity.py
+++ b/tests/schemas/entity.py
@@ -3,14 +3,14 @@
 
 from core import database_arango
 from core.schemas import tag
-from core.schemas.entities import (
-    attack_pattern,
-    malware,
-    threat_actor,
-    tool,
-    vulnerability,
+from core.schemas.entity import (
+    AttackPattern,
+    Entity,
+    Malware,
+    ThreatActor,
+    Tool,
+    Vulnerability,
 )
-from core.schemas.entity import Entity
 from core.schemas.observables import hostname
 
 
@@ -18,20 +18,18 @@ class EntityTest(unittest.TestCase):
     def setUp(self) -> None:
         database_arango.db.connect(database="yeti_test")
         database_arango.db.clear()
-        self.ta1 = threat_actor.ThreatActor(name="APT123", aliases=["CrazyFrog"]).save()
-        self.vuln1 = vulnerability.Vulnerability(
-            name="CVE-2018-1337", title="elite exploit"
-        ).save()
-        self.malware1 = malware.Malware(
+        self.ta1 = ThreatActor(name="APT123", aliases=["CrazyFrog"]).save()
+        self.vuln1 = Vulnerability(name="CVE-2018-1337", title="elite exploit").save()
+        self.malware1 = Malware(
             name="zeus", created=datetime.datetime(2020, 1, 1)
         ).save()
-        self.tool1 = tool.Tool(name="mimikatz").save()
+        self.tool1 = Tool(name="mimikatz").save()
 
     def tearDown(self) -> None:
         database_arango.db.clear()
 
     def test_create_entity(self) -> None:
-        result = threat_actor.ThreatActor(name="APT0").save()
+        result = ThreatActor(name="APT0").save()
         self.assertIsNotNone(result.id)
         self.assertIsNotNone(result.created)
         self.assertEqual(result.name, "APT0")
@@ -43,11 +41,11 @@ def test_entity_get_correct_type(self) -> None:
         result = Entity.get(self.ta1.id)
         assert result is not None
         self.assertIsNotNone(result)
-        self.assertIsInstance(result, threat_actor.ThreatActor)
+        self.assertIsInstance(result, ThreatActor)
         self.assertEqual(result.type, "threat-actor")
 
     def test_attack_pattern(self) -> None:
-        result = attack_pattern.AttackPattern(
+        result = AttackPattern(
             name="Abuse Elevation Control Mechanism",
             aliases=["T1548"],
             kill_chain_phases=["mitre-attack:Privilege Escalation"],
@@ -59,19 +57,19 @@ def test_attack_pattern(self) -> None:
         self.assertIn("T1548", result.aliases)
 
     def test_entity_dupe_name_type(self) -> None:
-        oldm = malware.Malware(name="APT123").save()
-        ta = threat_actor.ThreatActor.find(name="APT123")
-        m = malware.Malware.find(name="APT123")
+        oldm = Malware(name="APT123").save()
+        ta = ThreatActor.find(name="APT123")
+        m = Malware.find(name="APT123")
         self.assertEqual(ta.id, self.ta1.id)
         self.assertEqual(m.id, oldm.id)
-        self.assertIsInstance(m, malware.Malware)
-        self.assertIsInstance(ta, threat_actor.ThreatActor)
+        self.assertIsInstance(m, Malware)
+        self.assertIsInstance(ta, ThreatActor)
 
     def test_list_entities(self) -> None:
         all_entities = list(Entity.list())
-        threat_actor_entities = list(threat_actor.ThreatActor.list())
-        tool_entities = list(tool.Tool.list())
-        malware_entities = list(malware.Malware.list())
+        threat_actor_entities = list(ThreatActor.list())
+        tool_entities = list(Tool.list())
+        malware_entities = list(Malware.list())
 
         self.assertEqual(len(all_entities), 4)
 
@@ -124,7 +122,7 @@ def test_filter_entities_time(self):
         self.assertNotIn(self.malware1, entities)
 
     def test_entity_with_tags(self):
-        entity = threat_actor.ThreatActor(name="APT0").save()
+        entity = ThreatActor(name="APT0").save()
         entity.tag(["tag1", "tag2"])
         observable = hostname.Hostname(value="doman.com").save()
 
@@ -147,13 +145,13 @@ def test_entity_with_tags(self):
 
     def test_duplicate_name(self):
         """Tests that saving an entity with an existing name will return the existing entity."""
-        ta = threat_actor.ThreatActor(name="APT123").save()
+        ta = ThreatActor(name="APT123").save()
         self.assertEqual(ta.id, self.ta1.id)
 
     def test_entity_duplicate_name(self):
         """Tests that two entities of different types can have the same name."""
-        psexec_tool = tool.Tool(name="psexec").save()
-        psexec_ap = attack_pattern.AttackPattern(name="psexec").save()
+        psexec_tool = Tool(name="psexec").save()
+        psexec_ap = AttackPattern(name="psexec").save()
         self.assertNotEqual(psexec_tool.id, psexec_ap.id)
         self.assertEqual(psexec_tool.type, "tool")
         self.assertEqual(psexec_ap.type, "attack-pattern")
@@ -161,12 +159,12 @@ def test_entity_duplicate_name(self):
     def test_no_empty_name(self):
         """Tests that an entity with an empty name cannot be saved."""
         with self.assertRaises(ValueError):
-            threat_actor.ThreatActor(name="").save()
+            ThreatActor(name="").save()
 
     def test_bad_cve_name(self):
-        vuln = vulnerability.Vulnerability(name="1337-4242").save()
-        self.assertEqual(vulnerability.Vulnerability.is_valid(vuln), False)
+        vulnerability = Vulnerability(name="1337-4242").save()
+        self.assertEqual(Vulnerability.is_valid(vulnerability), False)
 
     def test_correct_cve_name(self):
-        vuln = vulnerability.Vulnerability(name="CVE-1337-4242").save()
-        self.assertEqual(vulnerability.Vulnerability.is_valid(vuln), True)
+        vulnerability = Vulnerability(name="CVE-1337-4242").save()
+        self.assertEqual(Vulnerability.is_valid(vulnerability), True)
diff --git a/tests/schemas/fixture.py b/tests/schemas/fixture.py
index b6e068061..774a16879 100644
--- a/tests/schemas/fixture.py
+++ b/tests/schemas/fixture.py
@@ -1,9 +1,8 @@
 import unittest
 
 from core import database_arango
-from core.schemas.entities import investigation, malware, threat_actor
-from core.schemas.indicator import DiamondModel
-from core.schemas.indicators import query, regex
+from core.schemas.entity import Investigation, Malware, ThreatActor
+from core.schemas.indicator import DiamondModel, Query, Regex
 from core.schemas.observables import (
     bic,
     generic,
@@ -41,7 +40,9 @@ def test_something(self):
         hacker = hostname.Hostname(value="hacker.com").save()
         sus_hacker = hostname.Hostname(value="sus.hacker.com").save()
         mac_address.MacAddress(value="00:11:22:33:44:55").save()
-        generic_obs = generic.Generic(value="SomeInterestingString").save()
+        generic_obs = generic.Generic(
+            value="SomeInterestingString"
+        ).save()
         generic_obs.add_context("test_source", {"test": "test"})
 
         hacker.link_to(www_hacker, "domain", "Domain")
@@ -58,22 +59,22 @@ def test_something(self):
         ibantest.link_to(bictest, "bic", "BIC")
         ibantest.tag(["example"])
 
-        ta = threat_actor.ThreatActor(name="HackerActor").save()
+        ta = ThreatActor(name="HackerActor").save()
         ta.tag(["Hack!ré T@ëst"])
         ta.link_to(hacker, "uses", "Uses domain")
 
-        regex_indicator = regex.Regex(
+        regex = Regex(
             name="hex",
             pattern="/tmp/[0-9a-f]",
             location="bodyfile",
             diamond=DiamondModel.capability,
         ).save()
-        regex_indicator.link_to(hacker, "indicates", "Domain dropped by this regex")
-        xmrig = malware.Malware(name="xmrig").save()
+        regex.link_to(hacker, "indicates", "Domain dropped by this regex")
+        xmrig = Malware(name="xmrig").save()
         xmrig.tag(["xmrig"])
-        regex_indicator.link_to(xmrig, "indicates", "Usual name for dropped binary")
+        regex.link_to(xmrig, "indicates", "Usual name for dropped binary")
 
-        query.Query(
+        Query(
             name="ssh succesful logins",
             location="syslogs",
             diamond=DiamondModel.capability,
@@ -82,7 +83,7 @@ def test_something(self):
             target_systems=["timesketch", "plaso"],
             relevant_tags=["ssh", "login"],
         ).save()
-        i = investigation.Investigation(
+        i = Investigation(
             name="coin mining case",
             reference="http://timesketch-server/sketch/12345",
             relevant_tags=["coin", "mining"],
diff --git a/tests/schemas/graph.py b/tests/schemas/graph.py
index 6682925a9..e55b50211 100644
--- a/tests/schemas/graph.py
+++ b/tests/schemas/graph.py
@@ -3,8 +3,7 @@
 from fastapi.testclient import TestClient
 
 from core import database_arango
-from core.schemas.entities import campaign, malware
-from core.schemas.entity import Entity
+from core.schemas.entity import Campaign, Entity, Malware
 from core.schemas.graph import GraphFilter, Relationship
 from core.schemas.observables import hostname, ipv4, user_agent
 from core.web import webapp
@@ -20,9 +19,9 @@ def setUp(self) -> None:
         self.observable2 = ipv4.IPv4(value="127.0.0.1").save()
         self.observable3 = ipv4.IPv4(value="8.8.8.8").save()
         self.observable4 = user_agent.UserAgent(value="Mozilla/5.0").save()
-        self.entity1 = malware.Malware(name="plugx").save()
-        self.entity2 = campaign.Campaign(name="campaign1").save()
-        self.entity3 = campaign.Campaign(name="campaign2").save()
+        self.entity1 = Malware(name="plugx").save()
+        self.entity2 = Campaign(name="campaign1").save()
+        self.entity3 = Campaign(name="campaign2").save()
 
     def tearDown(self) -> None:
         database_arango.db.clear()
diff --git a/tests/schemas/indicator.py b/tests/schemas/indicator.py
index d256116ee..ee1505ee0 100644
--- a/tests/schemas/indicator.py
+++ b/tests/schemas/indicator.py
@@ -1,8 +1,13 @@
 import unittest
 
 from core import database_arango
-from core.schemas.indicator import DiamondModel, Indicator
-from core.schemas.indicators import forensicartifact, query, regex
+from core.schemas.indicator import (
+    DiamondModel,
+    ForensicArtifact,
+    Indicator,
+    Query,
+    Regex,
+)
 
 
 class IndicatorTest(unittest.TestCase):
@@ -14,7 +19,7 @@ def tearDown(self) -> None:
         database_arango.db.clear()
 
     def test_create_indicator(self) -> None:
-        result = regex.Regex(
+        result = Regex(
             name="regex1",
             pattern="asd",
             location="any",
@@ -26,7 +31,7 @@ def test_create_indicator(self) -> None:
         self.assertEqual(result.type, "regex")
 
     def test_filter_entities_different_types(self) -> None:
-        regex_indicator = regex.Regex(
+        regex = Regex(
             name="regex1",
             pattern="asd",
             location="any",
@@ -34,55 +39,53 @@ def test_filter_entities_different_types(self) -> None:
         ).save()
 
         all_entities = list(Indicator.list())
-        regex_indicators = list(regex.Regex.list())
+        regex_entities = list(Regex.list())
 
         self.assertEqual(len(all_entities), 1)
-        self.assertEqual(len(regex_indicators), 1)
-        self.assertEqual(
-            regex_indicators[0].model_dump_json(), regex_indicator.model_dump_json()
-        )
+        self.assertEqual(len(regex_entities), 1)
+        self.assertEqual(regex_entities[0].model_dump_json(), regex.model_dump_json())
 
     def test_create_indicator_same_name_diff_types(self) -> None:
-        regex1 = regex.Regex(
+        regex = Regex(
             name="persistence1",
             pattern="asd",
             location="any",
             diamond=DiamondModel.capability,
         ).save()
-        regex2 = query.Query(
+        regex2 = Query(
             name="persistence1",
             pattern="asd",
             location="any",
             query_type="query",
             diamond=DiamondModel.capability,
         ).save()
-        self.assertNotEqual(regex1.id, regex2.id)
-        r = regex.Regex.find(name="persistence1")
-        q = query.Query.find(name="persistence1")
+        self.assertNotEqual(regex.id, regex2.id)
+        r = Regex.find(name="persistence1")
+        q = Query.find(name="persistence1")
         self.assertNotEqual(r.id, q.id)
 
     def test_regex_match(self) -> None:
-        regex_indicator = regex.Regex(
+        regex = Regex(
             name="regex1",
             pattern="Ba+dString",
             location="any",
             diamond=DiamondModel.capability,
         ).save()
 
-        result = regex_indicator.match("ThisIsAReallyBaaaadStringIsntIt")
+        result = regex.match("ThisIsAReallyBaaaadStringIsntIt")
         assert result is not None
         self.assertIsNotNone(result)
         self.assertEqual(result.name, "regex1")
         self.assertEqual(result.match, "BaaaadString")
 
     def test_regex_nomatch(self) -> None:
-        regex_indicator = regex.Regex(
+        regex = Regex(
             name="regex1",
             pattern="Blah",
             location="any",
             diamond=DiamondModel.capability,
         ).save()
-        result = regex_indicator.match("ThisIsAReallyBaaaadStringIsntIt")
+        result = regex.match("ThisIsAReallyBaaaadStringIsntIt")
         self.assertIsNone(result)
 
     def test_forensics_artifacts_indicator_extraction_file(self) -> None:
@@ -105,7 +108,7 @@ def test_forensics_artifacts_indicator_extraction_file(self) -> None:
         - Darwin
         - Linux"""
 
-        artifacts = forensicartifact.ForensicArtifact.from_yaml_string(pattern)
+        artifacts = ForensicArtifact.from_yaml_string(pattern)
         db_artifact = artifacts[0]
         self.assertIsNotNone(db_artifact.id)
         self.assertIsNotNone(db_artifact.created)
@@ -170,7 +173,7 @@ def test_forensics_artifacts_indicator_extraction_registry(self) -> None:
 supported_os:
 - Windows"""
 
-        artifacts = forensicartifact.ForensicArtifact.from_yaml_string(pattern)
+        artifacts = ForensicArtifact.from_yaml_string(pattern)
         db_artifact = artifacts[0]
         self.assertIsNotNone(db_artifact.id)
         self.assertIsNotNone(db_artifact.created)
@@ -244,9 +247,7 @@ def test_forensic_artifacts_parent_extraction(self):
     - blah3
 """
 
-        artifacts = forensicartifact.ForensicArtifact.from_yaml_string(
-            pattern, update_parents=True
-        )
+        artifacts = ForensicArtifact.from_yaml_string(pattern, update_parents=True)
         self.assertEqual(len(artifacts), 3)
 
         vertices, _, total = artifacts[0].neighbors()

From 0a6c3b157664b0f5e36f859a2d5373185f04352b Mon Sep 17 00:00:00 2001
From: Fred Baguelin 
Date: Wed, 18 Sep 2024 12:19:48 +0200
Subject: [PATCH 29/39] Revert analytics to main version

---
 plugins/analytics/public/censys.py        | 5 ++---
 plugins/analytics/public/github.py        | 5 ++---
 plugins/analytics/public/macaddress_io.py | 2 +-
 plugins/analytics/public/malshare.py      | 2 +-
 plugins/analytics/public/network_whois.py | 2 +-
 plugins/analytics/public/passive_total.py | 3 +--
 plugins/analytics/public/shodan.py        | 5 ++---
 plugins/analytics/public/shodan_api.py    | 2 +-
 8 files changed, 11 insertions(+), 15 deletions(-)

diff --git a/plugins/analytics/public/censys.py b/plugins/analytics/public/censys.py
index b1daaabb2..5c51a03a3 100644
--- a/plugins/analytics/public/censys.py
+++ b/plugins/analytics/public/censys.py
@@ -5,8 +5,7 @@
 
 from core import taskmanager
 from core.config.config import yeti_config
-from core.schemas import task
-from core.schemas.indicators.query import Query
+from core.schemas import indicator, task
 from core.schemas.observable import Observable
 
 
@@ -32,7 +31,7 @@ def run(self):
             api_secret=api_secret,
         )
 
-        censys_queries, _ = Query.filter({"query_type": "censys"})
+        censys_queries, _ = indicator.Query.filter({"query_type": "censys"})
 
         for query in censys_queries:
             ip_addresses = query_censys(hosts_api, query.pattern)
diff --git a/plugins/analytics/public/github.py b/plugins/analytics/public/github.py
index af3c21564..d2ee66abd 100644
--- a/plugins/analytics/public/github.py
+++ b/plugins/analytics/public/github.py
@@ -7,8 +7,7 @@
 
 from core import taskmanager
 from core.config.config import yeti_config
-from core.schemas import observable, task
-from core.schemas.indicators.query import Query
+from core.schemas import indicator, observable, task
 from core.schemas.observable import ObservableType
 
 
@@ -164,7 +163,7 @@ def run(self):
         auth = Auth.Token(github_token)
         self._github_api = Github(auth=auth)
 
-        github_query_indicators, _ = Query.filter({"query_type": "github"})
+        github_query_indicators, _ = indicator.Query.filter({"query_type": "github"})
         logging.info(
             f"[+] Found {len(github_query_indicators)} Github queries: {github_query_indicators}"
         )
diff --git a/plugins/analytics/public/macaddress_io.py b/plugins/analytics/public/macaddress_io.py
index a3cac4780..81137c2c2 100644
--- a/plugins/analytics/public/macaddress_io.py
+++ b/plugins/analytics/public/macaddress_io.py
@@ -9,7 +9,7 @@
 from core import taskmanager
 from core.config.config import yeti_config
 from core.schemas import task
-from core.schemas.entities.company import Company
+from core.schemas.entity import Company
 from core.schemas.observable import ObservableType
 from core.schemas.observables.mac_address import MacAddress
 
diff --git a/plugins/analytics/public/malshare.py b/plugins/analytics/public/malshare.py
index 4cb540ba0..ff81349e7 100644
--- a/plugins/analytics/public/malshare.py
+++ b/plugins/analytics/public/malshare.py
@@ -95,7 +95,7 @@ def each(self, observable: Observable):
             new_hash.add_context("malshare.com", context)
 
         if json_result["SSDEEP"]:
-            ssdeep_data = ssdeep.Ssdeep(value=json_result["SSDEEP"]).save()
+            ssdeep_data = ssdeep.SsdeepHash(value=json_result["SSDEEP"]).save()
             ssdeep_data.add_context("malshare.com", context)
             ssdeep_data.link_to(observable, "ssdeep", "malshare_query")
 
diff --git a/plugins/analytics/public/network_whois.py b/plugins/analytics/public/network_whois.py
index cd1927ed0..265f82dc9 100644
--- a/plugins/analytics/public/network_whois.py
+++ b/plugins/analytics/public/network_whois.py
@@ -2,7 +2,7 @@
 
 from core import taskmanager
 from core.schemas import task
-from core.schemas.entities.company import Company
+from core.schemas.entity import Company
 from core.schemas.observable import ObservableType
 from core.schemas.observables import email, ipv4
 
diff --git a/plugins/analytics/public/passive_total.py b/plugins/analytics/public/passive_total.py
index 243bdc6e1..68903c7bd 100644
--- a/plugins/analytics/public/passive_total.py
+++ b/plugins/analytics/public/passive_total.py
@@ -7,8 +7,7 @@
 from core import taskmanager
 from core.config.config import yeti_config
 from core.schemas import task
-from core.schemas.entities.company import Company
-from core.schemas.entities.phone import Phone
+from core.schemas.entity import Company, Phone
 from core.schemas.observable import Observable, ObservableType
 from core.schemas.observables import email, hostname, sha256
 
diff --git a/plugins/analytics/public/shodan.py b/plugins/analytics/public/shodan.py
index c5e8c6c53..e74b9f533 100644
--- a/plugins/analytics/public/shodan.py
+++ b/plugins/analytics/public/shodan.py
@@ -5,8 +5,7 @@
 
 from core import taskmanager
 from core.config.config import yeti_config
-from core.schemas import task
-from core.schemas.indicators.query import Query
+from core.schemas import indicator, task
 from core.schemas.observable import Observable
 
 
@@ -29,7 +28,7 @@ def run(self):
 
         shodan_api = Shodan(api_key)
 
-        shodan_queries, _ = Query.filter({"query_type": "shodan"})
+        shodan_queries, _ = indicator.Query.filter({"query_type": "shodan"})
 
         for query in shodan_queries:
             ip_addresses = query_shodan(shodan_api, query.pattern, result_limit)
diff --git a/plugins/analytics/public/shodan_api.py b/plugins/analytics/public/shodan_api.py
index 618c496b9..56e6ef8a9 100644
--- a/plugins/analytics/public/shodan_api.py
+++ b/plugins/analytics/public/shodan_api.py
@@ -5,7 +5,7 @@
 from core import taskmanager
 from core.config.config import yeti_config
 from core.schemas import task
-from core.schemas.entities.company import Company
+from core.schemas.entity import Company
 from core.schemas.observable import Observable, ObservableType
 from core.schemas.observables import asn, hostname, ipv4
 

From 0c074dcd16cb239c876a16af8ab36004040d9dbd Mon Sep 17 00:00:00 2001
From: Fred Baguelin 
Date: Wed, 18 Sep 2024 12:22:03 +0200
Subject: [PATCH 30/39] Replace ssdeep.SsdeepHash with ssdeep.Ssdeep

---
 plugins/analytics/public/malshare.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/plugins/analytics/public/malshare.py b/plugins/analytics/public/malshare.py
index ff81349e7..4cb540ba0 100644
--- a/plugins/analytics/public/malshare.py
+++ b/plugins/analytics/public/malshare.py
@@ -95,7 +95,7 @@ def each(self, observable: Observable):
             new_hash.add_context("malshare.com", context)
 
         if json_result["SSDEEP"]:
-            ssdeep_data = ssdeep.SsdeepHash(value=json_result["SSDEEP"]).save()
+            ssdeep_data = ssdeep.Ssdeep(value=json_result["SSDEEP"]).save()
             ssdeep_data.add_context("malshare.com", context)
             ssdeep_data.link_to(observable, "ssdeep", "malshare_query")
 

From 913f8e32197bf6332ef9e03832f42f147f12d423 Mon Sep 17 00:00:00 2001
From: Fred Baguelin 
Date: Wed, 18 Sep 2024 12:24:11 +0200
Subject: [PATCH 31/39] Revert feeds to main version

---
 plugins/feeds/public/abusech_malwarebazaar.py |  2 +-
 plugins/feeds/public/artifacts.py             |  5 ++--
 plugins/feeds/public/attack.py                | 29 +++++++------------
 plugins/feeds/public/cisa_kev.py              |  5 ++--
 plugins/feeds/public/et_open.py               | 18 +++++-------
 plugins/feeds/public/lolbas.py                |  9 ++----
 plugins/feeds/public/malpedia.py              | 16 +++++-----
 plugins/feeds/public/otx_alienvault.py        |  3 +-
 plugins/feeds/public/timesketch.py            |  2 +-
 .../public/wiz_cloud_threat_landscape.py      |  5 +---
 plugins/feeds/public/yaraify.py               |  3 +-
 11 files changed, 36 insertions(+), 61 deletions(-)

diff --git a/plugins/feeds/public/abusech_malwarebazaar.py b/plugins/feeds/public/abusech_malwarebazaar.py
index a9d8543ec..947ddb9e2 100644
--- a/plugins/feeds/public/abusech_malwarebazaar.py
+++ b/plugins/feeds/public/abusech_malwarebazaar.py
@@ -131,7 +131,7 @@ def analyze(self, block):
             imphash_data.tag(tags)
             malware_file.link_to(imphash_data, "imphash", self.name)
 
-        ssdeep_data = ssdeep.Ssdeep(value=context["ssdeep"]).save()
+        ssdeep_data = ssdeep.SsdeepHash(value=context["ssdeep"]).save()
         ssdeep_data.tag(tags)
         malware_file.link_to(ssdeep_data, "ssdeep", self.name)
 
diff --git a/plugins/feeds/public/artifacts.py b/plugins/feeds/public/artifacts.py
index c38f25aa1..c730e864c 100644
--- a/plugins/feeds/public/artifacts.py
+++ b/plugins/feeds/public/artifacts.py
@@ -9,8 +9,7 @@
 from artifacts.scripts import validator
 
 from core import taskmanager
-from core.schemas import task
-from core.schemas.indicators.forensicartifact import ForensicArtifact
+from core.schemas import indicator, task
 
 
 class ForensicArtifacts(task.FeedTask):
@@ -48,7 +47,7 @@ def run(self):
             with open(file, "r") as f:
                 yaml_string = f.read()
 
-            forensic_indicators = ForensicArtifact.from_yaml_string(
+            forensic_indicators = indicator.ForensicArtifact.from_yaml_string(
                 yaml_string, update_parents=False
             )
             for fi in forensic_indicators:
diff --git a/plugins/feeds/public/attack.py b/plugins/feeds/public/attack.py
index cc221b36d..56e502fdd 100644
--- a/plugins/feeds/public/attack.py
+++ b/plugins/feeds/public/attack.py
@@ -7,16 +7,7 @@
 from zipfile import ZipFile
 
 from core import taskmanager
-from core.schemas import task
-from core.schemas.entities.attack_pattern import AttackPattern
-from core.schemas.entities.campaign import Campaign
-from core.schemas.entities.course_of_action import CourseOfAction
-from core.schemas.entities.identity import Identity
-from core.schemas.entities.intrusion_set import IntrusionSet
-from core.schemas.entities.malware import Malware
-from core.schemas.entities.threat_actor import ThreatActor
-from core.schemas.entities.tool import Tool
-from core.schemas.entities.vulnerability import Vulnerability
+from core.schemas import entity, task
 
 
 def _format_context_from_obj(obj):
@@ -44,7 +35,7 @@ def _format_context_from_obj(obj):
 
 
 def _process_intrusion_set(obj):
-    intrusion_set = IntrusionSet(
+    intrusion_set = entity.IntrusionSet(
         name=obj["name"],
         aliases=obj.get("aliases", []) + obj.get("x_mitre_aliases", []),
         created=obj["created"],
@@ -56,7 +47,7 @@ def _process_intrusion_set(obj):
 
 
 def _process_malware(obj):
-    malware = Malware(
+    malware = entity.Malware(
         name=obj["name"],
         created=obj["created"],
         modified=obj["modified"],
@@ -77,7 +68,7 @@ def _process_attack_pattern(obj):
     for phase in obj["kill_chain_phases"]:
         kill_chain_phases.add(f'{phase["kill_chain_name"]}:{phase["phase_name"]}')
 
-    attack_pattern = AttackPattern(
+    attack_pattern = entity.AttackPattern(
         name=obj["name"],
         created=obj["created"],
         modified=obj["modified"],
@@ -90,7 +81,7 @@ def _process_attack_pattern(obj):
 
 
 def _process_course_of_action(obj):
-    course_of_action = CourseOfAction(
+    course_of_action = entity.CourseOfAction(
         name=obj["name"],
         created=obj["created"],
         modified=obj["modified"],
@@ -101,7 +92,7 @@ def _process_course_of_action(obj):
 
 
 def _process_identity(obj):
-    identity = Identity(
+    identity = entity.Identity(
         name=obj["name"],
         created=obj["created"],
         modified=obj["modified"],
@@ -115,7 +106,7 @@ def _process_identity(obj):
 
 
 def _process_threat_actor(obj):
-    threat_actor = ThreatActor(
+    threat_actor = entity.ThreatActor(
         name=obj["name"],
         created=obj["created"],
         modified=obj["modified"],
@@ -130,7 +121,7 @@ def _process_threat_actor(obj):
 
 
 def _process_campaign(obj):
-    campaign = Campaign(
+    campaign = entity.Campaign(
         name=obj["name"],
         created=obj["created"],
         modified=obj["modified"],
@@ -144,7 +135,7 @@ def _process_campaign(obj):
 
 
 def _process_tool(obj):
-    tool = Tool(
+    tool = entity.Tool(
         name=obj["name"],
         created=obj["created"],
         modified=obj["modified"],
@@ -161,7 +152,7 @@ def _process_tool(obj):
 
 
 def _process_vulnerability(obj):
-    vulnerabilty = Vulnerability(
+    vulnerabilty = entity.Vulnerability(
         name=obj["name"],
         created=obj["created"],
         modified=obj["modified"],
diff --git a/plugins/feeds/public/cisa_kev.py b/plugins/feeds/public/cisa_kev.py
index ccb188e53..0db462a15 100644
--- a/plugins/feeds/public/cisa_kev.py
+++ b/plugins/feeds/public/cisa_kev.py
@@ -3,8 +3,7 @@
 from typing import ClassVar
 
 from core import taskmanager
-from core.schemas import task
-from core.schemas.entities.vulnerability import Vulnerability
+from core.schemas import entity, task
 
 
 def _cves_as_dict(data):
@@ -152,7 +151,7 @@ def analyze_entry(self, entry: dict, cve_details: dict):
 
         name = f"{cve_id}"
         title = entry.get("vulnerabilityName", "")
-        vulnerability = Vulnerability(
+        vulnerability = entity.Vulnerability(
             name=name,
             title=title,
             description=description,
diff --git a/plugins/feeds/public/et_open.py b/plugins/feeds/public/et_open.py
index 6812107d5..0c3d12f2f 100644
--- a/plugins/feeds/public/et_open.py
+++ b/plugins/feeds/public/et_open.py
@@ -8,10 +8,6 @@
 from core import taskmanager
 from core.config.config import yeti_config
 from core.schemas import entity, indicator, task
-from core.schemas.entities.attack_pattern import AttackPattern
-from core.schemas.entities.malware import Malware
-from core.schemas.entities.vulnerability import Vulnerability
-from core.schemas.indicators.suricata import Suricata
 
 
 class ETOpen(task.FeedTask):
@@ -41,7 +37,7 @@ def run(self):
     def analyze(self, rule_suricata: rule.Rule):
         if not self._filter_rule(rule_suricata.metadata):
             return
-        ind_suricata_rule = Suricata(
+        ind_suricata_rule = indicator.Suricata(
             name=rule_suricata["msg"],
             pattern=rule_suricata["raw"],
             metadata=rule_suricata.metadata,
@@ -65,20 +61,20 @@ def analyze(self, rule_suricata: rule.Rule):
         if tags:
             ind_suricata_rule.tag(tags)
 
-    def _extract_cve(self, meta: str) -> Vulnerability:
+    def _extract_cve(self, meta: str) -> entity.Vulnerability:
         _, cve = meta.split(" ")
         if "_" in cve:
             cve = cve.replace("_", "-")
-        ind_cve = Vulnerability.find(name=cve)
+        ind_cve = entity.Vulnerability.find(name=cve)
         if not ind_cve:
-            ind_cve = Vulnerability(name=cve).save()
+            ind_cve = entity.Vulnerability(name=cve).save()
         return ind_cve
 
     def _extract_malware_family(self, meta: str):
         _, malware_family = meta.split(" ")
-        ind_malware_family = Malware.find(name=malware_family)
+        ind_malware_family = entity.Malware.find(name=malware_family)
         if not ind_malware_family:
-            ind_malware_family = Malware(name=malware_family).save()
+            ind_malware_family = entity.Malware(name=malware_family).save()
         return ind_malware_family
 
     def _extract_tags(self, metadata: list[str]) -> list[str]:
@@ -89,7 +85,7 @@ def _extract_tags(self, metadata: list[str]) -> list[str]:
                 tags.append(tag)
         return tags
 
-    def _extract_mitre_attack(self, meta: str) -> AttackPattern | None:
+    def _extract_mitre_attack(self, meta: str) -> entity.AttackPattern | None:
         _, mitre_id = meta.split(" ")
         ind_mitre_attack, nb_ent = entity.Entity.filter(
             query_args={"type": entity.EntityType.attack_pattern},
diff --git a/plugins/feeds/public/lolbas.py b/plugins/feeds/public/lolbas.py
index ab37c0b77..25a650e41 100644
--- a/plugins/feeds/public/lolbas.py
+++ b/plugins/feeds/public/lolbas.py
@@ -6,9 +6,6 @@
 
 from core import taskmanager
 from core.schemas import entity, indicator, task
-from core.schemas.entities.attack_pattern import AttackPattern
-from core.schemas.entities.tool import Tool
-from core.schemas.indicators.sigma import Sigma
 from core.schemas.observables import path
 
 
@@ -27,7 +24,7 @@ def run(self):
         if not response:
             return
         lolbas_json = response.json()
-        self._lolbas_attackpattern = AttackPattern(name="LOLBAS usage").save()
+        self._lolbas_attackpattern = entity.AttackPattern(name="LOLBAS usage").save()
         if not self._lolbas_attackpattern.description:
             self._lolbas_attackpattern.description = (
                 "Usage of living-off-the-land binaries and scripts"
@@ -75,7 +72,7 @@ def analyze_entry(self, entry: dict):
                         "Error processing sigma rule for %s: %s", entry["Name"], error
                     )
 
-    def process_sigma_rule(self, tool: Tool, detection: dict) -> None:
+    def process_sigma_rule(self, tool: entity.Tool, detection: dict) -> None:
         """Processes a Sigma rule as specified in the lolbas json."""
         url = detection["Sigma"]
         if not url:
@@ -95,7 +92,7 @@ def process_sigma_rule(self, tool: Tool, detection: dict) -> None:
         date = sigma_data["date"]
         date = datetime.strptime(date.strip(), "%Y/%m/%d")
         # create sigma indicator
-        sigma = Sigma(
+        sigma = indicator.Sigma(
             name=title,
             description=description,
             created=date,
diff --git a/plugins/feeds/public/malpedia.py b/plugins/feeds/public/malpedia.py
index 50c6af6e1..32799938b 100644
--- a/plugins/feeds/public/malpedia.py
+++ b/plugins/feeds/public/malpedia.py
@@ -3,9 +3,7 @@
 from typing import ClassVar
 
 from core import taskmanager
-from core.schemas import task
-from core.schemas.entities.intrusion_set import IntrusionSet
-from core.schemas.entities.malware import Malware
+from core.schemas import entity, task
 
 
 class MalpediaMalware(task.FeedTask):
@@ -34,9 +32,9 @@ def analyze_entry(self, malware_name: str, entry: dict):
         if not entry.get("common_name"):
             return
 
-        m = Malware.find(name=entry["common_name"])
+        m = entity.Malware.find(name=entry["common_name"])
         if not m:
-            m = Malware(name=entry["common_name"])
+            m = entity.Malware(name=entry["common_name"])
 
         m.aliases = entry.get("aliases", [])
         refs = entry.get("urls", [])
@@ -50,9 +48,9 @@ def analyze_entry(self, malware_name: str, entry: dict):
         m.add_context(context["source"], context)
         attributions = entry.get("attribution", [])
         for attribution in attributions:
-            intrusion_set = IntrusionSet.find(name=attribution)
+            intrusion_set = entity.IntrusionSet.find(name=attribution)
             if not intrusion_set:
-                intrusion_set = IntrusionSet(name=attribution).save()
+                intrusion_set = entity.IntrusionSet(name=attribution).save()
             intrusion_set.link_to(m, "uses", "Malpedia")
 
         tags = []
@@ -82,9 +80,9 @@ def run(self):
             self.analyze_entry(actor_name, entry)
 
     def analyze_entry(self, actor_name: str, entry: dict):
-        intrusion_set = IntrusionSet.find(name=entry["value"])
+        intrusion_set = entity.IntrusionSet.find(name=entry["value"])
         if not intrusion_set:
-            intrusion_set = IntrusionSet(name=entry["value"])
+            intrusion_set = entity.IntrusionSet(name=entry["value"])
 
         refs = entry.get("meta", {}).get("refs", [])
         context = {
diff --git a/plugins/feeds/public/otx_alienvault.py b/plugins/feeds/public/otx_alienvault.py
index 2fb8910cc..602cca5cf 100644
--- a/plugins/feeds/public/otx_alienvault.py
+++ b/plugins/feeds/public/otx_alienvault.py
@@ -10,7 +10,6 @@
 from core import taskmanager
 from core.config.config import yeti_config
 from core.schemas import entity, indicator, observable, task
-from core.schemas.entities.investigation import Investigation
 
 
 class OTXAlienvault(task.FeedTask):
@@ -68,7 +67,7 @@ def analyze(self, item):
         context["references"] = "\r\n".join(item["references"])
         context["description"] = item["description"]
         context["link"] = "https://otx.alienvault.com/pulse/%s" % item["id"]
-        investigation = Investigation(
+        investigation = entity.Investigation(
             name=item["name"],
             description=item["description"],
             reference=f"https://otx.alienvault.com/pulse/{item['id']}",
diff --git a/plugins/feeds/public/timesketch.py b/plugins/feeds/public/timesketch.py
index 1009bb488..628ffe9b7 100644
--- a/plugins/feeds/public/timesketch.py
+++ b/plugins/feeds/public/timesketch.py
@@ -6,7 +6,7 @@
 from core import taskmanager
 from core.config.config import yeti_config
 from core.schemas import observable, task
-from core.schemas.entities.investigation import Investigation
+from core.schemas.entity import Investigation
 from core.schemas.observables import hostname, ipv4, md5, path, sha1, sha256, url
 
 TIMESKETCH_TYPE_MAPPING = {
diff --git a/plugins/feeds/public/wiz_cloud_threat_landscape.py b/plugins/feeds/public/wiz_cloud_threat_landscape.py
index 62378aa85..b7cd0a982 100644
--- a/plugins/feeds/public/wiz_cloud_threat_landscape.py
+++ b/plugins/feeds/public/wiz_cloud_threat_landscape.py
@@ -7,10 +7,7 @@
 
 from core import taskmanager
 from core.schemas import task
-from core.schemas.entities.attack_pattern import AttackPattern
-from core.schemas.entities.campaign import Campaign
-from core.schemas.entities.intrusion_set import IntrusionSet
-from core.schemas.entities.tool import Tool
+from core.schemas.entity import AttackPattern, Campaign, IntrusionSet, Tool
 
 VALUE_PROPERTIES = [
     "tags",
diff --git a/plugins/feeds/public/yaraify.py b/plugins/feeds/public/yaraify.py
index 748530de0..51221a340 100644
--- a/plugins/feeds/public/yaraify.py
+++ b/plugins/feeds/public/yaraify.py
@@ -8,7 +8,6 @@
 
 from core import taskmanager
 from core.schemas import indicator, task
-from core.schemas.indicators.yara import Yara
 
 
 class YARAify(task.FeedTask):
@@ -41,7 +40,7 @@ def analyze_entry(self, entry: str):
             logging.error(f"Error compiling yara rule: {e}")
             return
         for r in yara_rules:
-            ind_obj = Yara(
+            ind_obj = indicator.Yara(
                 name=f"{r.identifier}",
                 pattern=entry,
                 diamond=indicator.DiamondModel.capability,

From fff01fdab3cc5a31f9fdd26a992d08cb0884c0d7 Mon Sep 17 00:00:00 2001
From: Fred Baguelin 
Date: Wed, 18 Sep 2024 12:24:26 +0200
Subject: [PATCH 32/39] Replace ssdeep.SsdeepHash with ssdeep.Ssdeep

---
 plugins/feeds/public/abusech_malwarebazaar.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/plugins/feeds/public/abusech_malwarebazaar.py b/plugins/feeds/public/abusech_malwarebazaar.py
index 947ddb9e2..a9d8543ec 100644
--- a/plugins/feeds/public/abusech_malwarebazaar.py
+++ b/plugins/feeds/public/abusech_malwarebazaar.py
@@ -131,7 +131,7 @@ def analyze(self, block):
             imphash_data.tag(tags)
             malware_file.link_to(imphash_data, "imphash", self.name)
 
-        ssdeep_data = ssdeep.SsdeepHash(value=context["ssdeep"]).save()
+        ssdeep_data = ssdeep.Ssdeep(value=context["ssdeep"]).save()
         ssdeep_data.tag(tags)
         malware_file.link_to(ssdeep_data, "ssdeep", self.name)
 

From 1c221ba7fde5e4e608bc6ac4c5260f648a18a9f9 Mon Sep 17 00:00:00 2001
From: Fred Baguelin 
Date: Wed, 18 Sep 2024 12:27:17 +0200
Subject: [PATCH 33/39] Restore apiv2 to main version

---
 core/web/apiv2/indicators.py  |  6 ++----
 core/web/apiv2/observables.py | 22 ++++++++++++++--------
 2 files changed, 16 insertions(+), 12 deletions(-)

diff --git a/core/web/apiv2/indicators.py b/core/web/apiv2/indicators.py
index 3638401c9..709802c4f 100644
--- a/core/web/apiv2/indicators.py
+++ b/core/web/apiv2/indicators.py
@@ -3,11 +3,11 @@
 
 from core.schemas import graph
 from core.schemas.indicator import (
+    ForensicArtifact,
     Indicator,
     IndicatorType,
     IndicatorTypes,
 )
-from core.schemas.indicators import forensicartifact
 from core.schemas.tag import MAX_TAGS_REQUEST
 
 
@@ -79,9 +79,7 @@ async def patch(request: PatchIndicatorRequest, indicator_id) -> IndicatorTypes:
 
     if db_indicator.type == IndicatorType.forensicartifact:
         if db_indicator.pattern != request.indicator.pattern:
-            return forensicartifact.ForensicArtifact.from_yaml_string(
-                request.indicator.pattern
-            )[0]
+            return ForensicArtifact.from_yaml_string(request.indicator.pattern)[0]
 
     update_data = request.indicator.model_dump(exclude_unset=True)
     updated_indicator = db_indicator.model_copy(update=update_data)
diff --git a/core/web/apiv2/observables.py b/core/web/apiv2/observables.py
index 4c1ae9de1..88704d11b 100644
--- a/core/web/apiv2/observables.py
+++ b/core/web/apiv2/observables.py
@@ -4,14 +4,20 @@
 from pydantic import BaseModel, ConfigDict, Field, conlist, field_validator
 
 from core.schemas import graph
-from core.schemas.observable import (
-    TYPE_MAPPING,
-    Observable,
-    ObservableType,
-    ObservableTypes,
-)
+from core.schemas.observable import TYPE_MAPPING, Observable, ObservableType
 from core.schemas.tag import MAX_TAG_LENGTH, MAX_TAGS_REQUEST
 
+ObservableTypes = ()
+
+for key in TYPE_MAPPING:
+    if key in ["observable", "observables"]:
+        continue
+    cls = TYPE_MAPPING[key]
+    if not ObservableTypes:
+        ObservableTypes = cls
+    else:
+        ObservableTypes |= cls
+
 
 class TagRequestMixin(BaseModel):
     tags: conlist(str, max_length=MAX_TAGS_REQUEST) = []
@@ -38,13 +44,13 @@ class NewObservableRequest(TagRequestMixin):
 class NewExtendedObservableRequest(TagRequestMixin):
     model_config = ConfigDict(extra="forbid")
 
-    observable: ObservableTypes = Field(discriminator="type")
+    observable: ObservableTypes = Field(discriminant="type")
 
 
 class PatchObservableRequest(BaseModel):
     model_config = ConfigDict(extra="forbid")
 
-    observable: ObservableTypes = Field(discriminator="type")
+    observable: ObservableTypes = Field(discriminant="type")
 
 
 class NewBulkObservableAddRequest(BaseModel):

From 11695615d90129a77d935d7456bc7b77d550f794 Mon Sep 17 00:00:00 2001
From: Fred Baguelin 
Date: Wed, 18 Sep 2024 12:35:16 +0200
Subject: [PATCH 34/39] Revert dfiq to main version

---
 core/schemas/dfiq.py | 11 +++++------
 1 file changed, 5 insertions(+), 6 deletions(-)

diff --git a/core/schemas/dfiq.py b/core/schemas/dfiq.py
index 2001aa645..282c5ea10 100644
--- a/core/schemas/dfiq.py
+++ b/core/schemas/dfiq.py
@@ -14,7 +14,6 @@
 from core.config.config import yeti_config
 from core.helpers import now
 from core.schemas import indicator
-from core.schemas.indicators import forensicartifact, query
 from core.schemas.model import YetiModel
 
 LATEST_SUPPORTED_DFIQ_VERSION = "1.1.0"
@@ -99,7 +98,7 @@ def extract_indicators(question: "DFIQQuestion") -> None:
                 continue
 
             if step.type in ("ForensicArtifact", "artifact"):
-                artifact = forensicartifact.ForensicArtifact.find(name=step.value)
+                artifact = indicator.ForensicArtifact.find(name=step.value)
                 if not artifact:
                     logging.warning(
                         "Missing artifact %s in %s", step.value, question.dfiq_id
@@ -109,9 +108,9 @@ def extract_indicators(question: "DFIQQuestion") -> None:
                 continue
 
             elif step.type and step.value and "query" in step.type:
-                query_indicator = query.Query.find(pattern=step.value)
-                if not query_indicator:
-                    query_indicator = query.Query(
+                query = indicator.Query.find(pattern=step.value)
+                if not query:
+                    query = indicator.Query(
                         name=f"{step.name} ({step.type})",
                         description=step.description or "",
                         pattern=step.value,
@@ -120,7 +119,7 @@ def extract_indicators(question: "DFIQQuestion") -> None:
                         location=step.type,
                         diamond=indicator.DiamondModel.victim,
                     ).save()
-                question.link_to(query_indicator, "query", "Uses query")
+                question.link_to(query, "query", "Uses query")
 
             else:
                 logging.warning(

From f4ccb2cc62872f6557e65bbda5230fda6aaa46f0 Mon Sep 17 00:00:00 2001
From: Fred Baguelin 
Date: Wed, 18 Sep 2024 12:36:04 +0200
Subject: [PATCH 35/39] create objects from indicator. Replace indicator
 variable with regex

---
 core/schemas/indicators/forensicartifact.py | 37 +++++++++------------
 1 file changed, 16 insertions(+), 21 deletions(-)

diff --git a/core/schemas/indicators/forensicartifact.py b/core/schemas/indicators/forensicartifact.py
index 793f16df1..b1f5744d6 100644
--- a/core/schemas/indicators/forensicartifact.py
+++ b/core/schemas/indicators/forensicartifact.py
@@ -9,7 +9,6 @@
 from pydantic import field_validator
 
 from core.schemas import indicator
-from core.schemas.indicators import regex
 
 
 class ForensicArtifact(indicator.Indicator):
@@ -19,9 +18,7 @@ class ForensicArtifact(indicator.Indicator):
     """
 
     _type_filter: ClassVar[str] = indicator.IndicatorType.forensicartifact
-    type: Literal[indicator.IndicatorType.forensicartifact] = (
-        indicator.IndicatorType.forensicartifact
-    )
+    type: Literal[indicator.IndicatorType.forensicartifact] = indicator.IndicatorType.forensicartifact
 
     sources: list[dict] = []
     aliases: list[str] = []
@@ -109,17 +106,17 @@ def save_indicators(self, create_links: bool = False):
                     pattern = re.escape(pattern).replace("\\*", ".*")
                     # Account for different path separators
                     pattern = re.sub(r"\\\\", r"[\\|/]", pattern)
-                    regex_indicator = regex.Regex.find(name=path)
-                    if not regex_indicator:
+                    regex = indicator.Regex.find(name=path)
+                    if not regex:
                         try:
-                            regex_indicator = regex.Regex(
+                            regex = indicator.Regex(
                                 name=path,
                                 pattern=pattern,
                                 location="filesystem",
                                 diamond=indicator.DiamondModel.victim,
                                 relevant_tags=self.relevant_tags,
                             ).save()
-                            indicators.append(regex_indicator)
+                            indicators.append(regex)
                         except Exception as error:
                             logging.error(
                                 f"Failed to create indicator for {path} (was: {source['attributes']['paths']}): {error}"
@@ -127,10 +124,10 @@ def save_indicators(self, create_links: bool = False):
                             continue
 
                     else:
-                        regex_indicator.relevant_tags = list(
-                            set(regex_indicator.relevant_tags + self.relevant_tags)
+                        regex.relevant_tags = list(
+                            set(regex.relevant_tags + self.relevant_tags)
                         )
-                        regex_indicator.save()
+                        regex.save()
         if source["type"] == definitions.TYPE_INDICATOR_WINDOWS_REGISTRY_KEY:
             for key in source["attributes"]["keys"]:
                 pattern = re.sub(r"\\\*$", "", key)
@@ -147,33 +144,31 @@ def save_indicators(self, create_links: bool = False):
                     )
                     pattern = pattern.replace("HKEY_LOCAL_MACHINE\\\\System\\\\", "")
 
-                regex_indicator = regex.Regex.find(name=key)
+                regex = indicator.Regex.find(name=key)
 
-                if not regex_indicator:
+                if not regex:
                     try:
-                        regex_indicator = regex.Regex(
+                        regex = indicator.Regex(
                             name=key,
                             pattern=pattern,
                             location="registry",
                             diamond=indicator.DiamondModel.victim,
                             relevant_tags=self.relevant_tags,
                         ).save()
-                        indicators.append(regex_indicator)
+                        indicators.append(regex)
                     except Exception as error:
                         logging.error(
                             f"Failed to create indicator for {key} (was: {source['attributes']['keys']}): {error}"
                         )
                         continue
                 else:
-                    regex_indicator.relevant_tags = list(
-                        set(regex_indicator.relevant_tags + self.relevant_tags)
+                    regex.relevant_tags = list(
+                        set(regex.relevant_tags + self.relevant_tags)
                     )
-                    regex_indicator.save()
+                    regex.save()
         if create_links:
             for indicator_obj in indicators:
-                indicator_obj.link_to(
-                    self, "indicates", f"Indicates {indicator_obj.name}"
-                )
+                indicator_obj.link_to(self, "indicates", f"Indicates {indicator_obj.name}")
         return indicators
 
 

From decbf6fa1dc4cb645fe626214ec7a7eee06c75b4 Mon Sep 17 00:00:00 2001
From: Fred Baguelin 
Date: Wed, 18 Sep 2024 12:37:36 +0200
Subject: [PATCH 36/39] Ruff format

---
 core/schemas/indicators/forensicartifact.py | 8 ++++++--
 tests/schemas/fixture.py                    | 4 +---
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/core/schemas/indicators/forensicartifact.py b/core/schemas/indicators/forensicartifact.py
index b1f5744d6..5e9398778 100644
--- a/core/schemas/indicators/forensicartifact.py
+++ b/core/schemas/indicators/forensicartifact.py
@@ -18,7 +18,9 @@ class ForensicArtifact(indicator.Indicator):
     """
 
     _type_filter: ClassVar[str] = indicator.IndicatorType.forensicartifact
-    type: Literal[indicator.IndicatorType.forensicartifact] = indicator.IndicatorType.forensicartifact
+    type: Literal[indicator.IndicatorType.forensicartifact] = (
+        indicator.IndicatorType.forensicartifact
+    )
 
     sources: list[dict] = []
     aliases: list[str] = []
@@ -168,7 +170,9 @@ def save_indicators(self, create_links: bool = False):
                     regex.save()
         if create_links:
             for indicator_obj in indicators:
-                indicator_obj.link_to(self, "indicates", f"Indicates {indicator_obj.name}")
+                indicator_obj.link_to(
+                    self, "indicates", f"Indicates {indicator_obj.name}"
+                )
         return indicators
 
 
diff --git a/tests/schemas/fixture.py b/tests/schemas/fixture.py
index 774a16879..4ba63a045 100644
--- a/tests/schemas/fixture.py
+++ b/tests/schemas/fixture.py
@@ -40,9 +40,7 @@ def test_something(self):
         hacker = hostname.Hostname(value="hacker.com").save()
         sus_hacker = hostname.Hostname(value="sus.hacker.com").save()
         mac_address.MacAddress(value="00:11:22:33:44:55").save()
-        generic_obs = generic.Generic(
-            value="SomeInterestingString"
-        ).save()
+        generic_obs = generic.Generic(value="SomeInterestingString").save()
         generic_obs.add_context("test_source", {"test": "test"})
 
         hacker.link_to(www_hacker, "domain", "Domain")

From 0dcd89b43fc45cbe350b64e44bd77d0fe6e0e1fe Mon Sep 17 00:00:00 2001
From: Fred Baguelin 
Date: Wed, 18 Sep 2024 12:54:40 +0200
Subject: [PATCH 37/39] Remove path_validator function and return regex matches
 in is_valid

---
 core/schemas/observables/path.py | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/core/schemas/observables/path.py b/core/schemas/observables/path.py
index dba75c59c..4c82641b9 100644
--- a/core/schemas/observables/path.py
+++ b/core/schemas/observables/path.py
@@ -9,13 +9,9 @@
 )
 
 
-def path_validator(value):
-    return LINUX_PATH_REGEX.match(value) or WINDOWS_PATH_REGEX.match(value)
-
-
 class Path(observable.Observable):
     type: Literal[observable.ObservableType.path] = observable.ObservableType.path
 
     @staticmethod
     def is_valid(value: str) -> bool:
-        return path_validator(value)
+        return LINUX_PATH_REGEX.match(value) or WINDOWS_PATH_REGEX.match(value)

From d920b00f6eb932764c629a7c2fb3034aad0ffbe4 Mon Sep 17 00:00:00 2001
From: Fred Baguelin 
Date: Thu, 19 Sep 2024 06:20:07 +0200
Subject: [PATCH 38/39] Update comments to detail where types are populated

---
 core/schemas/entity.py     | 3 ++-
 core/schemas/indicator.py  | 3 ++-
 core/schemas/observable.py | 3 ++-
 3 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/core/schemas/entity.py b/core/schemas/entity.py
index c4add2ecf..eee001bb7 100644
--- a/core/schemas/entity.py
+++ b/core/schemas/entity.py
@@ -9,7 +9,8 @@
 from core.schemas.model import YetiTagModel
 
 
-# forward declarations
+# Forward declarations
+# They are then populated by the load_entities function in __init__.py
 class EntityType(str, Enum): ...
 
 
diff --git a/core/schemas/indicator.py b/core/schemas/indicator.py
index c7434f74b..cb2adbbab 100644
--- a/core/schemas/indicator.py
+++ b/core/schemas/indicator.py
@@ -19,7 +19,8 @@ def future():
 DEFAULT_INDICATOR_VALIDITY_DAYS = 30
 
 
-# forward declarations
+# Forward declarations
+# They are then populated by the load_indicators function in __init__.py
 class IndicatorType(str, Enum): ...
 
 
diff --git a/core/schemas/observable.py b/core/schemas/observable.py
index 7f97b4e57..82dd5e1ae 100644
--- a/core/schemas/observable.py
+++ b/core/schemas/observable.py
@@ -17,7 +17,8 @@
 from core.schemas.model import YetiTagModel
 
 
-# forward declarations
+# Forward declarations
+# They are then populated by the load_observables function in __init__.py
 class ObservableType(str, Enum): ...
 
 

From d809ec22da543f18f088472547ea26d502ee8ea7 Mon Sep 17 00:00:00 2001
From: Fred Baguelin 
Date: Thu, 19 Sep 2024 06:22:54 +0200
Subject: [PATCH 39/39] Give more details about conversion from filename to
 entity name

---
 core/schemas/entities/private/README.md | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/core/schemas/entities/private/README.md b/core/schemas/entities/private/README.md
index 0f25cdf78..60db5d3db 100644
--- a/core/schemas/entities/private/README.md
+++ b/core/schemas/entities/private/README.md
@@ -3,3 +3,5 @@
 This directory is where you should place your private indicators. It could be
 named anything else, but this one has a `.gitignore` so you don't mess things
 up. ;-)
+
+Each entity defined with a filename containing `_` will be then represented in the API and the UI with `-`. For example, if you add a file `super_new_entity.py`, this entity will be defined as `super-new-entity`.  
\ No newline at end of file