Add support for using multiple checksum algorithms in AsdfWriter

dissect/evidence/asdf/asdf.py

@@ -42,6 +42,12 @@
 SPARSE_BYTES = b"\xa5\xdf"
 
 
+CHECKSUM_MAPPING: dict[int, str] = {
+    c_asdf.CHECKSUM.SHA256: 32,
+    c_asdf.CHECKSUM.SHA512: 64,
+}
+
+
 class AsdfWriter(io.RawIOBase):
     """ASDF file writer.
 
@@ -65,14 +71,20 @@ def __init__(
         guid: uuid.UUID | None = None,
         compress: bool = False,
         block_crc: bool = True,
+        checksum_algorithm: c_asdf.CHECKSUM = c_asdf.CHECKSUM.SHA256,
     ):
         self._fh = fh
         self.fh = self._fh
 
         if compress:
             self.fh = gzip.GzipFile(fileobj=self.fh, mode="wb")
 
-        self.fh = HashedStream(self.fh)
+        if checksum_algorithm not in CHECKSUM_MAPPING:
+            raise ValueError("Unsupported hashing algorithm used")
+
+        self.checksum = c_asdf.CHECKSUM(checksum_algorithm)
+        self.fh = HashedStream(self.fh, alg=self.checksum.name)
+
         self.guid = guid or uuid.uuid4()
 
         # Options
@@ -213,7 +225,8 @@ def _write_header(self) -> None:
         """Write the ASDF header to the destination file-like object."""
         header = c_asdf.header(
             magic=FILE_MAGIC,
-            flags=c_asdf.FILE_FLAG.SHA256,  # Currently the only option
+            flags=c_asdf.FILE_FLAG.CHECKSUM,  # Currently the only option
+            checksum=self.checksum,
             version=VERSION,
             timestamp=ts.unix_now(),
             guid=self.guid.bytes_le,
@@ -300,10 +313,10 @@ def _write_table(self) -> None:
 
     def _write_footer(self) -> None:
         """Write the ASDF footer to the destination file-like object."""
+        self.fh.write(self.fh.digest())
         footer = c_asdf.footer(
             magic=FOOTER_MAGIC,
             table_offset=self._table_offset,
-            sha256=self.fh.digest(),
         )
         footer.write(self.fh)
 

dissect/evidence/asdf/c_asdf.py

@@ -3,18 +3,26 @@
 from dissect.cstruct import cstruct
 
 asdf_def = """
