#!/usr/bin/env python3
"""
Hook Chrome's child CDM process to capture the Widevine RSA private key.

Chrome's EME uses its own bundled CDM running in sandboxed_process0 (or
privileged_process0 on some builds). We use Frida spawn-gating to intercept
every child Chrome spawns, enumerate its crypto exports, install RSA hooks in
the right process, then navigate Chrome to drm_test.html which auto-triggers
a Widevine challenge.

Prerequisite: Chrome --no-sandbox flag should be set so Frida can inject into
sandboxed processes. Run once before starting this script:
  adb shell "echo 'chrome --no-sandbox --disable-gpu-watchdog' > /data/local/tmp/chrome-command-line"

Usage:
  python3 chrome_cdm.py
"""
import frida, time, sys, subprocess, os, threading
from pathlib import Path

OUT = Path(__file__).parent / "L3"
OUT.mkdir(exist_ok=True)
ADB = os.path.expanduser("~/Android/Sdk/platform-tools/adb")

# Per-process state
proc_exports = {}   # pid → list of export dicts
proc_sessions = {}  # pid → frida session
proc_scripts  = {}  # pid → frida script
captured = [False]

DIAG_JS = r"""
(function() {
    send({type:'modules', pid: Process.id,
          names: Process.enumerateModules().map(function(m){ return m.name; })});

    Process.enumerateModules().forEach(function(mod) {
        if (mod.name.toLowerCase().indexOf('widevin') === -1 &&
            mod.name.toLowerCase().indexOf('crypto')  === -1 &&
            mod.name.toLowerCase().indexOf('ssl')     === -1) return;
        try {
            mod.enumerateExports().forEach(function(e) {
                var n = e.name.toLowerCase();
                if (n.indexOf('rsa') !== -1 || n.indexOf('sign') !== -1 ||
                    n.indexOf('private') !== -1) {
                    send({type:'export', mod: mod.name, name: e.name,
                          addr: e.address.toString()});
                }
            });
        } catch(ex) {}
    });
    send({type:'diag_done', pid: Process.id});
})();
"""

