diff --git a/dissect/evidence/asdf/asdf.py b/dissect/evidence/asdf/asdf.py index 7622316..3221f99 100644 --- a/dissect/evidence/asdf/asdf.py +++ b/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,6 +71,7 @@ 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 @@ -72,7 +79,12 @@ def __init__( 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) diff --git a/dissect/evidence/asdf/c_asdf.py b/dissect/evidence/asdf/c_asdf.py index 8da56ce..c6845ec 100644 --- a/dissect/evidence/asdf/c_asdf.py +++ b/dissect/evidence/asdf/c_asdf.py @@ -3,8 +3,8 @@ from dissect.cstruct import cstruct asdf_def = """ -flag FILE_FLAG : uint32 { - SHA256 = 0x01, +flag FILE_FLAG : uint24 { + CHECKSUM = 0x01 }; flag BLOCK_FLAG : uint8 { @@ -12,9 +12,17 @@ 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 }; """ diff --git a/dissect/evidence/asdf/c_asdf.pyi b/dissect/evidence/asdf/c_asdf.pyi index 58c6f12..cdfc77b 100644 --- a/dissect/evidence/asdf/c_asdf.pyi +++ b/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, /): ... diff --git a/dissect/evidence/tools/asdf/verify.py b/dissect/evidence/tools/asdf/verify.py index 6237613..eeef73c 100644 --- a/dissect/evidence/tools/asdf/verify.py +++ b/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") diff --git a/tests/tools/test_asdf_verify.py b/tests/tools/test_asdf_verify.py new file mode 100644 index 0000000..4fdfc8d --- /dev/null +++ b/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