Load schemas dynamically

core/schemas/__init__.py

@@ -0,0 +1,120 @@
+import importlib
+import inspect
+import logging
+from pathlib import Path
+
+import aenum
+
+from core.schemas import entity, indicator, observable
+
+logger = logging.getLogger(__name__)
+
+
+def load_entities():
+    logger.info("Registering entities")
+    modules = dict()
+    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)
+        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):
+                entity.TYPE_MAPPING[enum_value] = obj
+                setattr(entity, obj.__name__, 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
+
+
+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
+                setattr(indicator, obj.__name__, 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
+                setattr(observable, obj.__name__, 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()

core/schemas/entities/attack_pattern.py

@@ -0,0 +1,10 @@
+from typing import ClassVar, Literal
+
+from core.schemas import entity
+
+
+class AttackPattern(entity.Entity):
+    _type_filter: ClassVar[str] = entity.EntityType.attack_pattern
+    type: Literal[entity.EntityType.attack_pattern] = entity.EntityType.attack_pattern
+    aliases: list[str] = []
+    kill_chain_phases: list[str] = []

core/schemas/entities/campaign.py

@@ -0,0 +1,16 @@
+import datetime
+from typing import ClassVar, Literal
+
+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: Literal[entity.EntityType.campaign] = entity.EntityType.campaign
+
+    aliases: list[str] = []
+    first_seen: datetime.datetime = Field(default_factory=now)
+    last_seen: datetime.datetime = Field(default_factory=now)

core/schemas/entities/company.py

@@ -0,0 +1,8 @@
+from typing import ClassVar, Literal
+
+from core.schemas import entity
+
+
+class Company(entity.Entity):
+    type: Literal[entity.EntityType.company] = entity.EntityType.company
+    _type_filter: ClassVar[str] = entity.EntityType.company

core/schemas/entities/course_of_action.py

@@ -0,0 +1,10 @@
+from typing import ClassVar, Literal
+
+from core.schemas import entity
+
+
+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
+    )

core/schemas/entities/identity.py

@@ -0,0 +1,12 @@
+from typing import ClassVar, Literal
+
+from core.schemas import entity
+
+
+class Identity(entity.Entity):
+    _type_filter: ClassVar[str] = entity.EntityType.identity
+    type: Literal[entity.EntityType.identity] = entity.EntityType.identity
+
+    identity_class: str = ""
+    sectors: list[str] = []
+    contact_information: str = ""

core/schemas/entities/intrusion_set.py

@@ -0,0 +1,16 @@
+import datetime
+from typing import ClassVar, Literal
+
+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: Literal[entity.EntityType.intrusion_set] = entity.EntityType.intrusion_set
+
+    aliases: list[str] = []
+    first_seen: datetime.datetime = Field(default_factory=now)
+    last_seen: datetime.datetime = Field(default_factory=now)

core/schemas/entities/investigation.py

@@ -0,0 +1,10 @@
+from typing import ClassVar, Literal
+
+from core.schemas import entity
+
+
+class Investigation(entity.Entity):
+    _type_filter: ClassVar[str] = entity.EntityType.investigation
+    type: Literal[entity.EntityType.investigation] = entity.EntityType.investigation
+
+    reference: str = ""

core/schemas/entities/malware.py

@@ -0,0 +1,12 @@
+from typing import ClassVar, Literal
+
+from core.schemas import entity
+
+
+class Malware(entity.Entity):
+    _type_filter: ClassVar[str] = entity.EntityType.malware
+    type: Literal[entity.EntityType.malware] = entity.EntityType.malware
+
+    kill_chain_phases: list[str] = []
+    aliases: list[str] = []
+    family: str = ""

core/schemas/entities/note.py

@@ -0,0 +1,8 @@
+from typing import ClassVar, Literal
+
+from core.schemas import entity
+
+
+class Note(entity.Entity):
+    type: Literal[entity.EntityType.note] = entity.EntityType.note
+    _type_filter: ClassVar[str] = entity.EntityType.note

core/schemas/entities/phone.py

@@ -0,0 +1,8 @@
+from typing import ClassVar, Literal
+
+from core.schemas import entity
+
+
+class Phone(entity.Entity):
+    type: Literal[entity.EntityType.phone] = entity.EntityType.phone
+    _type_filter: ClassVar[str] = entity.EntityType.phone

core/schemas/entities/private/.gitignore

@@ -0,0 +1,4 @@
+*
+!.gitignore
+!README.md
+!__init__.py

core/schemas/entities/private/README.md

@@ -0,0 +1,7 @@
+### 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. ;-)
+
+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`.  

core/schemas/entities/threat_actor.py

@@ -0,0 +1,17 @@
+import datetime
+from typing import ClassVar, Literal
+
+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: Literal[entity.EntityType.threat_actor] = 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)

core/schemas/entities/tool.py

@@ -0,0 +1,12 @@
+from typing import ClassVar, Literal
+
+from core.schemas import entity
+
+
+class Tool(entity.Entity):
+    _type_filter: ClassVar[str] = entity.EntityType.tool
+    type: Literal[entity.EntityType.tool] = entity.EntityType.tool
+
+    aliases: list[str] = []
+    kill_chain_phases: list[str] = []
+    tool_version: str = ""

core/schemas/entities/vulnerability.py

@@ -0,0 +1,46 @@
+import re
+from enum import Enum
+from typing import ClassVar, Literal
+
+from pydantic import Field
+
+from core.schemas import entity
+
+vulnerability_matcher = re.compile(
+    r"(?P<pre>\W?)(?P<search>CVE-\d{4}-\d{4,7})(?P<post>\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: Literal[entity.EntityType.vulnerability] = 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