Add write_sid function (#111)

Co-authored-by: Schamper <1254028+Schamper@users.noreply.github.com>

Author: joost-j

dissect/util/sid.py

@@ -48,3 +48,34 @@ def read_sid(fh: BinaryIO | bytes, endian: str = "<", swap_last: bool = False) -
     ]
     sid_elements.extend(map(str, sub_authorities))
     return "-".join(sid_elements)
+
+
+def write_sid(sid: str, endian: str = "<", swap_last: bool = False) -> bytes:
+    """Write a Windows SID string to bytes.
+
+    Args:
+        sid: SID in the form ``S-Revision-Authority-SubAuth1-...``.
+        endian: Optional endianness for reading the sub authorities.
+        swap_last: Optional flag for swapping the endianess of the _last_ sub authority entry.
+    """
+    if not sid:
+        return b""
+
+    parts = sid.split("-")
+    if len(parts) < 3 or parts[0].upper() != "S":
+        raise ValueError("Invalid SID string format: insufficient parts")
+
+    revision = int(parts[1]).to_bytes(1, "little")
+    authority = int(parts[2]).to_bytes(6, "big")
+    sub_authorities = [int(x) for x in parts[3:]]
+
+    header = revision + len(sub_authorities).to_bytes(1, "little") + authority
+
+    if not sub_authorities:
+        return header
+
+    sub_bytes = bytearray(struct.pack(f"{endian}{len(sub_authorities)}I", *sub_authorities))
+    if swap_last:
+        sub_bytes[-4:] = sub_bytes[-4:][::-1]
+
+    return header + bytes(sub_bytes)

tests/test_sid.py

@@ -84,9 +84,13 @@ def id_fn(val: bytes | str) -> str:
     ],
     ids=id_fn,
 )
-def test_read_sid(binary_sid: bytes | BinaryIO, endian: str, swap_last: bool, readable_sid: str) -> None:
+def test_read_write_sid(binary_sid: bytes | BinaryIO, endian: str, swap_last: bool, readable_sid: str) -> None:
     assert readable_sid == sid.read_sid(binary_sid, endian, swap_last)
 
+    if isinstance(binary_sid, io.BytesIO):
+        binary_sid = binary_sid.getvalue()
+    assert binary_sid == sid.write_sid(readable_sid, endian, swap_last)
+
 
 @pytest.mark.benchmark
 @pytest.mark.parametrize(
@@ -105,3 +109,22 @@ def test_read_sid(binary_sid: bytes | BinaryIO, endian: str, swap_last: bool, re
 )
 def test_read_sid_benchmark(benchmark: BenchmarkFixture, binary_sid: bytes, swap_last: bool) -> None:
     benchmark(sid.read_sid, binary_sid, "<", swap_last)
+
+
+@pytest.mark.benchmark
+@pytest.mark.parametrize(
+    ("readable_sid", "swap_last"),
+    [
+        (
+            "S-1-5-21-123456789-268435456-500",
+            False,
+        ),
+        (
+            "S-1-5-21-123456789-268435456-500",
+            True,
+        ),
+    ],
+    ids=lambda x: x if isinstance(x, str) else str(x),
+)
+def test_write_sid_benchmark(benchmark: BenchmarkFixture, readable_sid: str, swap_last: bool) -> None:
+    benchmark(sid.write_sid, readable_sid, "<", swap_last)