#!/usr/bin/env python3
"""Corrupt the hole of one full-page image in a copy of a WAL directory.

Usage: corrupt_fpi.py <pg_wal copy> <pg_waldump> <start lsn>
Finds the first record whose block 0 carries an FPI with a hole, moves the
hole so it no longer fits in BLCKSZ, recomputes the record CRC, and prints
the record's LSN.  The WAL is modified in place (use a copy).
"""
import re, struct, subprocess, sys

BLCKSZ, XLOG_BLCKSZ, SEGSZ = 8192, 8192, 16 * 1024 * 1024
SHORT_PHD, LONG_PHD, SIZEOF_XLOGRECORD = 24, 40, 24
BKPBLOCK_HAS_IMAGE, BKPIMAGE_HAS_HOLE, BKPIMAGE_COMPRESSED = 0x10, 0x01, 0x04 | 0x08 | 0x10

# CRC-32C (Castagnoli), reflected, as PostgreSQL's pg_crc32c
TABLE = []
for i in range(256):
    c = i
    for _ in range(8):
        c = (c >> 1) ^ 0x82F63B78 if c & 1 else c >> 1
    TABLE.append(c)

def crc_update(crc, data):
    for b in data:
        crc = TABLE[(crc ^ b) & 0xFF] ^ (crc >> 8)
    return crc

def record_crc(rec):
    # same order as XLogRecordAssemble/ValidXLogRecord: payload first, then header up to xl_crc
    crc = crc_update(0xFFFFFFFF, rec[SIZEOF_XLOGRECORD:])
    crc = crc_update(crc, rec[:20])
    return crc ^ 0xFFFFFFFF

def segfile(waldir, lsn, tli=1):
    segno = lsn // SEGSZ
    return f"{waldir}/{tli:08X}{segno // 256:08X}{segno % 256:08X}"

def physical_offsets(lsn, length):
    """Map logical record bytes to physical offsets in the segment, skipping page headers."""
    off, out = lsn % SEGSZ, []
    while len(out) < length:
        if off % XLOG_BLCKSZ == 0:
            off += LONG_PHD if off == 0 else SHORT_PHD
        out.append(off)
        off += 1
    if out[-1] >= SEGSZ:
        raise ValueError("record crosses a segment boundary")
    return out

waldir, waldump, start = sys.argv[1:4]
dump = subprocess.run([waldump, "-p", waldir, "-s", start, "-b"], capture_output=True, text=True).stdout
records = re.split(r"\n(?=rmgr:)", dump)
for r in records:
    m = re.search(r"len \(rec/tot\):\s*\d+/\s*(\d+).*?lsn: ([0-9A-F]+)/([0-9A-F]+)", r)
    b = re.search(r"blkref #0:.*?FPW.*?hole: offset: (\d+), length: (\d+)", r)
    if not (m and b):
        continue
    tot, lsn = int(m.group(1)), (int(m.group(2), 16) << 32) | int(m.group(3), 16)
    if lsn % XLOG_BLCKSZ > XLOG_BLCKSZ - 64:          # keep the headers we edit on one page
        continue
    try:
        offs = physical_offsets(lsn, tot)
    except ValueError:
        continue
    path = segfile(waldir, lsn)
    data = bytearray(open(path, "rb").read())
    rec = bytearray(data[o] for o in offs)
    stored = struct.unpack_from("<I", rec, 20)[0]
    assert record_crc(rec) == stored, "CRC implementation/extraction mismatch"
    block_id, fork_flags = rec[24], rec[25]
    assert block_id == 0 and fork_flags & BKPBLOCK_HAS_IMAGE
    bimg_len, hole_offset, bimg_info = struct.unpack_from("<HHB", rec, 28)
    assert bimg_info & BKPIMAGE_HAS_HOLE
    compressed = bool(bimg_info & BKPIMAGE_COMPRESSED)
    hole_length = struct.unpack_from("<H", rec, 33)[0] if compressed else BLCKSZ - bimg_len
    # push the hole past the end of the page, leaving everything else (incl. compressed data) intact
    new_offset = BLCKSZ - hole_length + 1000
    struct.pack_into("<H", rec, 30, new_offset)
    struct.pack_into("<I", rec, 20, record_crc(rec))
    for i in range(40):                            # write back the (single-page) header region
        data[offs[i]] = rec[i]
    open(path, "wb").write(data)
    print(f"lsn={lsn >> 32:X}/{lsn & 0xFFFFFFFF:08X} compressed={compressed} bimg_len={bimg_len} "
          f"hole_offset {hole_offset} -> {new_offset}, hole_length={hole_length} "
          f"(offset+length={new_offset + hole_length} > BLCKSZ) crc {stored:08x} -> {struct.unpack_from('<I', rec, 20)[0]:08x}")
    break
else:
    sys.exit("no suitable FPI record found")