From cb7ebfbdc92a7d39dd5e283ac038fd6cdcda7c8b Mon Sep 17 00:00:00 2001 From: Benny Date: Tue, 24 Dec 2024 22:27:52 +0100 Subject: [PATCH] Beginnings of NibbleRepository etc --- octopoes/nibbles/definitions.py | 6 +++++- octopoes/nibbles/runner.py | 25 ++++++++++++++++------ octopoes/octopoes/core/app.py | 13 +++++++++-- octopoes/octopoes/core/service.py | 4 +++- octopoes/tests/conftest.py | 15 +++++++++++-- octopoes/tests/integration/test_nibbles.py | 6 ++++++ 6 files changed, 56 insertions(+), 13 deletions(-) diff --git a/octopoes/nibbles/definitions.py b/octopoes/nibbles/definitions.py index 73eba39ec48..1d031c1f569 100644 --- a/octopoes/nibbles/definitions.py +++ b/octopoes/nibbles/definitions.py @@ -8,7 +8,7 @@ import structlog from pydantic import BaseModel -from xxhash import xxh3_128_hexdigest as xxh3 # INFO: xxh3_64_hexdigest is faster but hash more collision probabilities +from xxhash import xxh3_128_hexdigest as xxh3 from octopoes.models import OOI, Reference @@ -36,6 +36,7 @@ class NibbleDefinition(BaseModel): id: str signature: list[NibbleParameter] query: str | Callable[[list[Reference | None]], str] | None = None + enabled: bool = True _payload: MethodType | None = None _checksum: str | None = None @@ -48,6 +49,9 @@ def __call__(self, args: Iterable[OOI]) -> OOI | Iterable[OOI | None] | None: def __hash__(self): return hash(self.id) + def _ini(self) -> dict[str, Any]: + return {"id": self.id, "enabled": self.enabled, "checksum": self._checksum} + def get_nibble_definitions() -> dict[str, NibbleDefinition]: nibble_definitions = {} diff --git a/octopoes/nibbles/runner.py b/octopoes/nibbles/runner.py index 7bd15b1b3d2..664ecbdc267 100644 --- a/octopoes/nibbles/runner.py +++ b/octopoes/nibbles/runner.py @@ -3,19 +3,25 @@ from datetime import datetime from typing import Any -from xxhash import xxh3_128_hexdigest as xxh3 # INFO: xxh3_64_hexdigest is faster but hash more collision probabilities +from xxhash import xxh3_128_hexdigest as xxh3 from nibbles.definitions import NibbleDefinition, get_nibble_definitions from octopoes.models import OOI, Reference from octopoes.models.origin import Origin, OriginType +from octopoes.repositories.nibble_repository import NibbleRepository from octopoes.repositories.ooi_repository import OOIRepository from octopoes.repositories.origin_repository import OriginRepository -from octopoes.repositories.scan_profile_repository import ScanProfileRepository def merge_results( d1: dict[OOI, dict[str, dict[tuple[Any, ...], set[OOI]]]], d2: dict[OOI, dict[str, dict[tuple[Any, ...], set[OOI]]]] ) -> dict[OOI, dict[str, dict[tuple[Any, ...], set[OOI]]]]: + """ + Merge new runner results with old runner results + d1: runner_results + d2: runner_results + --> runner_results + """ return { key: { nibble_id: { @@ -31,6 +37,9 @@ def merge_results( def flatten(items: Iterable[Any | Iterable[Any | None] | None]) -> Iterable[OOI]: + """ + Retrieve OOIs as returned from the nibble + """ for item in items: if isinstance(item, OOI): yield item @@ -43,6 +52,9 @@ def flatten(items: Iterable[Any | Iterable[Any | None] | None]) -> Iterable[OOI] def nibble_hasher(data: Iterable, additional: str | None = None) -> str: + """ + Hash the nibble generated data with its content together with the nibble checksum + """ return xxh3( "".join( [ @@ -58,16 +70,15 @@ def nibble_hasher(data: Iterable, additional: str | None = None) -> str: class NibblesRunner: def __init__( - self, - ooi_repository: OOIRepository, - origin_repository: OriginRepository, - scan_profile_repository: ScanProfileRepository, + self, ooi_repository: OOIRepository, origin_repository: OriginRepository, nibble_repository: NibbleRepository ): self.ooi_repository = ooi_repository self.origin_repository = origin_repository - self.scan_profile_repository = scan_profile_repository self.cache: dict[OOI, dict[str, dict[tuple[Any, ...], set[OOI]]]] = {} + self.nibble_repository = nibble_repository self.nibbles: dict[str, NibbleDefinition] = get_nibble_definitions() + inis = [nibble._ini() for nibble in self.nibbles.values()] + nibble_repository.put_many(inis, datetime.now()) def __del__(self): self._write(datetime.now()) diff --git a/octopoes/octopoes/core/app.py b/octopoes/octopoes/core/app.py index af61d604f11..a9384634623 100644 --- a/octopoes/octopoes/core/app.py +++ b/octopoes/octopoes/core/app.py @@ -4,6 +4,7 @@ from octopoes.config.settings import GATHER_BIT_METRICS, QUEUE_NAME_OCTOPOES, Settings from octopoes.core.service import OctopoesService from octopoes.events.manager import EventManager, get_rabbit_channel +from octopoes.repositories.nibble_repository import XTDBNibbleRepository from octopoes.repositories.ooi_repository import XTDBOOIRepository from octopoes.repositories.origin_parameter_repository import XTDBOriginParameterRepository from octopoes.repositories.origin_repository import XTDBOriginRepository @@ -37,10 +38,18 @@ def bootstrap_octopoes(settings: Settings, client: str, xtdb_session: XTDBSessio origin_repository = XTDBOriginRepository(event_manager, xtdb_session) origin_param_repository = XTDBOriginParameterRepository(event_manager, xtdb_session) scan_profile_repository = XTDBScanProfileRepository(event_manager, xtdb_session) + nibble_repository = XTDBNibbleRepository(xtdb_session) if GATHER_BIT_METRICS: return OctopoesService( - ooi_repository, origin_repository, origin_param_repository, scan_profile_repository, xtdb_session + ooi_repository, + origin_repository, + origin_param_repository, + scan_profile_repository, + nibble_repository, + xtdb_session, ) else: - return OctopoesService(ooi_repository, origin_repository, origin_param_repository, scan_profile_repository) + return OctopoesService( + ooi_repository, origin_repository, origin_param_repository, scan_profile_repository, nibble_repository + ) diff --git a/octopoes/octopoes/core/service.py b/octopoes/octopoes/core/service.py index 88a4fd41b41..8881283b99b 100644 --- a/octopoes/octopoes/core/service.py +++ b/octopoes/octopoes/core/service.py @@ -42,6 +42,7 @@ ) from octopoes.models.transaction import TransactionRecord from octopoes.models.tree import ReferenceTree +from octopoes.repositories.nibble_repository import NibbleRepository from octopoes.repositories.ooi_repository import OOIRepository from octopoes.repositories.origin_parameter_repository import OriginParameterRepository from octopoes.repositories.origin_repository import OriginRepository @@ -71,13 +72,14 @@ def __init__( origin_repository: OriginRepository, origin_parameter_repository: OriginParameterRepository, scan_profile_repository: ScanProfileRepository, + nibble_repository: NibbleRepository, session: XTDBSession | None = None, ): self.ooi_repository = ooi_repository self.origin_repository = origin_repository self.origin_parameter_repository = origin_parameter_repository self.scan_profile_repository = scan_profile_repository - self.nibbler = NibblesRunner(ooi_repository, origin_repository, scan_profile_repository) + self.nibbler = NibblesRunner(ooi_repository, origin_repository, nibble_repository) self.session = session @overload diff --git a/octopoes/tests/conftest.py b/octopoes/tests/conftest.py index ed978733267..5315b6ab0c8 100644 --- a/octopoes/tests/conftest.py +++ b/octopoes/tests/conftest.py @@ -34,6 +34,7 @@ Service, Website, ) +from octopoes.repositories.nibble_repository import XTDBNibbleRepository from octopoes.repositories.ooi_repository import OOIRepository, XTDBOOIRepository from octopoes.repositories.origin_parameter_repository import XTDBOriginParameterRepository from octopoes.repositories.origin_repository import XTDBOriginRepository @@ -202,7 +203,7 @@ def app_settings(): @pytest.fixture def octopoes_service() -> OctopoesService: - return OctopoesService(Mock(), Mock(), Mock(), Mock()) + return OctopoesService(Mock(), Mock(), Mock(), Mock(), Mock()) @pytest.fixture @@ -291,15 +292,25 @@ def xtdb_scan_profile_repository(xtdb_session: XTDBSession, event_manager) -> It yield XTDBScanProfileRepository(event_manager, xtdb_session) +@pytest.fixture +def xtdb_nibble_repository(xtdb_session: XTDBSession) -> Iterator[XTDBNibbleRepository]: + yield XTDBNibbleRepository(xtdb_session) + + @pytest.fixture def xtdb_octopoes_service( xtdb_ooi_repository: XTDBOOIRepository, xtdb_origin_repository: XTDBOriginRepository, xtdb_origin_parameter_repository: XTDBOriginParameterRepository, xtdb_scan_profile_repository: XTDBScanProfileRepository, + xtdb_nibble_repository: XTDBNibbleRepository, ) -> OctopoesService: return OctopoesService( - xtdb_ooi_repository, xtdb_origin_repository, xtdb_origin_parameter_repository, xtdb_scan_profile_repository + xtdb_ooi_repository, + xtdb_origin_repository, + xtdb_origin_parameter_repository, + xtdb_scan_profile_repository, + xtdb_nibble_repository, ) diff --git a/octopoes/tests/integration/test_nibbles.py b/octopoes/tests/integration/test_nibbles.py index eb30a105168..0199e77cad2 100644 --- a/octopoes/tests/integration/test_nibbles.py +++ b/octopoes/tests/integration/test_nibbles.py @@ -343,3 +343,9 @@ def test_nibbles_update(xtdb_octopoes_service: OctopoesService, event_manager: M assert xtdb_octopoes_service.ooi_repository.list_oois({Finding}, valid_time).count == 1 assert xtdb_octopoes_service.ooi_repository.list_oois({KATFindingType}, valid_time).count == 1 + + +def test_nibble_states(xtdb_octopoes_service: OctopoesService, event_manager: Mock, valid_time: datetime): + assert len(xtdb_octopoes_service.nibbler.nibble_repository.get_all(valid_time)) == len( + xtdb_octopoes_service.nibbler.nibbles + )