More changes

Author: Schamper

dissect/database/ese/cursor.py

@@ -3,7 +3,7 @@
 from typing import TYPE_CHECKING
 
 from dissect.database.ese.btree import BTree
-from dissect.database.ese.exception import NoNeighbourPageError
+from dissect.database.ese.exception import KeyNotFoundError, NoNeighbourPageError
 from dissect.database.ese.record import Record
 
 if TYPE_CHECKING:
@@ -64,7 +64,28 @@ def reset(self) -> None:
         if self._secondary:
             self._secondary.reset()
 
-    def search(self, **kwargs: RecordValue) -> Record:
+    def make_key(self, *args: RecordValue, **kwargs: RecordValue) -> bytes:
+        """Generate a key for this index from the given values.
+
+        Args:
+            *args: The values to generate a key for.
+            **kwargs: The columns and values to generate a key for.
+
+        Returns:
+            The generated key as bytes.
+        """
+        if not args and not kwargs:
+            raise ValueError("At least one value must be provided")
+
+        if args and kwargs:
+            raise ValueError("Cannot mix positional and keyword arguments in make_key")
+
+        if args and not len(args) == 1 and not isinstance(args[0], list):
+            raise ValueError("When using positional arguments, provide a single list of values")
+
+        return self.index.make_key(args[0] if args else kwargs)
+
+    def search(self, *args: RecordValue, **kwargs: RecordValue) -> Record:
         """Search the index for the requested values.
 
         Searching modifies the cursor state. Searching again will search from the current position.
@@ -76,7 +97,7 @@ def search(self, **kwargs: RecordValue) -> Record:
         Returns:
             A :class:`~dissect.database.ese.record.Record` object of the found record.
         """
-        key = self.index.make_key(kwargs)
+        key = self.make_key(*args, **kwargs)
         return self.search_key(key, exact=True)
 
     def search_key(self, key: bytes, exact: bool = True) -> Record:
@@ -90,13 +111,13 @@ def search_key(self, key: bytes, exact: bool = True) -> Record:
         self._primary.search(key, exact)
         return self._record()
 
-    def seek(self, **kwargs: RecordValue) -> None:
+    def seek(self, *args: RecordValue, **kwargs: RecordValue) -> None:
         """Seek to the record with the given values.
 
         Args:
             **kwargs: The columns and values to seek to.
         """
-        key = self.index.make_key(kwargs)
+        key = self.make_key(*args, **kwargs)
         self.search_key(key, exact=False)
 
     def seek_key(self, key: bytes) -> None:
@@ -130,7 +151,10 @@ def find_all(self, **kwargs: RecordValue) -> Iterator[Record]:
         other_columns = kwargs
 
         # We need at least an exact match on the indexed columns
-        self.search(**indexed_columns)
+        try:
+            self.search(**indexed_columns)
+        except KeyNotFoundError:
+            return
 
         current_key = self._primary.node().key
 

dissect/database/ese/ese.py

@@ -36,6 +36,7 @@ def __init__(self, fh: BinaryIO, impacket_compat: bool = False):
         self.fh = fh
         self.impacket_compat = impacket_compat
 
+        self.fh.seek(0)
         self.header = c_ese.DBFILEHDR(fh)
         if self.header.ulMagic != ulDAEMagic:
             raise InvalidDatabase("invalid file header signature")

dissect/database/ese/index.py

@@ -105,20 +105,27 @@ def key_from_record(self, record: Record) -> bytes:
         values = {c.name: record[c.name] for c in self.columns}
         return self.make_key(values)
 
-    def make_key(self, values: dict[str, RecordValue]) -> bytes:
+    def make_key(self, values: list[RecordValue] | dict[str, RecordValue]) -> bytes:
         """Generate a key out of the given values.
 
         Args:
-            values: A map of the column names and values to generate a key for.
+            values: A map of the column names and values to generate a key for, or a list of values in the order of the
+                    index columns.
         """
         key_buf = []
         key_remaining = self._key_most
 
-        for column in self.columns:
-            if column.name not in values:
-                break
+        if isinstance(values, dict):
+            tmp = []
+            for column in self.columns:
+                if column.name not in values:
+                    break
+                tmp.append(values[column.name])
+
+            values = tmp
 
-            key_part = encode_key(self, column, values[column.name], self._var_seg_mac)
+        for column, value in zip(self.columns, values, strict=False):
+            key_part = encode_key(self, column, value, self._var_seg_mac)
             key_buf.append(key_part)
             key_remaining -= len(key_part)
 

dissect/database/ese/ntds/__init__.py

@@ -1,12 +1,13 @@
 from __future__ import annotations
 
 from dissect.database.ese.ntds.ntds import NTDS
-from dissect.database.ese.ntds.object import Computer, Group, Server, User
+from dissect.database.ese.ntds.object import Computer, Group, Object, Server, User
 
 __all__ = [
     "NTDS",
     "Computer",
     "Group",
+    "Object",
     "Server",
     "User",
 ]