-flag FILE_FLAG : uint32 {
-    SHA256      = 0x01,
+flag FILE_FLAG : uint24 {
+    CHECKSUM      = 0x01
 };
 
 flag BLOCK_FLAG : uint8 {
     CRC32       = 0x01,
     COMPRESS    = 0x02,
 };
 
+enum CHECKSUM : uint8 {
+    SHA256 = 0x1,
+    SHA512 = 0x2,
+    NONE = 0xFF,
+};
+
+
 struct header {
     char        magic[4];       // File magic, must be "ASDF"
     FILE_FLAG   flags;          // File flags
+    CHECKSUM    checksum;       // Checksum used for the flags
     uint8       version;        // File version
     char        reserved1[7];   // Reserved
     uint64      timestamp;      // Creation timestamp of the file
@@ -45,7 +53,6 @@
     char        magic[4];       // Footer magic, must be "FT\\xa5\\xdf"
     char        reserved[4];    // Reserved
     uint64      table_offset;   // Offset in file to start of block table
-    char        sha256[32];     // SHA256 of this file up until this hash
 };
 """
 

dissect/evidence/asdf/c_asdf.pyi

@@ -1,19 +1,25 @@
 # Generated by cstruct-stubgen
-from typing import BinaryIO, TypeAlias, overload
+from typing import BinaryIO, Literal, TypeAlias, overload
 
 import dissect.cstruct as __cs__
 
 class _c_asdf(__cs__.cstruct):
     class FILE_FLAG(__cs__.Flag):
-        SHA256 = ...
+        CHECKSUM = ...
 
     class BLOCK_FLAG(__cs__.Flag):
         CRC32 = ...
         COMPRESS = ...
 
+    class CHECKSUM(__cs__.Enum):
+        SHA256 = ...
+        SHA512 = ...
+        NONE = ...
+
     class header(__cs__.Structure):
         magic: __cs__.CharArray
         flags: _c_asdf.FILE_FLAG
+        checksum: _c_asdf.CHECKSUM
         version: _c_asdf.uint8
         reserved1: __cs__.CharArray
         timestamp: _c_asdf.uint64
@@ -24,6 +30,7 @@ class _c_asdf(__cs__.cstruct):
             self,
             magic: __cs__.CharArray | None = ...,
             flags: _c_asdf.FILE_FLAG | None = ...,
+            checksum: _c_asdf.CHECKSUM | None = ...,
             version: _c_asdf.uint8 | None = ...,
             reserved1: __cs__.CharArray | None = ...,
             timestamp: _c_asdf.uint64 | None = ...,
@@ -79,14 +86,12 @@ class _c_asdf(__cs__.cstruct):
         magic: __cs__.CharArray
         reserved: __cs__.CharArray
         table_offset: _c_asdf.uint64
-        sha256: __cs__.CharArray
         @overload
         def __init__(
             self,
             magic: __cs__.CharArray | None = ...,
             reserved: __cs__.CharArray | None = ...,
             table_offset: _c_asdf.uint64 | None = ...,
-            sha256: __cs__.CharArray | None = ...,
         ): ...
         @overload
         def __init__(self, fh: bytes | memoryview | bytearray | BinaryIO, /): ...

dissect/evidence/tools/asdf/verify.py

@@ -63,6 +63,7 @@ def main() -> int:
     with Path(args.file).open("rb") as fh:
         header = None
         footer = None
+        checksum = None
         footer_offset = 0
 
         with status("Checking header", args.verbose):
@@ -71,9 +72,12 @@ def main() -> int:
                 print("[!] Invalid header magic")
                 return 1
 
+        checksum_size = asdf.CHECKSUM_MAPPING.get(header.checksum, 0)
+
         with status("Checking footer", args.verbose):
-            fh.seek(-len(asdf.c_asdf.footer), io.SEEK_END)
+            fh.seek(-(len(asdf.c_asdf.footer) + checksum_size), io.SEEK_END)
             footer_offset = fh.tell()
+            checksum = fh.read(checksum_size)
             footer = asdf.c_asdf.footer(fh)
             if footer.magic != asdf.FOOTER_MAGIC:
                 footer = None
@@ -83,9 +87,10 @@ def main() -> int:
         if not args.skip_hash and footer:
             with status("Checking file hash", args.verbose):
                 hashstream = RangeStream(fh, 0, footer_offset)
-                res = hash_fileobj(hashstream)
-                if res != footer.sha256:
-                    print(f"[!] File hash doesn't match. Expected {footer.sha256.hex()}, got {res.hex()}")
+                res = hash_fileobj(hashstream, alg=header.checksum.name.lower())
+
+                if res != checksum:
+                    print(f"[!] File hash doesn't match. Expected {checksum.hex()}, got {res.hex()}")
                     return 1
         else:
             print("[@] Skipping file hash")

tests/tools/test_asdf_verify.py

@@ -0,0 +1,36 @@
+from pathlib import Path
+
+import pytest
+
+from dissect.evidence.asdf import AsdfWriter
+from dissect.evidence.asdf.asdf import CHECKSUM_MAPPING
+from dissect.evidence.tools.asdf.verify import main as asdf_verify
+
+
+@pytest.fixture(params=CHECKSUM_MAPPING.keys())
+def asdf_file(tmp_path: Path, request: pytest.FixtureRequest) -> Path:
+    asdf_file = tmp_path.joinpath("asdf.asdf")
+    fh = asdf_file.open("wb")
+    with AsdfWriter(fh, checksum_algorithm=request.param) as asdf_writer:
+        asdf_writer.add_bytes(b"\x00" * 0x1000, idx=0, base=0)
+        asdf_writer.add_bytes(b"\x02" * 0x1000, idx=0, base=0x4000)
+        asdf_writer.add_bytes(b"\x04" * 0x1000, idx=0, base=0x8000)
+        asdf_writer.add_bytes(b"\x06" * 0x1000, idx=0, base=0x10000)
+        asdf_writer.add_bytes(b"\xff" * 0x1000, idx=0, base=0x14000)
+
+        asdf_writer.add_bytes(b"\x08" * 0x1000, idx=1, base=0x2000)
+        asdf_writer.add_bytes(b"\x10" * 0x1000, idx=1, base=0x5000)
+        asdf_writer.add_bytes(b"\x12" * 0x1000, idx=1, base=0x8000)
+        asdf_writer.add_bytes(b"\x14" * 0x1000, idx=1, base=0xB000)
+        asdf_writer.add_bytes(b"\xff" * 0x1000, idx=1, base=0xE000)
+    return asdf_file
+
+
+def test_asdf_verify(asdf_file: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    with monkeypatch.context() as m:
+        m.setattr(
+            "sys.argv",
+            ["asdf-verify", str(asdf_file)],
+        )
+
+        assert asdf_verify() == 0