Changes

Author: Schamper

dissect/target/helpers/regutil.py

@@ -30,8 +30,16 @@
 RE_GLOB_MAGIC = re.compile(r"[*?[]")
 RE_REGFLEX_NAME_VALUE = re.compile(r'^"(?P<name>(?:[^"\\]|\\.)*?)"=(?P<value>.*)')
 
-REG_FLEX_HEADER = "Windows Registry Editor Version 5.00"
-_REG_LINE_LIMIT = 80
+REGFLEX_HEADER = "Windows Registry Editor Version 5.00"
+_REGFLEX_LINE_LIMIT = 80
+
+SHORTNAMES = {
+    "HKLM": "HKEY_LOCAL_MACHINE",
+    "HKCC": "HKEY_CURRENT_CONFIG",
+    "HKCU": "HKEY_CURRENT_USER",
+    "HKCR": "HKEY_CLASSES_ROOT",
+    "HKU": "HKEY_USERS",
+}
 
 
 KeyType = regf.IndexLeaf | regf.FastLeaf | regf.HashLeaf | regf.IndexRoot | regf.KeyNode
@@ -798,110 +806,11 @@ def map_definition(self, fh: TextIO) -> None:
 
         vhive.map_key(vkey.path, vkey)
 
-    def dump_definition(self, fh: TextIO, paths: list[str] | None = None) -> None:
-        """Dump the parsed registry hives to a .reg file format.
-
-        This is the inverse of :meth:`map_definition`. When ``paths`` is
-        given, only the listed registry subtrees are written; otherwise every
-        hive stored in :attr:`hives` is exported in full.
-
-        The header always contains a comment line indicating that the file was
-        generated by RegFlex. When ``paths`` is provided each path is also
-        listed as an individual comment line so that the exported file is
-        self-documenting.
-
-        Args:
-            fh: A file-like object opened in text mode to write the registry export to.
-            paths: Optional list of full registry paths to export
-                (e.g. ``["HKEY_LOCAL_MACHINE\\\\SOFTWARE\\\\Test"]``).
-                When ``None`` all hives are exported in full.
-        """
-        fh.write(f"{REG_FLEX_HEADER}\n\n")
-        fh.write("; This .reg file was generated by RegFlex (part of Dissect Target)\n")
-        if paths:
-            fh.write("; Registry paths exported:\n")
-            fh.writelines(f";   {path}\n" for path in paths)
-        fh.write("\n")
-
-        if not paths:
-            for hive_name, hive in self.hives.items():
-                for subkey in hive.root().subkeys():
-                    write_reg_key(fh, subkey, f"{hive_name}\\{subkey.name}")
-            return
-
-        for path in paths:
-            hive_name, _, key_path = path.partition("\\")
-            hive_name = hive_name.upper()
-            hive = self.hives.get(hive_name)
-            if hive is None:
-                log.warning("RegFlex: hive %r not found, skipping path %r", hive_name, path)
-                continue
-            try:
-                key = hive.key(key_path) if key_path else hive.root()
-            except RegistryKeyNotFoundError:
-                log.warning("RegFlex: key %r not found in hive %r, skipping", key_path, hive_name)
-                continue
-            write_reg_key(fh, key, path)
-
 
 class RegFlexHive(VirtualHive):
     pass
 
 
-def _escape_reg_str(value: str) -> str:
-    """Escape special characters in a string for .reg file format.
-
-    Backslashes are doubled and double-quote characters are backslash-escaped,
-    matching the encoding used by the Windows Registry Editor.
-
-    Args:
-        value: The string to escape.
-
-    Returns:
-        The escaped string, safe for use inside a quoted .reg value.
-    """
-    return value.replace("\\", "\\\\").replace('"', '\\"')
-
-
-def _format_hex_data(data: bytes, prefix: str) -> str:
-    """Format binary data as hex for .reg format, wrapping long lines with continuation.
-
-    Produces comma-separated lowercase hex byte pairs. Lines longer than
-    :data:`_REG_LINE_LIMIT` characters are split with a trailing backslash and
-    the continuation line is indented by two spaces, matching the layout
-    produced by the Windows Registry Editor.
-
-    Args:
-        data: The bytes to encode.
-        prefix: The line prefix (e.g. ``'"name"=hex:'``). Included verbatim at
-            the start of the first line and used when computing line length.
-
-    Returns:
-        The formatted string, potentially spanning multiple lines.
-    """
-    if not data:
-        return prefix
-
-    hex_pairs = [f"{b:02x}" for b in data]
-    lines = []
-    current = prefix
-
-    for i, pair in enumerate(hex_pairs):
-        is_last = i == len(hex_pairs) - 1
-        entry = pair if is_last else f"{pair},"
-        # Reserve one column for the backslash continuation character on non-final segments
-        fits = len(current) + len(entry) + (0 if is_last else 1) <= _REG_LINE_LIMIT
-
-        if not fits and i > 0:
-            lines.append(f"{current}\\")
-            current = f"  {entry}"
-        else:
-            current += entry
-
-    lines.append(current)
-    return "\n".join(lines)
-
-
 class RegFlexKey(VirtualKey):
     pass
 
