refactor: clean up comments and docstrings

pack.py

@@ -42,10 +42,7 @@ def make_payload_block(filename: str, raw: bytes) -> bytes:
     return header + compressed + crc
 
 
-# ── PNG chunk helpers ─────────────────────────────────────────────
-
 def parse_chunks(data: bytes):
-    """Parse PNG data into list of (chunk_type, chunk_data)."""
     if data[:8] != PNG_SIG:
         raise ValueError("Not a valid PNG")
     pos = 8
@@ -60,22 +57,18 @@ def parse_chunks(data: bytes):
 
 
 def make_png_chunk(chunk_type: bytes, chunk_data: bytes) -> bytes:
-    """Build a valid PNG chunk: length + type + data + CRC32."""
     length = struct.pack(">I", len(chunk_data))
     crc = struct.pack(">I", zlib.crc32(chunk_type + chunk_data) & 0xFFFFFFFF)
     return length + chunk_type + chunk_data + crc
 
 
 def assemble_png(chunks) -> bytes:
-    """Reassemble PNG from list of (type, data) tuples."""
     result = bytearray(PNG_SIG)
     for chunk_type, chunk_data in chunks:
         result += make_png_chunk(chunk_type, chunk_data)
     return bytes(result)
 
 
-# ── Chunk mode ────────────────────────────────────────────────────
-
 def pack_chunk(cover_bytes: bytes, payload_block: bytes) -> bytes:
     chunks = parse_chunks(cover_bytes)
 
@@ -95,8 +88,6 @@ def pack_chunk(cover_bytes: bytes, payload_block: bytes) -> bytes:
     return assemble_png(result)
 
 
-# ── LSB mode ──────────────────────────────────────────────────────
-
 def _paeth(a, b, c):
     p = a + b - c
     pa, pb, pc = abs(p - a), abs(p - b), abs(p - c)
@@ -130,7 +121,6 @@ def _unfilter_row(ftype, row, prev, bpp):
 def pack_lsb(cover_bytes: bytes, payload_block: bytes) -> bytes:
     chunks = parse_chunks(cover_bytes)
 
-    # ── IHDR ──
     ihdr = next(cd for ct, cd in chunks if ct == b"IHDR")
     width, height = struct.unpack(">II", ihdr[:8])
     bit_depth, color_type = ihdr[8], ihdr[9]
@@ -144,11 +134,9 @@ def pack_lsb(cover_bytes: bytes, payload_block: bytes) -> bytes:
     channels = 3 if color_type == 2 else 4
     bpp = channels
 
-    # ── Collect & decompress IDAT ──
     idat_data = b"".join(cd for ct, cd in chunks if ct == b"IDAT")
     raw = zlib.decompress(idat_data)
 
-    # ── Unfilter ──
     stride = width * bpp
     pixels = bytearray()
     prev = bytearray(stride)
@@ -159,7 +147,6 @@ def pack_lsb(cover_bytes: bytes, payload_block: bytes) -> bytes:
         prev = _unfilter_row(ftype, row, prev, bpp)
         pixels.extend(prev)
 
-    # ── Capacity check ──
     capacity = len(pixels)  # 1 bit per byte
     needed = len(payload_block) * 8
     if needed > capacity:
@@ -169,21 +156,19 @@ def pack_lsb(cover_bytes: bytes, payload_block: bytes) -> bytes:
               file=sys.stderr)
         sys.exit(1)
 
-    # ── Embed bits into LSB ──
     bit_idx = 0
     for byte in payload_block:
         for shift in range(7, -1, -1):
             pixels[bit_idx] = (pixels[bit_idx] & 0xFE) | ((byte >> shift) & 1)
             bit_idx += 1
 
-    # ── Re-filter (type 0 = None) & compress ──
+    # re-filter with type 0 (None) and compress
     filtered = bytearray()
     for y in range(height):
         filtered.append(0)
         filtered.extend(pixels[y * stride:(y + 1) * stride])
     compressed = zlib.compress(bytes(filtered), 6)
 
-    # ── Reassemble PNG ──
     result = []
     idat_done = False
     for ct, cd in chunks:
@@ -196,10 +181,7 @@ def pack_lsb(cover_bytes: bytes, payload_block: bytes) -> bytes:
     return assemble_png(result)
 
 
-# ── Clean mode ────────────────────────────────────────────────────
-
 def clean_png(data: bytes) -> bytes:
-    """Strip all embedded payloads (append + chunk) and bootstrap metadata from a PNG."""
     marker = BOOTSTRAP_MARKER.encode("latin-1")
     chunks = parse_chunks(data)
     cleaned = []
@@ -217,7 +199,6 @@ def clean_png(data: bytes) -> bytes:
 
 
 def build_exif_user_comment(text: str) -> bytes:
-    """Build minimal eXIf chunk data with UserComment field."""
     user_comment = b"ASCII\x00\x00\x00" + text.encode("ascii")
     uc_len = len(user_comment)
     # TIFF header (8) + IFD0: count(2) + entry(12) + next(4) = 18
@@ -242,7 +223,6 @@ def build_exif_user_comment(text: str) -> bytes:
 
 
 def _make_py_b64(fname: str, last: bool = False) -> str:
