Skip to content

Commit d5ce750

Browse files
tomchoptomchop
authored
refactor(schemas): register observable/entity/indicator types statically (#1285)
1 parent 7b642bf commit d5ce750

9 files changed

Lines changed: 442 additions & 133 deletions

File tree

‎core/schemas/__init__.py‎

Lines changed: 5 additions & 89 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,5 @@
1-
import importlib
2-
import inspect
3-
import logging
4-
import re
5-
from pathlib import Path
6-
7-
import aenum
8-
91
from core.events import message
10-
from core.schemas import (
2+
from core.schemas import ( # noqa: F401 (imported for their side-effect registration)
113
dfiq,
124
entity,
135
graph,
@@ -20,84 +12,8 @@
2012
user,
2113
)
2214

23-
logger = logging.getLogger(__name__)
24-
logging.getLogger().setLevel(logging.INFO)
25-
26-
27-
def register_module(module_name, base_module):
28-
"""
29-
Register the classes for the schema implementation files
30-
31-
module_name: The module name to load
32-
base_module: The base module to register the classes in (entity, indicator, observable)
33-
"""
34-
module = importlib.import_module(module_name)
35-
module_base_name = base_module.__name__.split(".")[-1]
36-
schema_base_class = getattr(base_module, module_base_name.capitalize())
37-
schema_type_mapping = getattr(base_module, "TYPE_MAPPING")
38-
schema_types = getattr(base_module, f"{module_base_name.capitalize()}Types", None)
39-
schema_enum = getattr(base_module, f"{module_base_name.capitalize()}Type")
40-
for _, obj in inspect.getmembers(module, inspect.isclass):
41-
if issubclass(obj, schema_base_class) and "type" in obj.model_fields:
42-
obs_type = obj.model_fields["type"].default
43-
logger.info(f"Registering class {obj.__name__} defining type <{obs_type}>")
44-
aenum.extend_enum(schema_enum, obs_type, obs_type)
45-
schema_type_mapping[obs_type] = obj
46-
setattr(base_module, obj.__name__, obj)
47-
if not schema_types:
48-
schema_types = obj
49-
else:
50-
schema_types |= obj
51-
setattr(base_module, f"{module_base_name.capitalize()}Types", schema_types)
52-
53-
54-
def register_classes(schema_root_type, base_module):
55-
"""
56-
Register the classes for the schema root type
57-
58-
schema_root_type: The schema root type to work with (entities, indicators, observables)
59-
base_module: The base module to register the classes in (entity, indicator, observable)
60-
"""
61-
module_base_name = base_module.__name__.split(".")[-1]
62-
logger.info(f"Registering {module_base_name} classes")
63-
for schema_file in Path(__file__).parent.glob(f"{schema_root_type}/**/*.py"):
64-
if schema_file.stem == "__init__":
65-
continue
66-
if schema_file.parent.stem == schema_root_type:
67-
module_name = f"core.schemas.{schema_root_type}.{schema_file.stem}"
68-
elif schema_file.parent.stem == "private":
69-
module_name = f"core.schemas.{schema_root_type}.private.{schema_file.stem}"
70-
try:
71-
register_module(module_name, base_module)
72-
except Exception:
73-
logger.exception(f"Failed to register classes from {module_name}")
74-
75-
76-
def load_entities():
77-
entity.TYPE_MAPPING = {"entity": entity.Entity, "entities": entity.Entity}
78-
register_classes("entities", entity)
79-
80-
81-
def load_indicators():
82-
indicator.TYPE_MAPPING = {
83-
"indicator": indicator.Indicator,
84-
"indicators": indicator.Indicator,
85-
}
86-
register_classes("indicators", indicator)
87-
88-
89-
def load_observables():
90-
observable.TYPE_MAPPING = {
91-
"observable": observable.Observable,
92-
"observables": observable.Observable,
93-
}
94-
if "guess" not in observable.ObservableType.__members__:
95-
aenum.extend_enum(observable.ObservableType, "guess", "guess")
96-
register_classes("observables", observable)
97-
98-
99-
load_observables()
100-
load_entities()
101-
load_indicators()
102-
15+
# Importing the schema modules above triggers their static type registration
16+
# (see the "Static type registry" sections in observable.py / entity.py /
17+
# indicator.py; dfiq.py registers its own types inline). Rebuild the event
18+
# message model now that every referenced schema type is importable.
10319
message.EventMessage.model_rebuild()

‎core/schemas/entity.py‎

Lines changed: 85 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,21 +1,19 @@
1+
from __future__ import annotations
2+
13
import datetime
24
from enum import Enum
3-
from typing import ClassVar, List, Literal
5+
from typing import ClassVar, List, Literal, Union
46

57
from pydantic import ConfigDict, Field, computed_field
68

79
from core import database_arango
810
from core.helpers import now
911
from core.schemas.model import YetiAclModel, YetiContextModel, YetiTagModel
1012

11-
12-
# Forward declarations
13-
# They are then populated by the load_entities function in __init__.py
14-
class EntityType(str, Enum): ...
15-
16-
17-
EntityTypes = ()
18-
TYPE_MAPPING = {}
13+
# EntityType, EntityTypes and TYPE_MAPPING are defined statically at the bottom
14+
# of this module (see "Static type registry"). TYPE_MAPPING must exist before
15+
# the functions below are *called*, which is always the case at runtime.
16+
TYPE_MAPPING: dict[str, type["Entity"]] = {}
1917

2018

2119
class Entity(
@@ -80,3 +78,81 @@ def save(*, name: str, type: str, tags: List[str] = None, **kwargs):
8078

8179
def find(*, name: str, **kwargs) -> "EntityTypes":
8280
return Entity.find(name=name, **kwargs)
81+
82+
83+
# ---------------------------------------------------------------------------
84+
# Static type registry (see observable.py for the rationale).
85+
# ---------------------------------------------------------------------------
86+
from core.schemas.entities.attack_pattern import AttackPattern # noqa: E402
87+
from core.schemas.entities.campaign import Campaign # noqa: E402
88+
from core.schemas.entities.company import Company # noqa: E402
89+
from core.schemas.entities.course_of_action import CourseOfAction # noqa: E402
90+
from core.schemas.entities.identity import Identity # noqa: E402
91+
from core.schemas.entities.intrusion_set import IntrusionSet # noqa: E402
92+
from core.schemas.entities.investigation import Investigation # noqa: E402
93+
from core.schemas.entities.malware import Malware # noqa: E402
94+
from core.schemas.entities.note import Note # noqa: E402
95+
from core.schemas.entities.phone import Phone # noqa: E402
96+
from core.schemas.entities.threat_actor import ThreatActor # noqa: E402
97+
from core.schemas.entities.tool import Tool # noqa: E402
98+
from core.schemas.entities.vulnerability import Vulnerability # noqa: E402
99+
from core.schemas.loader import load_private_types # noqa: E402
100+
101+
102+
class EntityType(str, Enum):
103+
# Member names must be valid identifiers, so the four hyphenated types use
104+
# underscores here; the *values* keep the wire-format hyphens.
105+
attack_pattern = "attack-pattern"
106+
campaign = "campaign"
107+
company = "company"
108+
course_of_action = "course-of-action"
109+
identity = "identity"
110+
intrusion_set = "intrusion-set"
111+
investigation = "investigation"
112+
malware = "malware"
113+
note = "note"
114+
phone = "phone"
115+
threat_actor = "threat-actor"
116+
tool = "tool"
117+
vulnerability = "vulnerability"
118+
119+
120+
_ENTITY_CLASSES: list[type[Entity]] = [
121+
AttackPattern,
122+
Campaign,
123+
Company,
124+
CourseOfAction,
125+
Identity,
126+
IntrusionSet,
127+
Investigation,
128+
Malware,
129+
Note,
130+
Phone,
131+
ThreatActor,
132+
Tool,
133+
Vulnerability,
134+
]
135+
136+
_private_entity_classes = load_private_types("core.schemas.entities", Entity)
137+
138+
TYPE_MAPPING = {"entity": Entity, "entities": Entity}
139+
for _cls in (*_ENTITY_CLASSES, *_private_entity_classes):
140+
TYPE_MAPPING[str(_cls.model_fields["type"].default)] = _cls
141+
142+
EntityTypes = Union[
143+
AttackPattern,
144+
Campaign,
145+
Company,
146+
CourseOfAction,
147+
Identity,
148+
IntrusionSet,
149+
Investigation,
150+
Malware,
151+
Note,
152+
Phone,
153+
ThreatActor,
154+
Tool,
155+
Vulnerability,
156+
]
157+
if _private_entity_classes:
158+
EntityTypes = Union[(EntityTypes, *_private_entity_classes)]

‎core/schemas/indicator.py‎

Lines changed: 55 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,9 @@
1+
from __future__ import annotations
2+
13
import datetime
24
import logging
35
from enum import Enum
4-
from typing import Any, ClassVar, Generator, List, Literal
6+
from typing import Any, ClassVar, Generator, List, Literal, Union
57

68
from pydantic import ConfigDict, Field, computed_field
79

@@ -18,14 +20,10 @@ def future():
1820

1921
DEFAULT_INDICATOR_VALIDITY_DAYS = 30
2022

21-
22-
# Forward declarations
23-
# They are then populated by the load_indicators function in __init__.py
24-
class IndicatorType(str, Enum): ...
25-
26-
27-
IndicatorTypes = ()
28-
TYPE_MAPPING = {}
23+
# IndicatorType, IndicatorTypes and TYPE_MAPPING are defined statically at the
24+
# bottom of this module (see "Static type registry"). TYPE_MAPPING must exist
25+
# before the functions below are *called*, which is always the case at runtime.
26+
TYPE_MAPPING: dict[str, type["Indicator"]] = {}
2927

3028

3129
class DiamondModel(Enum):
@@ -130,3 +128,51 @@ def save(
130128

131129
def find(*, name: str, **kwargs) -> "IndicatorTypes":
132130
return Indicator.find(name=name, **kwargs)
131+
132+
133+
# ---------------------------------------------------------------------------
134+
# Static type registry (see observable.py for the rationale).
135+
# ---------------------------------------------------------------------------
136+
from core.schemas.indicators.forensicartifact import ForensicArtifact # noqa: E402
137+
from core.schemas.indicators.query import Query # noqa: E402
138+
from core.schemas.indicators.regex import Regex # noqa: E402
139+
from core.schemas.indicators.sigma import Sigma # noqa: E402
140+
from core.schemas.indicators.suricata import Suricata # noqa: E402
141+
from core.schemas.indicators.yara import Yara # noqa: E402
142+
from core.schemas.loader import load_private_types # noqa: E402
143+
144+
145+
class IndicatorType(str, Enum):
146+
forensicartifact = "forensicartifact"
147+
query = "query"
148+
regex = "regex"
149+
sigma = "sigma"
150+
suricata = "suricata"
151+
yara = "yara"
152+
153+
154+
_INDICATOR_CLASSES: list[type[Indicator]] = [
155+
ForensicArtifact,
156+
Query,
157+
Regex,
158+
Sigma,
159+
Suricata,
160+
Yara,
161+
]
162+
163+
_private_indicator_classes = load_private_types("core.schemas.indicators", Indicator)
164+
165+
TYPE_MAPPING = {"indicator": Indicator, "indicators": Indicator}
166+
for _cls in (*_INDICATOR_CLASSES, *_private_indicator_classes):
167+
TYPE_MAPPING[str(_cls.model_fields["type"].default)] = _cls
168+
169+
IndicatorTypes = Union[
170+
ForensicArtifact,
171+
Query,
172+
Regex,
173+
Sigma,
174+
Suricata,
175+
Yara,
176+
]
177+
if _private_indicator_classes:
178+
IndicatorTypes = Union[(IndicatorTypes, *_private_indicator_classes)]

‎core/schemas/loader.py‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
"""Helper to load user-supplied ("private") schema types.
2+
3+
The observable/entity/indicator packages each expose a `private` subpackage
4+
where deployments can drop in their own subtypes. Unlike the public types,
5+
these are not known statically, so they are discovered at import time by
6+
globbing the `private` subpackage. This is the only remaining dynamic hook in
7+
the schema layer — public types are all registered explicitly.
8+
"""
9+
10+
import importlib
11+
import inspect
12+
import logging
13+
import pkgutil
14+
15+
logger = logging.getLogger(__name__)
16+
17+
18+
def load_private_types(package_name: str, base_class: type) -> list[type]:
19+
"""Return the schema subclasses defined in `<package_name>.private`.
20+
21+
Args:
22+
package_name: e.g. "core.schemas.observables".
23+
base_class: the family base class (Observable/Entity/Indicator);
24+
only its subclasses that declare a `type` field are returned.
25+
"""
26+
classes: list[type] = []
27+
try:
28+
private_pkg = importlib.import_module(f"{package_name}.private")
29+
except ModuleNotFoundError:
30+
return classes
31+
32+
for _, modname, _ in pkgutil.iter_modules(private_pkg.__path__):
33+
full_name = f"{package_name}.private.{modname}"
34+
try:
35+
module = importlib.import_module(full_name)
36+
except Exception:
37+
logger.exception(f"Failed to import private schema module {full_name}")
38+
continue
39+
for _, obj in inspect.getmembers(module, inspect.isclass):
40+
if (
41+
obj.__module__ == module.__name__
42+
and issubclass(obj, base_class)
43+
and "type" in getattr(obj, "model_fields", {})
44+
):
45+
logger.info(f"Registering private type {obj.__name__} from {full_name}")
46+
classes.append(obj)
47+
return classes

0 commit comments

Comments
 (0)