@@ -977,12 +886,66 @@ def parse_flex_value(value: str) -> tuple[RegistryValueType, ValueType]:
     raise ValueError(f"Unsupported registry flex value type: {value!r}")
 
 
-def write_flex_value(value: RegistryValue) -> str:
+def _escape_flex_str(value: str) -> str:
+    """Escape special characters in a string for ``.reg`` file format.
+
+    Backslashes are doubled and double-quote characters are backslash-escaped,
+    matching the encoding used by the Windows Registry Editor.
+
+    Args:
+        value: The string to escape.
+
+    Returns:
+        The escaped string, safe for use inside a quoted ``.reg`` value.
+    """
+    return value.replace("\\", "\\\\").replace('"', '\\"')
+
+
+def _format_flex_hex(data: bytes, prefix: str) -> str:
+    """Format binary data as hex for .reg format, wrapping long lines with continuation.
+
+    Produces comma-separated lowercase hex byte pairs. Lines longer than
+    :data:`_REGFLEX_LINE_LIMIT` characters are split with a trailing backslash and
+    the continuation line is indented by two spaces, matching the layout
+    produced by the Windows Registry Editor.
+
+    Args:
+        data: The bytes to encode.
+        prefix: The line prefix (e.g. ``'"name"=hex:'``). Included verbatim at
+            the start of the first line and used when computing line length.
+
+    Returns:
+        The formatted string, potentially spanning multiple lines.
+    """
+    if not data:
+        return prefix
+
+    hex_pairs = [f"{b:02x}" for b in data]
+    lines = []
+    current = prefix
+
+    for i, pair in enumerate(hex_pairs):
+        is_last = i == len(hex_pairs) - 1
+        entry = pair if is_last else f"{pair},"
+        # Reserve one column for the backslash continuation character on non-final segments
+        fits = len(current) + len(entry) + (0 if is_last else 1) <= _REGFLEX_LINE_LIMIT
+
+        if not fits and i > 0:
+            lines.append(f"{current}\\")
+            current = f"  {entry}"
+        else:
+            current += entry
+
+    lines.append(current)
+    return "\n".join(lines)
+
+
+def format_flex_value(value: RegistryValue) -> str:
     """Format a :class:`RegistryValue` as a ``.reg`` file value line.
 
     This is the inverse of :func:`parse_flex_value`.
 
-    When the registry type is :attr:`RegistryValueType.NONE` the Python type
+    When the value type is :attr:`RegistryValueType.NONE` the Python type
     of the value data is used to infer an appropriate registry type:
 
     - :class:`str`   → :attr:`RegistryValueType.SZ`
@@ -1001,7 +964,7 @@ def write_flex_value(value: RegistryValue) -> str:
     data = value.value
     vtype = RegistryValueType(value.type)
 
-    name_str = "@" if name == "" else f'"{_escape_reg_str(name)}"'
+    name_str = "@" if name == "" else f'"{_escape_flex_str(name)}"'
 
     # Infer a sensible registry type from the Python value when the type is NONE.
     # This covers VirtualValues which always report NONE.
@@ -1014,58 +977,79 @@ def write_flex_value(value: RegistryValue) -> str:
             vtype = RegistryValueType.MULTI_SZ
         elif isinstance(data, bytes):
             vtype = RegistryValueType.BINARY
+
     if vtype == RegistryValueType.SZ:
-        return f'{name_str}="{_escape_reg_str(str(data))}"'
+        return f'{name_str}="{_escape_flex_str(str(data))}"'
 
     if vtype == RegistryValueType.DWORD:
         return f"{name_str}=dword:{int(data):08x}"
 
     if vtype == RegistryValueType.BINARY:
-        return _format_hex_data(data if isinstance(data, bytes) else b"", f"{name_str}=hex:")
+        return _format_flex_hex(data if isinstance(data, bytes) else b"", f"{name_str}=hex:")
 
     if vtype == RegistryValueType.EXPAND_SZ:
         raw = data.encode("utf-16-le") + b"\x00\x00" if isinstance(data, str) else (data or b"")
-        return _format_hex_data(raw, f"{name_str}=hex(2):")
+        return _format_flex_hex(raw, f"{name_str}=hex(2):")
 
     if vtype == RegistryValueType.MULTI_SZ:
         if isinstance(data, list):
             joined = "".join(f"{s}\x00" for s in data) + "\x00"
             raw = joined.encode("utf-16-le")
         else:
             raw = data if isinstance(data, bytes) else b""
-        return _format_hex_data(raw, f"{name_str}=hex(7):")
+        return _format_flex_hex(raw, f"{name_str}=hex(7):")
 
     if vtype == RegistryValueType.DWORD_BIG_ENDIAN:
         raw = int(data).to_bytes(4, "big") if isinstance(data, int) else (data or b"")
-        return _format_hex_data(raw, f"{name_str}=hex(5):")
+        return _format_flex_hex(raw, f"{name_str}=hex(5):")
 
     if vtype == RegistryValueType.QWORD:
         raw = int(data).to_bytes(8, "little") if isinstance(data, int) else (data or b"")
-        return _format_hex_data(raw, f"{name_str}=hex(b):")
+        return _format_flex_hex(raw, f"{name_str}=hex(b):")
 
     if vtype == RegistryValueType.NONE:
         raw = data if isinstance(data, bytes) else b""
-        return _format_hex_data(raw, f"{name_str}=hex(0):")
+        return _format_flex_hex(raw, f"{name_str}=hex(0):")
 
     # Fallback for LINK, RESOURCE_LIST, FULL_RESOURCE_DESCRIPTOR, RESOURCE_REQUIREMENTS_LIST, etc.
     raw = data if isinstance(data, bytes) else b""
-    return _format_hex_data(raw, f"{name_str}=hex({int(vtype):x}):")
+    return _format_flex_hex(raw, f"{name_str}=hex({int(vtype):x}):")
 
 
-def write_reg_key(fh: TextIO, key: RegistryKey, path: str) -> None:
-    """Recursively write a registry key and its subkeys in .reg format.
+def format_flex_key(
+    key: RegistryKey,
+    path: str | None = None,
+) -> Iterator[str]:
+    """Recursively format a registry key and its subkeys in ``.reg`` format.
 
     Args:
-        fh: A file-like object opened in text mode.
-        key: The registry key to write.
-        path: The full registry path for this key
-            (e.g. ``HKEY_LOCAL_MACHINE\\SOFTWARE\\Test``).
+        key: The registry key to format.
+        path: The full registry path for this key (e.g. ``HKEY_LOCAL_MACHINE\\SOFTWARE\\Test``).
+              If ``None``, ``key.path`` is used.
     """
