Towards better inference: bits → nibbles

octopoes/nibbles/definitions.py

@@ -0,0 +1,95 @@
+import importlib
+import pkgutil
+from logging import getLogger
+from pathlib import Path
+from types import MethodType, ModuleType
+
+from pydantic import BaseModel
+
+from octopoes.models import OOI
+
+NIBBLES_DIR = Path(__file__).parent
+NIBBLE_ATTR_NAME = "NIBBLE"
+NIBBLE_FUNC_NAME = "nibble"
+logger = getLogger(__name__)
+
+
+class NibbleParameterDefinition(BaseModel):
+    ooi_type: type[OOI]
+    relation_path: str | None = None
+
+    def __eq__(self, other):
+        if isinstance(other, NibbleParameterDefinition):
+            return vars(self) == vars(other)
+        elif isinstance(other, type):
+            return self.ooi_type == other
+        else:
+            return False
+
+    def __hash__(self):
+        return hash(str(self.ooi_type) + self.relation_path if self.relation_path else "\0")
+
+
+class NibbleDefinition:
+    id: str
+    signature: list[NibbleParameterDefinition]
+    min_scan_level: int = 1
+    default_enabled: bool = True
+    config_ooi_relation_path: str | None = None
+    payload: MethodType | None = None
+
+    def __init__(
+        self,
+        name: str,
+        signature: list,
+        min_scan_level: int = 1,
+        default_enabled: bool = True,
+        config_ooi_relation_path: str | None = None,
+    ):
+        self.id = name
+        self.signature = signature
+        self.min_scan_level = min_scan_level
+        self.default_enabled = default_enabled
+        self.config_ooi_relation_path = config_ooi_relation_path
+
+    def __call__(self, args):
+        if self.payload is None:
+            raise NotImplementedError
+        else:
+            return self.payload(*args)
+
+
+def get_nibble_definitions() -> list[NibbleDefinition]:
+    nibble_definitions = []
+
+    for package in pkgutil.walk_packages([str(NIBBLES_DIR)]):
+        if package.name in ["definitions", "runner"]:
+            continue
+
+        try:
+            module: ModuleType = importlib.import_module(".nibble", f"{NIBBLES_DIR.name}.{package.name}")
+
+            if hasattr(module, NIBBLE_ATTR_NAME):
+                nibble_definition: NibbleDefinition = getattr(module, NIBBLE_ATTR_NAME)
+
+                try:
+                    payload: ModuleType = importlib.import_module(
+                        f".{package.name}", f"{NIBBLES_DIR.name}.{package.name}"
+                    )
+                    if hasattr(payload, NIBBLE_FUNC_NAME):
+                        nibble_definition.payload = getattr(payload, NIBBLE_FUNC_NAME)
+                    else:
+                        logger.warning('module "%s" has no function %s', package.name, NIBBLE_FUNC_NAME)
+
+                except ModuleNotFoundError:
+                    logger.warning('package "%s" has no function nibble', package.name)
+
+                nibble_definitions.append(nibble_definition)
+
+            else:
+                logger.warning('module "%s" has no attribute %s', package.name, NIBBLE_ATTR_NAME)
+
+        except ModuleNotFoundError:
+            logger.warning('package "%s" has no module nibble', package.name)
+
+    return nibble_definitions

octopoes/nibbles/runner.py

