"""Port defining the common interface for data import services. All import services (Atomic Red Team, Sigma, CALDERA, etc.) follow the same contract: they receive a database session and return a summary dict with import statistics. New import sources can be added by: 1. Implementing the ``ImportService`` protocol in a new module 2. Registering the handler in the ``IMPORT_REGISTRY`` This satisfies the Open/Closed Principle — the system is open for new import sources without modifying existing code. """ # Enable future language features for compatibility from __future__ import annotations # Import Any, Protocol, runtime_checkable from typing from typing import Any, Protocol, runtime_checkable # Import Session from sqlalchemy.orm from sqlalchemy.orm import Session # Apply the @runtime_checkable decorator @runtime_checkable # Define class ImportService class ImportService(Protocol): """Contract for any data-import operation. Each implementation is a callable ``(Session) -> dict`` that downloads, parses, and upserts records from an external source. """ # Define function __call__ def __call__(self, db: Session) -> dict[str, Any]: """Execute the import operation against the given database session. Args: db (Session): Active SQLAlchemy session to use for all DB operations. Returns: dict[str, Any]: Summary statistics for the import run (e.g. created, updated, skipped counts). """ # ... ... # Define class ImportServiceEntry class ImportServiceEntry: """Lazy-loading wrapper that resolves a module-level function on first call.""" # Assign __slots__ = ("_module_path", "_func_name", "_resolved") __slots__ = ("_module_path", "_func_name", "_resolved") # Define function __init__ def __init__(self, module_path: str, func_name: str) -> None: """Initialise the lazy entry with the module path and function name to resolve later. Args: module_path (str): Dotted Python module path, e.g. ``"app.services.atomic_import_service"``. func_name (str): Name of the callable to import from *module_path*. Returns: None """ # Assign self._module_path = module_path self._module_path = module_path # Assign self._func_name = func_name self._func_name = func_name # Assign self._resolved = None self._resolved: ImportService | None = None # Define function __call__ def __call__(self, db: Session) -> dict[str, Any]: """Resolve the import function on first call and invoke it with *db*. Args: db (Session): SQLAlchemy session passed through to the underlying import function. Returns: dict[str, Any]: Import statistics returned by the underlying function (e.g. counts of created/updated/skipped records). """ # Check: self._resolved is None if self._resolved is None: # Import importlib import importlib # Assign mod = importlib.import_module(self._module_path) mod = importlib.import_module(self._module_path) # Assign self._resolved = getattr(mod, self._func_name) self._resolved = getattr(mod, self._func_name) # Return self._resolved(db) return self._resolved(db) # Apply the @property decorator @property # Define function source_info def source_info(self) -> str: """Return a human-readable identifier for this import entry. Returns: str: The fully qualified function reference as ``"."``. """ # Return f"{self._module_path}.{self._func_name}" return f"{self._module_path}.{self._func_name}" # Assign IMPORT_REGISTRY = { IMPORT_REGISTRY: dict[str, ImportServiceEntry] = { # Literal argument value "atomic_red_team": ImportServiceEntry( # Literal argument value "app.services.atomic_import_service", "import_atomic_red_team", ), # Literal argument value "sigma": ImportServiceEntry( # Literal argument value "app.services.sigma_import_service", "sync", ), # Literal argument value "lolbas": ImportServiceEntry( # Literal argument value "app.services.lolbas_import_service", "sync", ), # Literal argument value "gtfobins": ImportServiceEntry( # Literal argument value "app.services.lolbas_import_service", "sync_gtfobins", ), # Literal argument value "caldera": ImportServiceEntry( # Literal argument value "app.services.caldera_import_service", "sync", ), # Literal argument value "elastic_rules": ImportServiceEntry( # Literal argument value "app.services.elastic_import_service", "sync", ), # Literal argument value "mitre_cti": ImportServiceEntry( # Literal argument value "app.services.threat_actor_import_service", "sync", ), # Literal argument value "d3fend": ImportServiceEntry( # Literal argument value "app.services.d3fend_import_service", "sync", ), } # Define function get_import_handler def get_import_handler(source_name: str) -> ImportServiceEntry | None: """Look up the import handler for *source_name*. Returns ``None`` when no handler is registered. """ # Return IMPORT_REGISTRY.get(source_name) return IMPORT_REGISTRY.get(source_name)