-    fh.write(f"[{path}]\n")
-    fh.writelines(f"{write_flex_value(reg_value)}\n" for reg_value in key.values())
-    fh.write("\n")
+    yield f"[{path or key.path}]"
+    for value in key.values():
+        yield f"{format_flex_value(value)}"
+    yield ""
     for subkey in key.subkeys():
-        write_reg_key(fh, subkey, f"{path}\\{subkey.name}")
+        yield from format_flex_key(subkey, f"{path or key.path}\\{subkey.name}")
+
+
+def format_flex_header(keys: list[str] | None = None) -> Iterator[str]:
+    """Format the header for a ``.reg`` file.
+
+    Args:
+        keys: Optional list of registry paths that will be exported in the file.
+              If provided, each path is included as a comment line in the header.
+    """
+    yield f"{REGFLEX_HEADER}"
+    yield ""
+    yield "; Generated by dissect.target"
+    if keys:
+        yield "; Registry paths exported:"
+        for key in keys:
+            yield f";   {key}"
+    yield ""
 
 
 def has_glob_magic(pattern: str) -> bool:
@@ -1080,6 +1064,33 @@ def has_glob_magic(pattern: str) -> bool:
     return RE_GLOB_MAGIC.search(pattern) is not None
 
 
+def split_key_path(path: str) -> tuple[str, str]:
+    """Split a registry key path into the canonical ``HKEY_*`` hive name and the rest of the path.
+
+    Args:
+        path: A registry path starting with a shortname such as ``HKLM`` or ``HKCU``.
+
+    Returns:
+        A tuple containing the canonical ``HKEY_*`` hive name and the rest of the path.
+    """
+    hive, _, rest = path.partition("\\")
+    expanded = SHORTNAMES.get(hive.upper(), hive.upper())
+    return expanded, rest
+
+
+def expand_key_path(path: str) -> str:
+    """Expand a registry shortname prefix to the canonical ``HKEY_*`` name.
+
+    Args:
+        path: A registry path starting with a shortname such as ``HKLM`` or ``HKCU``.
+
+    Returns:
+        The path with the leading component expanded to its full name.
+    """
+    hive, rest = split_key_path(path)
+    return f"{hive}\\{rest}"
+
+
 def glob_split(pattern: str) -> tuple[str]:
     """Split a key path with glob patterns on the first key path part with glob patterns.