png-zip/unpack.py

#!/usr/bin/env python3
"""Extract file(s) hidden inside a PNG image.

Supports three embedding modes (auto-detected):
    append — data after IEND
    chunk  — data in private pnZp PNG chunk
    lsb    — data in least-significant bits of pixels

Usage:
    python3 unpack.py packed.png [-o output_dir] [-a]
"""
import argparse
import os
import struct
import sys
import zlib

MAGIC = b"PNGZIP\x00\x01"
PNG_SIG = b"\x89PNG\r\n\x1a\n"


# ── Shared helpers ────────────────────────────────────────────────

def find_iend(data: bytes) -> int:
    """Return offset of the first byte AFTER the IEND chunk, or -1."""
    idx = data.rfind(b"IEND")
    if idx == -1:
        return -1
    chunk_start = idx - 4
    length = struct.unpack(">I", data[chunk_start:chunk_start + 4])[0]
    return chunk_start + 4 + 4 + length + 4


def parse_chunks(data: bytes):
    """Parse PNG data into list of (chunk_type, chunk_data)."""
    pos = 8  # skip PNG signature
    chunks = []
    while pos + 8 <= len(data):
        length = struct.unpack(">I", data[pos:pos + 4])[0]
        chunk_type = data[pos + 4:pos + 8]
        chunk_data = data[pos + 8:pos + 8 + length]
        pos += 4 + 4 + length + 4
        chunks.append((chunk_type, chunk_data))
    return chunks


def parse_payload_blocks(data: bytes):
    """Yield (filename, raw_bytes) for each payload block in data."""
    offset = 0
    while True:
        idx = data.find(MAGIC, offset)
        if idx == -1:
            break
        pos = idx + len(MAGIC)

        if pos + 2 > len(data):
            break
        name_len = struct.unpack(">H", data[pos:pos + 2])[0]
        pos += 2

        if pos + name_len > len(data):
            break
        filename = data[pos:pos + name_len].decode("utf-8")
        pos += name_len

        if pos + 8 > len(data):
            break
        payload_len = struct.unpack(">Q", data[pos:pos + 8])[0]
        pos += 8

        if pos + payload_len > len(data):
            break
        compressed = data[pos:pos + payload_len]
        pos += payload_len

        if pos + 4 > len(data):
            break
        stored_crc = struct.unpack(">I", data[pos:pos + 4])[0]
        pos += 4

        raw = zlib.decompress(compressed)
        actual_crc = zlib.crc32(raw) & 0xFFFFFFFF
        if actual_crc != stored_crc:
            print(f"Warning: CRC mismatch for {filename}", file=sys.stderr)

        yield filename, raw
        offset = pos


# ── Append mode ───────────────────────────────────────────────────

def find_append_payloads(data: bytes):
    """Find payloads appended after IEND."""
    iend = find_iend(data)
    if iend == -1 or iend >= len(data):
        return []
    return list(parse_payload_blocks(data[iend:]))


# ── Chunk mode ────────────────────────────────────────────────────

def find_chunk_payloads(data: bytes):
    """Find payloads stored in pnZp chunks."""
    if data[:8] != PNG_SIG:
        return []
    chunks = parse_chunks(data)
    pnzp = [(cd[0], cd[1:]) for ct, cd in chunks if ct == b"pnZp" and len(cd) > 0]
    if not pnzp:
        return []
    pnzp.sort(key=lambda x: x[0])
    combined = b"".join(frag for _, frag in pnzp)
    return list(parse_payload_blocks(combined))


# ── LSB mode ──────────────────────────────────────────────────────

def _paeth(a, b, c):
    p = a + b - c
    pa, pb, pc = abs(p - a), abs(p - b), abs(p - c)
    if pa <= pb and pa <= pc:
        return a
    return b if pb <= pc else c