def make_hook_js(exports):
    addrs = {e['name']: e['addr'] for e in exports}
    i2d_addr = addrs.get('i2d_RSAPrivateKey', '0')
    rtk_addr = addrs.get('RSA_private_key_to_bytes', '0')
    sign_a   = addrs.get('RSA_sign', '0')
    raw_a    = addrs.get('RSA_sign_raw', '0')
    priv_a   = addrs.get('RSA_private_encrypt', '0')
    from_a   = addrs.get('RSA_private_key_from_bytes', '0')

    return f"""
'use strict';
var alreadySent = false;

var i2dAddr  = '{i2d_addr}';
var rtkAddr  = '{rtk_addr}';
var signA    = '{sign_a}';
var rawA     = '{raw_a}';
var privA    = '{priv_a}';
var fromA    = '{from_a}';

// Also scan for RSA key in memory (Chrome's BoringSSL is bundled, not exported)
function scanForRSAKey(label) {{
    if (alreadySent) return;
    var pattern = '30 82 ?? ?? 02 01 00 02 82';
    var found = 0;
    Process.enumerateRanges({{protection:'rw-', coalesce:true}}).forEach(function(r) {{
        if (r.size < 200) return;
        try {{
            Memory.scan(r.base, r.size, pattern, {{
                onMatch: function(addr) {{
                    if (alreadySent) return 'stop';
                    var innerLen = (addr.add(2).readU8() << 8) | addr.add(3).readU8();
                    var total = 4 + innerLen;
                    if (total < 200 || total > 5000) return;
                    var v = addr.add(4).readByteArray(3);
                    var va = new Uint8Array(v);
                    if (va[0] !== 0x02 || va[1] !== 0x01 || va[2] !== 0x00) return;
                    found++;
                    alreadySent = true;
                    send({{type:'private_key', src: label + '@' + addr,
                          data: Array.from(new Uint8Array(addr.readByteArray(total)))}});
                    return 'stop';
                }},
                onError: function() {{}},
                onComplete: function() {{}}
            }});
        }} catch(e) {{}}
    }});
    send({{type:'log', msg:'Scan(' + label + '): ' + found + ' RSA key candidates'}});
}}

var dumpKey;
if (rtkAddr !== '0') {{
    var rtk = new NativeFunction(ptr(rtkAddr), 'int', ['pointer','pointer','pointer']);
    dumpKey = function(rsa, label) {{
        if (alreadySent || !rsa || rsa.isNull()) return;
        var pBuf = Memory.alloc(8), pLen = Memory.alloc(8);
        pBuf.writePointer(ptr('0')); pLen.writeU64(0);
        if (rtk(pBuf, pLen, rsa) !== 1) return;
        var buf = pBuf.readPointer(), len = pLen.readU64().toNumber();
        if (len < 100 || buf.isNull()) return;
        alreadySent = true;
        send({{type:'private_key', src: label,
              data: Array.from(new Uint8Array(buf.readByteArray(len)))}});
    }};
}} else if (i2dAddr !== '0') {{
    var i2d = new NativeFunction(ptr(i2dAddr), 'int', ['pointer','pointer']);
    dumpKey = function(rsa, label) {{
        if (alreadySent || !rsa || rsa.isNull()) return;
        var pp = Memory.alloc(8); pp.writePointer(ptr('0'));
        var len = i2d(rsa, pp); if (len <= 0) return;
        var buf = pp.readPointer(); if (buf.isNull()) return;
        alreadySent = true;
        send({{type:'private_key', src: label,
              data: Array.from(new Uint8Array(buf.readByteArray(len)))}});
    }};
}} else {{
    dumpKey = function(rsa, label) {{
        send({{type:'log', msg: label + ' fired — no serializer, trying memory scan'}});
        scanForRSAKey(label + '_scan');
    }};
}}

if (signA !== '0') {{
    Interceptor.attach(ptr(signA), {{
        onEnter: function(a) {{ this.rsa = a[5]; }},
        onLeave: function(r) {{ if (r.toInt32()===1) dumpKey(this.rsa,'RSA_sign'); }}
    }}); send({{type:'log', msg:'Hooked RSA_sign'}});
}}
if (rawA !== '0') {{
    Interceptor.attach(ptr(rawA), {{
        onEnter: function(a) {{ this.rsa = a[0]; }},
        onLeave: function(r) {{ if (r.toInt32()===1) dumpKey(this.rsa,'RSA_sign_raw'); }}
    }}); send({{type:'log', msg:'Hooked RSA_sign_raw'}});
}}
if (privA !== '0') {{
    Interceptor.attach(ptr(privA), {{
        onEnter: function(a) {{ this.rsa = a[3]; }},
        onLeave: function(r) {{ if (r.toInt32()>0) dumpKey(this.rsa,'RSA_private_encrypt'); }}
    }}); send({{type:'log', msg:'Hooked RSA_private_encrypt'}});
}}
if (fromA !== '0') {{
    Interceptor.attach(ptr(fromA), {{
        onEnter: function(a) {{ this.inPtr = a[0]; }},
        onLeave: function(r) {{
            if (!r.isNull()) {{
                send({{type:'log', msg:'RSA_private_key_from_bytes → RSA*=' + r}});
                dumpKey(r, 'RSA_private_key_from_bytes');
            }}
        }}
    }}); send({{type:'log', msg:'Hooked RSA_private_key_from_bytes'}});
}}
send({{type:'hooks_ready'}});
"""

def save_key(raw, src):
    captured[0] = True
    out = OUT / 'private_key.der'
    out.write_bytes(raw)
    print(f'\n[+] Private key saved → {out}  ({len(raw)} bytes)')
    print(f'    Source: {src}')
    print(f'    Header: {raw[:16].hex()}')

def navigate_chrome():
    url = 'http://localhost:8765/drm_test.html'
    subprocess.run([ADB, 'shell', 'am', 'start', '-a', 'android.intent.action.VIEW',
                    '-d', url, 'com.android.chrome'], capture_output=True)
    print(f'  Navigated Chrome to {url}')