@@ -0,0 +1,83 @@
+from collections.abc import Callable
+from datetime import datetime
+from itertools import chain, product
+from typing import TypeVar
+
+from nibbles.definitions import NibbleDefinition, get_nibble_definitions
+from octopoes.models import OOI
+from octopoes.models.types import type_by_name
+from octopoes.repositories.ooi_repository import OOIRepository
+from octopoes.repositories.origin_parameter_repository import OriginParameterRepository
+from octopoes.repositories.scan_profile_repository import ScanProfileRepository
+
+T = TypeVar("T")
+U = TypeVar("U")
+
+
+def otype(ooi: OOI) -> type[OOI]:
+    return type_by_name(ooi.get_ooi_type())
+
+
+def mergewith(func: Callable[[set[T], set[T]], set[T]], d1: dict[U, set[T]], d2: dict[U, set[T]]) -> dict[U, set[T]]:
+    return {k: func(d1.get(k, set()), d2.get(k, set())) for k in set(d1) | set(d2)}
+
+
+class NibblesRunner:
+    def __init__(
+        self,
+        ooi_repository: OOIRepository,
+        scan_profile_repository: ScanProfileRepository,
+        origin_parameter_repository: OriginParameterRepository,
+    ):
+        self.ooi_repository = ooi_repository
+        self.scan_profile_repository = scan_profile_repository
+        self.origin_parameter_repository = origin_parameter_repository
+        self.objects_by_type_cache: dict[type[OOI], set[OOI]]
+        self.update_nibbles()
+
+    def _retrieve(self, types: set[type[OOI]], valid_time: datetime) -> None:
+        cached_types = set(self.objects_by_type_cache)
+        target_types = set(filter(lambda x: x not in cached_types, types))
+        objects = self.ooi_repository.list_oois_by_object_types(target_types, valid_time)
+        objects_by_type = {t: {x for x in objects if isinstance(otype(x), t)} for t in set(map(otype, objects))}
+        self.objects_by_type_cache = mergewith(set.union, self.objects_by_type_cache, objects_by_type)
+
+    def _run(self, ooi: OOI, valid_time: datetime) -> dict[str, set[OOI]]:
+        retval: dict[str, set[OOI]] = {}
+        target_nibbles = list(filter(lambda x: type(ooi) in x.signature, self.nibbles))
+        self._retrieve(
+            set(map(lambda x: x.ooi_type, chain.from_iterable(map(lambda x: x.signature, target_nibbles)))), valid_time
+        )
+        for nibble in target_nibbles:
+            # TODO: filter OOI not abiding the parameters from radix
+            radix = [self.objects_by_type_cache[sgn.ooi_type] for sgn in nibble.signature]
+            results = set(
+                filter(lambda ooi: ooi is not None, chain(map(nibble, filter(lambda x: ooi in x, product(*radix)))))
+            )
+            if results:
+                retval |= {nibble.id: results}
+        return retval
+
+    def update_nibbles(self):
+        self.nibbles: list[NibbleDefinition] = get_nibble_definitions()
+
+    def _cleared(self, ooi: OOI, valid_time: datetime) -> bool:
+        ooi_level = self.scan_profile_repository.get(ooi.reference, valid_time).level.value
+        target_nibbles = list(filter(lambda x: type(ooi) in x.signature, self.nibbles))
+        return any(nibble.min_scan_level < ooi_level for nibble in target_nibbles)
+
+    def infer(self, stack: list[OOI], valid_time: datetime) -> dict[OOI, dict[str, set[OOI]]]:
+        retval: dict[OOI, dict[str, set[OOI]]] = {}
+        blockset = set(stack)
+        self.objects_by_type_cache = {}
+        if self._cleared(stack[-1], valid_time):
+            while stack:
+                ooi = stack.pop()
+                self.objects_by_type_cache = mergewith(set.union, self.objects_by_type_cache, {otype(ooi): {ooi}})
+                results = self._run(ooi, valid_time)
+                if results:
+                    blocks = set(chain.from_iterable(results.values()))
+                    stack += [ooi for ooi in blocks if ooi not in blockset]
+                    blockset |= blocks
+                    retval |= {ooi: results}
+        return retval

octopoes/octopoes/core/service.py

@@ -8,6 +8,7 @@
 import structlog
 from bits.definitions import get_bit_definitions
 from bits.runner import BitRunner