def _unfilter_row(ftype, row, prev, bpp):
    out = bytearray(len(row))
    for i in range(len(row)):
        x = row[i]
        a = out[i - bpp] if i >= bpp else 0
        b = prev[i]
        c = prev[i - bpp] if i >= bpp else 0
        if ftype == 0:
            out[i] = x
        elif ftype == 1:
            out[i] = (x + a) & 0xFF
        elif ftype == 2:
            out[i] = (x + b) & 0xFF
        elif ftype == 3:
            out[i] = (x + (a + b) // 2) & 0xFF
        elif ftype == 4:
            out[i] = (x + _paeth(a, b, c)) & 0xFF
        else:
            raise ValueError(f"Unknown PNG filter type {ftype}")
    return out


def find_lsb_payloads(data: bytes):
    """Find payloads encoded in LSB of pixel data."""
    if data[:8] != PNG_SIG:
        return []
    try:
        chunks = parse_chunks(data)
    except Exception:
        return []

    # Parse IHDR
    ihdr = None
    for ct, cd in chunks:
        if ct == b"IHDR":
            ihdr = cd
            break
    if ihdr is None or len(ihdr) < 13:
        return []

    width, height = struct.unpack(">II", ihdr[:8])
    bit_depth, color_type = ihdr[8], ihdr[9]
    interlace = ihdr[12]
    if bit_depth != 8 or color_type not in (2, 6) or interlace != 0:
        return []
    channels = 3 if color_type == 2 else 4
    bpp = channels

    # Decompress IDAT
    idat_data = b"".join(cd for ct, cd in chunks if ct == b"IDAT")
    if not idat_data:
        return []
    try:
        raw = zlib.decompress(idat_data)
    except zlib.error:
        return []

    # Unfilter
    stride = width * bpp
    pixels = bytearray()
    prev = bytearray(stride)
    pos = 0
    for _ in range(height):
        if pos >= len(raw):
            return []
        ftype = raw[pos]; pos += 1
        if pos + stride > len(raw):
            return []
        row = raw[pos:pos + stride]; pos += stride
        prev = _unfilter_row(ftype, row, prev, bpp)
        pixels.extend(prev)

    # Quick check: extract MAGIC-length bytes first
    total_pixels = len(pixels)
    magic_len = len(MAGIC)
    if total_pixels < magic_len * 8:
        return []

    quick = bytearray(magic_len)
    for i in range(magic_len):
        b = 0
        base = i * 8
        for bit in range(8):
            b = (b << 1) | (pixels[base + bit] & 1)
        quick[i] = b

    if quick != bytearray(MAGIC):
        return []

    # MAGIC found — extract all available bytes
    total_bytes = total_pixels // 8
    extracted = bytearray(total_bytes)
    for i in range(total_bytes):
        b = 0
        base = i * 8
        for bit in range(8):
            b = (b << 1) | (pixels[base + bit] & 1)
        extracted[i] = b

    return list(parse_payload_blocks(bytes(extracted)))


# ── Auto-detect all modes ────────────────────────────────────────

def find_all_payloads(data: bytes):
    """Try all detection methods: append → chunk → lsb."""
    results = find_append_payloads(data)
    if results:
        return results
    results = find_chunk_payloads(data)
    if results:
        return results
    results = find_lsb_payloads(data)
    if results:
        return results
    return []


# ── Main ──────────────────────────────────────────────────────────

def main():
    ap = argparse.ArgumentParser(description="Extract file from PNG")
    ap.add_argument("image", help="PNG with hidden payload")
    ap.add_argument("-o", "--outdir", default=".", help="Output directory")
    ap.add_argument("-a", "--all", action="store_true",
                    help="Extract all payloads (including embedded unpack.py)")
    ap.add_argument("-l", "--list", action="store_true",
                    help="List payloads without extracting")
    args = ap.parse_args()

    with open(args.image, "rb") as f:
        data = f.read()

    payloads = list(find_all_payloads(data))
    if not payloads:
        print("No payload found in this PNG.", file=sys.stderr)
        sys.exit(3)

    if args.list:
        for name, raw in payloads:
            print(f"  {name} ({len(raw)} bytes)")
        return

    # By default extract only the last payload (the main one).
    # With --all, extract everything.
    to_extract = payloads if args.all else [payloads[-1]]

    os.makedirs(args.outdir, exist_ok=True)
    for filename, raw in to_extract:
        safe_name = os.path.basename(filename)  # path traversal protection
        out_path = os.path.join(args.outdir, safe_name)
        with open(out_path, "wb") as f:
            f.write(raw)
        print(f"Extracted: {out_path} ({len(raw)} bytes)")


if __name__ == "__main__":
    main()