def make_handler(pid):
    def on_message(msg, data):
        if captured[0]:
            return
        if msg['type'] == 'send':
            p = msg['payload']
            t = p.get('type')
            if t == 'modules':
                interesting = [n for n in p['names']
                               if any(k in n.lower() for k in
                                      ['widevin','crypto','ssl','boringssl','monochrome'])]
                if interesting:
                    print(f'  [PID {pid}] Interesting modules: {interesting}')
            elif t == 'export':
                proc_exports.setdefault(pid, []).append(p)
                print(f'  [PID {pid}] Export: {p["mod"]}!{p["name"]}')
            elif t == 'diag_done':
                exps = proc_exports.get(pid, [])
                print(f'  [PID {pid}] Diagnostic done. {len(exps)} relevant exports.')
                if exps:
                    install_hooks(pid)
                else:
                    print(f'  [PID {pid}] No exports found — skipping hooks for this process.')
            elif t == 'hooks_ready':
                print(f'  [PID {pid}] Hooks ready! Navigating Chrome...')
                time.sleep(1)   # let Chrome stabilize before navigating
                navigate_chrome()
            elif t == 'log':
                print(f'  [PID {pid}] {p["msg"]}')
            elif t == 'private_key' and not captured[0]:
                raw = bytes(p['data'])
                if len(raw) > 100:
                    save_key(raw, p.get('src', '?'))
        elif msg['type'] == 'error':
            print(f'  [PID {pid}] JS Error: {msg["description"]}')
    return on_message

def install_hooks(pid):
    session = proc_sessions.get(pid)
    if not session:
        return
    exps = proc_exports.get(pid, [])
    print(f'  [PID {pid}] Installing hooks based on {len(exps)} exports...')
    js = make_hook_js(exps)
    script = session.create_script(js)
    script.on('message', make_handler(pid))
    script.load()
    proc_scripts[pid] = script

def on_child_added(child):
    pid   = child.pid
    ident = child.identifier or ''
    print(f'[*] Child spawned: PID={pid}  id={ident}')

    # Skip GPU/browser processes by name — they cause instability when hooked at startup
    # Focus on sandboxed (renderer/CDM) and privileged (utility/CDM) processes
    # With --no-sandbox, sandboxed_process0 can be injected into
    skip_keywords = []  # hook everything now that --no-sandbox is set
    if any(k in ident for k in skip_keywords):
        print(f'    Skipping {ident}')
        dev.resume(pid)
        return

    try:
        session = dev.attach(pid)
        proc_sessions[pid] = session
    except Exception as e:
        print(f'    Could not attach to PID {pid}: {e}')
        dev.resume(pid)
        return

    proc_exports[pid] = []

    # Delay injection slightly so the process can initialize its native libraries
    def inject_delayed():
        time.sleep(0.5)
        try:
            script = session.create_script(DIAG_JS)
            script.on('message', make_handler(pid))
            script.load()
            proc_scripts[pid] = script
        except Exception as e:
            print(f'    Injection failed for PID {pid}: {e}')

    threading.Thread(target=inject_delayed, daemon=True).start()
    dev.resume(pid)

# ── Main ──────────────────────────────────────────────────────────────────────
print('[*] Setting up adb reverse so localhost:8765 reaches the HTTP server...')
subprocess.run([ADB, 'reverse', 'tcp:8765', 'tcp:8765'], capture_output=True)

dev = frida.get_usb_device()

print('[*] Force-stopping Chrome...')
subprocess.run([ADB, 'shell', 'am', 'force-stop', 'com.android.chrome'], capture_output=True)
time.sleep(1)

print('[*] Enabling spawn gating...')
dev.enable_spawn_gating()
dev.on('spawn-added', on_child_added)

print('[*] Spawning com.android.chrome...')
try:
    pid = dev.spawn('com.android.chrome')
    print(f'[*] Chrome PID={pid}')
    proc_sessions[pid] = dev.attach(pid)
    dev.resume(pid)
    print('[*] Chrome started. Waiting for child processes...\n')
except Exception as e:
    print(f'[!] Spawn failed: {e}')
    sys.exit(1)

print('[*] Monitoring (Ctrl+C to quit)...\n')
try:
    while not captured[0]:
        time.sleep(1)
    print('\n[+] Key captured! Done.')
except KeyboardInterrupt:
    print('\n[!] Stopped.')