diff --git a/app/service/data_svc.py b/app/service/data_svc.py index d5e9d8ccf..56c91ba3e 100644 --- a/app/service/data_svc.py +++ b/app/service/data_svc.py @@ -46,11 +46,19 @@ class DataService(DataServiceInterface, BaseService): + _INDEXED_FIELDS = { + 'agents': ['paw'], + 'abilities': ['ability_id'], + 'planners': ['name'], + 'obfuscators': ['name'], + } + def __init__(self): self.log = self.add_service('data_svc', self) self.schema = dict(agents=[], planners=[], adversaries=[], abilities=[], sources=[], operations=[], schedules=[], plugins=[], obfuscators=[], objectives=[], data_encoders=[]) self.ram = copy.deepcopy(self.schema) + self._index = {} # collection -> field -> {value: [objects]} @staticmethod def _iter_data_files(): @@ -127,14 +135,52 @@ async def load_data(self, plugins=()): async def reload_data(self, plugins=()): await self._load(plugins) + def _index_object(self, collection, obj): + """Index an object by its indexed fields for fast lookup.""" + indexed_fields = self._INDEXED_FIELDS.get(collection, []) + for field in indexed_fields: + val = getattr(obj, field, None) + if val is not None: + if collection not in self._index: + self._index[collection] = {} + if field not in self._index[collection]: + self._index[collection][field] = {} + if val not in self._index[collection][field]: + self._index[collection][field][val] = [] + if obj not in self._index[collection][field][val]: + self._index[collection][field][val].append(obj) + + def _remove_from_index(self, collection, obj): + """Remove an object from the index.""" + indexed_fields = self._INDEXED_FIELDS.get(collection, []) + for field in indexed_fields: + val = getattr(obj, field, None) + if val is not None and collection in self._index and field in self._index[collection]: + objs = self._index[collection][field].get(val, []) + if obj in objs: + objs.remove(obj) + async def store(self, c_object): try: - return c_object.store(self.ram) + result = c_object.store(self.ram) + # Index the object for fast lookups + for collection, objs in self.ram.items(): + if result in objs if isinstance(objs, list) else False: + self._index_object(collection, result) + break + return result except Exception as e: self.log.error('[!] can only store first-class objects: %s' % e) async def locate(self, object_name, match=None): try: + # Use index for single-key lookups on indexed fields + if match and len(match) == 1 and object_name in self._INDEXED_FIELDS: + key, val = next(iter(match.items())) + if key in self._INDEXED_FIELDS[object_name]: + idx = self._index.get(object_name, {}).get(key, {}) + candidates = idx.get(val, []) + return [obj for obj in candidates if obj.match(match)] return [obj for obj in self.ram[object_name] if obj.match(match)] except Exception as e: self.log.error('[!] LOCATE: %s' % e) diff --git a/tests/security/test_indexed_lookups.py b/tests/security/test_indexed_lookups.py new file mode 100644 index 000000000..beba12a66 --- /dev/null +++ b/tests/security/test_indexed_lookups.py @@ -0,0 +1,35 @@ +import unittest +from unittest.mock import MagicMock + + +class TestIndexedLookups(unittest.TestCase): + def test_index_and_lookup(self): + from app.service.data_svc import DataService + svc = DataService.__new__(DataService) + svc._index = {} + svc.log = MagicMock() + svc.ram = {'agents': []} + svc.schema = {'agents': []} + + obj = MagicMock() + obj.paw = 'testpaw' + obj.match = lambda m: m is None or all(getattr(obj, k, None) == v for k, v in m.items()) + obj.store = lambda ram: (ram['agents'].append(obj) or obj) if obj not in ram['agents'] else obj + + import asyncio + loop = asyncio.new_event_loop() + try: + loop.run_until_complete(svc.store(obj)) + result = loop.run_until_complete(svc.locate('agents', dict(paw='testpaw'))) + self.assertEqual(len(result), 1) + self.assertEqual(result[0].paw, 'testpaw') + + # Non-existent lookup + result2 = loop.run_until_complete(svc.locate('agents', dict(paw='nonexistent'))) + self.assertEqual(len(result2), 0) + finally: + loop.close() + + +if __name__ == '__main__': + unittest.main()