-    """Generate base64-encoded Python bootstrap script for CMD."""
     import base64
     find = "d.rfind(m)" if last else "d.find(m)"
     script = (
@@ -259,7 +239,6 @@ def _make_py_b64(fname: str, last: bool = False) -> str:
     return base64.b64encode(script.encode()).decode()
 
 
-# Common PowerShell payload-extraction tail (after $i is set to MAGIC position)
 _PS_EXTRACT_TAIL = (
     "$p=$i+8;"
     "$nl=$d[$p]*256+$d[$p+1];$p+=2;"
@@ -274,14 +253,12 @@ _PS_EXTRACT_TAIL = (
     'Write-Host "Extracted: $n"'
 )
 
-# PowerShell: find FIRST MAGIC (for bootstrap — extracts unpacker)
 _PS_FIND_FIRST = (
     "$s='PNGZIP';$i=-1;"
     "do{$i=$t.IndexOf($s,$i+1,[StringComparison]::Ordinal)}"
     "while($i-ge0-and($d[$i+6]-ne0-or$d[$i+7]-ne1));"
 )
 
-# PowerShell: find LAST MAGIC (for direct extraction — extracts payload)
 _PS_FIND_LAST = (
     "$s='PNGZIP';$i=$d.Length;"
     "do{$i=$t.LastIndexOf($s,$i-1,[StringComparison]::Ordinal)}"
@@ -295,7 +272,6 @@ _PS_READ_HEAD = (
 
 
 def _make_ps_encoded(fname: str, last: bool = False) -> str:
-    """Generate base64-encoded PowerShell bootstrap script for CMD."""
     import base64
     find = _PS_FIND_LAST if last else _PS_FIND_FIRST
     script = _PS_READ_HEAD.replace("IMG", fname) + find + _PS_EXTRACT_TAIL
@@ -303,7 +279,6 @@ def _make_ps_encoded(fname: str, last: bool = False) -> str:
 
 
 def _py_oneliner(fname: str, last: bool = False) -> str:
-    """Python one-liner for shell (Linux/macOS/PowerShell)."""
     find = "d.rfind" if last else "d.find"
     return (
         "python3 -c \"d=open('IMG','rb').read();"
@@ -318,13 +293,11 @@ def _py_oneliner(fname: str, last: bool = False) -> str:
 
 
 def _ps_oneliner(fname: str, last: bool = False) -> str:
-    """PowerShell one-liner."""
     find = _PS_FIND_LAST if last else _PS_FIND_FIRST
     return (_PS_READ_HEAD + find + _PS_EXTRACT_TAIL).replace("IMG", fname)
 
 
 def _add_oneliner_block(parts: list, fname: str, last: bool, label: str):
-    """Append a full set of 4 one-liners to parts list."""
     parts.append("")
     parts.append(f"=== {label} ===")
     parts.append("")
@@ -344,7 +317,6 @@ def _add_oneliner_block(parts: list, fname: str, last: bool, label: str):
 
 
 def build_bootstrap_comment(output_name: str, has_self_extract: bool = True) -> bytes:
-    """Build tEXt chunk data with all bootstrap one-liners for PNG metadata."""
     fname = os.path.basename(output_name)
     parts = [BOOTSTRAP_MARKER]
     if has_self_extract:
@@ -356,8 +328,6 @@ def build_bootstrap_comment(output_name: str, has_self_extract: bool = True) ->
     return b"Comment\x00" + text.encode("latin-1")
 
 
-# ── Main ──────────────────────────────────────────────────────────
-
 def main():
     ap = argparse.ArgumentParser(description="Pack file into PNG")
     ap.add_argument("cover", help="Cover PNG image")
@@ -380,7 +350,6 @@ def main():
         print("Error: not a valid PNG", file=sys.stderr)
         sys.exit(1)
 
-    # ── Clean mode ──
     if args.clean:
         result = clean_png(cover)
         out = args.output if args.payload is None else args.output
@@ -391,7 +360,6 @@ def main():
         print(f"Cleaned: {out} ({after_kb:.1f} KB, was {before_kb:.1f} KB)")
         return
 
-    # ── Pack mode — payload required ──
     if args.payload is None:
         print("Error: payload argument is required (unless --clean)", file=sys.stderr)
         sys.exit(1)

unpack.py

@@ -19,8 +19,6 @@ 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")
@@ -87,20 +85,14 @@ def parse_payload_blocks(data: bytes):
         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)
@@ -112,8 +104,6 @@ def find_chunk_payloads(data: bytes):
     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)
@@ -145,7 +135,6 @@ def _unfilter_row(ftype, row, prev, bpp):
 
 
 def find_lsb_payloads(data: bytes):
-    """Find payloads encoded in LSB of pixel data."""
     if data[:8] != PNG_SIG:
         return []
     try:
@@ -224,10 +213,7 @@ def find_lsb_payloads(data: bytes):
     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
@@ -240,8 +226,6 @@ def find_all_payloads(data: bytes):
     return []
 
 
-# ── Main ──────────────────────────────────────────────────────────
-
 def main():
     ap = argparse.ArgumentParser(description="Extract file from PNG")
     ap.add_argument("image", help="PNG with hidden payload")