+from nibbles.runner import NibblesRunner
 from pydantic import TypeAdapter
 
 from octopoes.config.settings import (
@@ -77,6 +78,7 @@ def __init__(
         self.origin_parameter_repository = origin_parameter_repository
         self.scan_profile_repository = scan_profile_repository
         self.session = session
+        self.nibbles = NibblesRunner(ooi_repository, scan_profile_repository, origin_parameter_repository)
 
     @overload
     def _populate_scan_profiles(self, oois: ValuesView[OOI], valid_time: datetime) -> ValuesView[OOI]: ...
@@ -170,10 +172,10 @@ def save_origin(
             self.ooi_repository.get(origin.source, valid_time)
         except ObjectNotFoundException:
             if (
-                origin.origin_type not in [OriginType.DECLARATION, OriginType.AFFIRMATION]
+                origin.origin_type not in [OriginType.DECLARATION, OriginType.AFFIRMATION, OriginType.NIBBLE]
                 and origin.source not in origin.result
             ):
-                raise ValueError("Origin source of observation does not exist")
+                raise ValueError(f"Origin source [{origin.source}] does not exist")
             elif origin.origin_type == OriginType.AFFIRMATION:
                 logger.debug("Affirmation source %s already deleted", origin.source)
                 return
@@ -200,6 +202,7 @@ def save_origin(
             self.origin_repository.delete(origin, valid_time=valid_time)
 
     def _run_inference(self, origin: Origin, valid_time: datetime) -> None:
+        # The bit part of inferring
         bit_definition = get_bit_definitions().get(origin.method, None)
 
         if bit_definition is None:
@@ -234,6 +237,7 @@ def _run_inference(self, origin: Origin, valid_time: datetime) -> None:
             if len(configs) != 0:
                 config = configs[-1].config
 
+        resulting_oois: list[OOI] = []
         try:
             if isinstance(self.session, XTDBSession):
                 start = perf_counter()
@@ -252,10 +256,24 @@ def _run_inference(self, origin: Origin, valid_time: datetime) -> None:
                 self.session.client.submit_transaction(ops)
             else:
                 resulting_oois = BitRunner(bit_definition).run(source, parameters, config=config)
-            self.save_origin(origin, resulting_oois, valid_time)
         except Exception as e:
             logger.exception("Error running inference", exc_info=e)
 
+        self.save_origin(origin, resulting_oois, valid_time)
+
+        # The nibble part of inferring
+        resulting_nibble_oois = self.nibbles.infer([source], valid_time)
+        for source_ooi, results in resulting_nibble_oois.items():
+            self.ooi_repository.save(source_ooi, valid_time)
+            for nibble_id, result_oois in results.items():
+                inference_origin = Origin(
+                    method=nibble_id,
+                    origin_type=OriginType.NIBBLE,
+                    result=[ooi.reference for ooi in result_oois],
+                    source=source_ooi.reference,
+                )
+                self.save_origin(inference_origin, list(result_oois), valid_time)
+
     @staticmethod
     def check_path_level(path_level: int | None, current_level: int):
         return path_level is not None and path_level >= current_level

octopoes/octopoes/models/origin.py

@@ -11,6 +11,7 @@ class OriginType(Enum):
     OBSERVATION = "observation"
     INFERENCE = "inference"
     AFFIRMATION = "affirmation"
+    NIBBLE = "nibble"
 
 
 class Origin(BaseModel):

octopoes/tests/integration/test_nibbles.py

@@ -0,0 +1,47 @@
+import os
+import sys
+from datetime import datetime
+from unittest.mock import Mock
+
+import pytest
+from nibbles.definitions import NibbleDefinition, NibbleParameterDefinition
+
+from octopoes.core.service import OctopoesService
+from octopoes.models import OOI, ScanLevel
+from octopoes.models.ooi.network import Network
+
+if os.environ.get("CI") != "1":
+    pytest.skip("Needs XTDB multinode container.", allow_module_level=True)
+
+NMAX = 13
+dummy_nibble = NibbleDefinition(name="dummy", signature=[NibbleParameterDefinition(ooi_type=Network)])
+
+
+def dummy(network: Network):
+    global NMAX
+    if len(network.name) < NMAX:
+        new_name = network.name + "I"
+        return Network(name=new_name)
+
+
+dummy_nibble.payload = getattr(sys.modules[__name__], "dummy")
+
+
+def test_dummy_nibble(xtdb_octopoes_service: OctopoesService, event_manager: Mock, valid_time: datetime):
+    xtdb_octopoes_service.nibbles.nibbles = [dummy_nibble]
+    network = Network(name="internet")
+    xtdb_octopoes_service.ooi_repository.save(network, valid_time)
+    event_manager.complete_process_events(xtdb_octopoes_service)
+
+    assert xtdb_octopoes_service.ooi_repository.list_oois({Network}, valid_time).count == 1
+    assert xtdb_octopoes_service.ooi_repository.list_oois({OOI}, valid_time).count == 3
+
+    sp = xtdb_octopoes_service.scan_profile_repository.get(network.reference, valid_time)
+    new_sp = sp.model_copy()
+    new_sp.level = ScanLevel.L2
+    xtdb_octopoes_service.scan_profile_repository.save(sp, new_sp, valid_time)
+    event_manager.complete_process_events(xtdb_octopoes_service)
+
+    ctx = 1 + NMAX - len(network.name)
+    assert xtdb_octopoes_service.ooi_repository.list_oois({Network}, valid_time).count == ctx
+    assert xtdb_octopoes_service.ooi_repository.list_oois({OOI}, valid_time).count == 3 * ctx