#!/usr/bin/env python3
"""Carve common file signatures from an authorized binary image.

This starter supports bounded carving for PNG, JPG, PDF, ZIP, and PE files. It is
intended for lab, CTF, or professional evidence workflows where carving is in scope.
"""

from __future__ import annotations

import argparse
from dataclasses import dataclass
from pathlib import Path


@dataclass(frozen=True)
class Signature:
    name: str
    extension: str
    header: bytes
    footer: bytes | None
    max_size: int


SIGNATURES = [
    Signature("png", "png", b"\x89PNG\r\n\x1a\n", b"IEND\xaeB`\x82", 25 * 1024 * 1024),
    Signature("jpg", "jpg", b"\xff\xd8\xff", b"\xff\xd9", 25 * 1024 * 1024),
    Signature("pdf", "pdf", b"%PDF-", b"%%EOF", 50 * 1024 * 1024),
    Signature("zip", "zip", b"PK\x03\x04", b"PK\x05\x06", 100 * 1024 * 1024),
    Signature("pe", "exe", b"MZ", None, 50 * 1024 * 1024),
]


def carve(blob: bytes, output_dir: Path) -> int:
    output_dir.mkdir(parents=True, exist_ok=True)
    count = 0
    for signature in SIGNATURES:
        offset = 0
        while True:
            start = blob.find(signature.header, offset)
            if start == -1:
                break
            if signature.footer:
                end = blob.find(signature.footer, start + len(signature.header))
                if end == -1:
                    offset = start + 1
                    continue
                end += len(signature.footer)
            else:
                end = min(start + signature.max_size, len(blob))
            if end - start <= signature.max_size:
                count += 1
                out = output_dir / f"carved_{count:04d}_{start:x}.{signature.extension}"
                out.write_bytes(blob[start:end])
                print(f"{signature.name}: offset=0x{start:x} size={end - start} -> {out}")
            offset = start + len(signature.header)
    return count


def main() -> int:
    parser = argparse.ArgumentParser(description="Carve known file types from an evidence image.")
    parser.add_argument("image", type=Path)
    parser.add_argument("--output-dir", type=Path, default=Path("carved_files"))
    args = parser.parse_args()
    count = carve(args.image.read_bytes(), args.output_dir)
    print(f"Carved {count} candidate files")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())