From 41c6e9cfc033bebdaf189cccb9622c4fe97bd82e Mon Sep 17 00:00:00 2001 From: Struchkov Mark Date: Sun, 1 Mar 2026 23:10:05 +0300 Subject: [PATCH] refactor: clean up comments and docstrings --- pack.py | 34 +--------------------------------- unpack.py | 16 ---------------- 2 files changed, 1 insertion(+), 49 deletions(-) diff --git a/pack.py b/pack.py index f18cd6b..cf731aa 100644 --- a/pack.py +++ b/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) diff --git a/unpack.py b/unpack.py index a5698e7..7c3fe6c 100644 --- a/unpack.py +++ b/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")