#!/usr/bin/env python3
"""
Decode the headless-Chrome screenshot (PNG) without external deps and
compare the rendered pixels of boxes A..F to prove the cross-origin
EXIF-orientation clamp is bypassed inside cross-fade().

Box layout (CSS px, matching poc.html, --force-device-scale-factor=1):
  A=( 20, 40)  B=(160, 40)  C=(300, 40)   — row 1, image-orientation:from-image
  D=( 20,160)  E=(160,160)  F=(300,160)   — row 2, image-orientation:none
Each box is 100x50.

Security invariant being tested:
  Opaque (no-CORS) cross-origin images must always be painted with EXIF
  orientation applied; an attacker page must not be able to observe the
  raw, unrotated bitmap.

Pass criteria for the bug:
  * B != A   (cross-fade[0] paints something different from the clamped direct image)
  * E != D   (cross-fade[1] with image-orientation:none escapes the clamp)
  * B == F and E == F  (both leaks paint the SAME pixels as the same-origin
                        raw reference — i.e. the unrotated bitmap)
"""
import struct, sys, zlib

def load_png(path):
    with open(path, 'rb') as f:
        data = f.read()
    assert data[:8] == b'\x89PNG\r\n\x1a\n', 'not a PNG'
    p = 8
    idat = b''
    w = h = bpp = None
    while p < len(data):
        ln = struct.unpack('>I', data[p:p+4])[0]
        typ = data[p+4:p+8]
        chunk = data[p+8:p+8+ln]
        p += 12 + ln
        if typ == b'IHDR':
            w, h, bd, ct, _, _, il = struct.unpack('>IIBBBBB', chunk)
            assert bd == 8 and il == 0, f'unsupported PNG (bd={bd}, interlace={il})'
            bpp = {0:1, 2:3, 4:2, 6:4}[ct]
        elif typ == b'IDAT':
            idat += chunk
        elif typ == b'IEND':
            break
    raw = zlib.decompress(idat)
    stride = w * bpp
    out = bytearray(h * stride)
    pos = 0
    for y in range(h):
        f = raw[pos]; pos += 1
        row = bytearray(raw[pos:pos+stride]); pos += stride
        prev = out[(y-1)*stride:y*stride] if y else bytes(stride)
        if f == 0:
            pass
        elif f == 1:  # Sub
            for i in range(bpp, stride):
                row[i] = (row[i] + row[i-bpp]) & 0xFF
        elif f == 2:  # Up
            for i in range(stride):
                row[i] = (row[i] + prev[i]) & 0xFF
        elif f == 3:  # Average
            for i in range(stride):
                a = row[i-bpp] if i >= bpp else 0
                row[i] = (row[i] + ((a + prev[i]) >> 1)) & 0xFF
        elif f == 4:  # Paeth
            for i in range(stride):
                a = row[i-bpp] if i >= bpp else 0
                b = prev[i]
                c = prev[i-bpp] if i >= bpp else 0
                pa, pb, pc = abs(b-c), abs(a-c), abs(a+b-2*c)
                pr = a if pa <= pb and pa <= pc else (b if pb <= pc else c)
                row[i] = (row[i] + pr) & 0xFF
        else:
            raise ValueError(f'filter {f}')
        out[y*stride:(y+1)*stride] = row
    return w, h, bpp, bytes(out)

def sample(out, w, bpp, x0, y0):
    """Sample an 8x4 grid of RGB triples inside a 100x50 box at (x0,y0)."""
    pts = []
    for gy in range(4):
        for gx in range(8):
            x = x0 + 10 + gx * 10
            y = y0 + 10 + gy * 10
            o = (y * w + x) * bpp
            pts.append(tuple(out[o:o+3]))
    return pts

def diff(a, b):
    """Mean per-channel abs difference between two sample grids."""
    s = 0
    for pa, pb in zip(a, b):
        s += sum(abs(x-y) for x, y in zip(pa, pb))
    return s / (len(a) * 3)

def main():
    path = sys.argv[1] if len(sys.argv) > 1 else 'screenshot.png'
    w, h, bpp, out = load_png(path)
    print(f'screenshot: {w}x{h}, {bpp} bytes/px')

    boxes = {
        'A': ( 20,  40), 'B': (160,  40), 'C': (300,  40),
        'D': ( 20, 160), 'E': (160, 160), 'F': (300, 160),
    }
    S = {k: sample(out, w, bpp, *xy) for k, xy in boxes.items()}

    TOL = 8.0  # JPEG noise / blending rounding tolerance
    def eq(x, y): return diff(S[x], S[y]) <= TOL

    print()
    for pair in ['AB','AC','AD','BE','BF','EF','DE','DF']:
        a,b = pair
        d = diff(S[a], S[b])
        print(f'  diff({a},{b}) = {d:6.2f}  -> {"SAME" if d<=TOL else "DIFFERENT"}')

    print()
    # Sanity: the test image must actually be orientation-sensitive.
    if eq('A','F'):
        print('SANITY FAIL: oriented and raw renderings are indistinguishable.')
        sys.exit(2)
    if not (eq('A','C') and eq('A','D')):
        print('SANITY FAIL: clamped/oriented references disagree (A,C,D should match).')
        sys.exit(2)

    leak0 = (not eq('B','A')) and eq('B','F')
    leak1 = (not eq('E','D')) and eq('E','F')

    print(f'variant 1 (cross-fade image[0], hardcoded kDoNotRespect): '
          f'{"LEAK CONFIRMED" if leak0 else "no leak"}')
    print(f'variant 2 (cross-fade image[1], missing ForceOrientationIfNecessary): '
          f'{"LEAK CONFIRMED" if leak1 else "no leak"}')

    if leak0 or leak1:
        print()
        print('>>> BUG REPRODUCED: opaque cross-origin image painted with raw '
              '(unrotated) pixels via cross-fade().')
        sys.exit(0)
    print('>>> Bug NOT reproduced.')
    sys.exit(1)

if __name__ == '__main__':
    main()