From 57c4bd0e7cf767a773e851e0286c80d72ba4d3bd Mon Sep 17 00:00:00 2001 From: Salad Dais Date: Fri, 22 Dec 2023 21:23:34 +0000 Subject: [PATCH] Improve AIS support --- hippolyzer/lib/base/inventory.py | 24 ++++++++++++++++++++---- hippolyzer/lib/base/legacy_schema.py | 20 +++++++++++++++----- tests/base/test_legacy_schema.py | 17 +++++++++++++++++ 3 files changed, 52 insertions(+), 9 deletions(-) diff --git a/hippolyzer/lib/base/inventory.py b/hippolyzer/lib/base/inventory.py index 84dc87f..8c65bdf 100644 --- a/hippolyzer/lib/base/inventory.py +++ b/hippolyzer/lib/base/inventory.py @@ -17,7 +17,6 @@ import inspect import logging import secrets import struct -import typing import weakref from io import StringIO from typing import * @@ -190,7 +189,7 @@ class InventoryBase(SchemaBase): writer.write("\t}\n") -class InventoryDifferences(typing.NamedTuple): +class InventoryDifferences(NamedTuple): changed: List[InventoryNodeBase] removed: List[InventoryNodeBase] @@ -400,7 +399,6 @@ class InventoryNodeBase(InventoryBase, _HasName): @dataclasses.dataclass class InventoryContainerBase(InventoryNodeBase): - # TODO: Not a string in AIS type: AssetType = schema_field(SchemaEnumField(AssetType)) @property @@ -461,7 +459,6 @@ class InventoryCategory(InventoryContainerBase): VERSION_NONE: ClassVar[int] = -1 cat_id: UUID = schema_field(SchemaUUID) - # TODO: not a string in AIS pref_type: FolderType = schema_field(SchemaEnumField(FolderType), llsd_name="preferred_type") name: str = schema_field(SchemaMultilineStr) owner_id: Optional[UUID] = schema_field(SchemaUUID, default=None) @@ -488,6 +485,18 @@ class InventoryCategory(InventoryContainerBase): type=AssetType.CATEGORY, ) + @classmethod + def _get_fields_dict(cls, llsd_flavor: Optional[str] = None): + fields = super()._get_fields_dict(llsd_flavor) + if llsd_flavor == "ais": + # AIS is smart enough to know that all categories are asset type category... + fields.pop("type") + # These have different names though + fields["type_default"] = fields.pop("preferred_type") + fields["agent_id"] = fields.pop("owner_id") + fields["category_id"] = fields.pop("cat_id") + return fields + __hash__ = InventoryNodeBase.__hash__ @@ -573,5 +582,12 @@ class InventoryItem(InventoryNodeBase): creation_date=block["CreationDate"], ) + def to_llsd(self, flavor: str = "legacy"): + val = super().to_llsd(flavor=flavor) + if flavor == "ais": + # There's little chance this differs from owner ID, just place it. + val["agent_id"] = val["permissions"]["owner_id"] + return val + INVENTORY_TYPES: Tuple[Type[InventoryNodeBase], ...] = (InventoryCategory, InventoryObject, InventoryItem) diff --git a/hippolyzer/lib/base/legacy_schema.py b/hippolyzer/lib/base/legacy_schema.py index 1f2d39f..b6b7da1 100644 --- a/hippolyzer/lib/base/legacy_schema.py +++ b/hippolyzer/lib/base/legacy_schema.py @@ -104,6 +104,13 @@ class SchemaStr(SchemaFieldSerializer[str]): class SchemaUUID(SchemaFieldSerializer[UUID]): + @classmethod + def from_llsd(cls, val: Any, flavor: str) -> UUID: + # FetchInventory2 will return a string, but we want a UUID. It's not an issue + # for us to return a UUID later there because it'll just cast to string if + # that's what it wants + return UUID(val) + @classmethod def deserialize(cls, val: str) -> UUID: return UUID(val) @@ -157,11 +164,11 @@ def parse_schema_line(line: str): @dataclasses.dataclass class SchemaBase(abc.ABC): @classmethod - def _get_fields_dict(cls, llsd=False): + def _get_fields_dict(cls, llsd_flavor: Optional[str] = None): fields_dict = {} for field in dataclasses.fields(cls): field_name = field.name - if llsd: + if llsd_flavor: field_name = field.metadata.get("llsd_name") or field_name fields_dict[field_name] = field return fields_dict @@ -181,7 +188,7 @@ class SchemaBase(abc.ABC): @classmethod def from_llsd(cls, inv_dict: Dict, flavor: str = "legacy"): - fields = cls._get_fields_dict(llsd=True) + fields = cls._get_fields_dict(llsd_flavor=flavor) obj_dict = {} for key, val in inv_dict.items(): if key in fields: @@ -205,7 +212,10 @@ class SchemaBase(abc.ABC): else: raise ValueError(f"Unsupported spec for {key!r}, {spec!r}") else: - LOG.warning(f"Unknown key {key!r}") + if flavor != "ais": + # AIS has a number of different fields that are irrelevant depending on + # what exactly sent the payload + LOG.warning(f"Unknown key {key!r}") return cls._obj_from_dict(obj_dict) def to_bytes(self) -> bytes: @@ -219,7 +229,7 @@ class SchemaBase(abc.ABC): def to_llsd(self, flavor: str = "legacy"): obj_dict = {} - for field_name, field in self._get_fields_dict(llsd=True).items(): + for field_name, field in self._get_fields_dict(llsd_flavor=flavor).items(): spec = field.metadata.get("spec") # Not meant to be serialized if not spec: diff --git a/tests/base/test_legacy_schema.py b/tests/base/test_legacy_schema.py index 4d0900c..eb976c2 100644 --- a/tests/base/test_legacy_schema.py +++ b/tests/base/test_legacy_schema.py @@ -54,6 +54,7 @@ INV_CATEGORY = """\tinv_category\t0 \t\ttype\tlsltext \t\tpref_type\tlsltext \t\tname\tScripts| +\t\towner_id\ta2e76fcd-9360-4f6d-a924-000000000003 \t} """ @@ -160,6 +161,22 @@ class TestLegacyInv(unittest.TestCase): ] ) + def test_llsd_serialization_ais(self): + model = InventoryModel.from_str(INV_CATEGORY) + self.assertEqual( + [ + { + 'agent_id': UUID('a2e76fcd-9360-4f6d-a924-000000000003'), + 'category_id': UUID('f4d91477-def1-487a-b4f3-6fa201c17376'), + 'name': 'Scripts', + 'parent_id': UUID('00000000-0000-0000-0000-000000000000'), + 'type_default': 10, + 'version': -1 + } + ], + model.to_llsd("ais") + ) + def test_llsd_legacy_equality(self): new_model = InventoryModel.from_llsd(self.model.to_llsd()) self.assertEqual(self.model, new_model)