#!/usr/bin/env python3
"""
Use Frida to extract the Widevine private key and ClientIdentification
from the running android.hardware.drm@1.4-service.widevine process,
then package them into a .wvd device file.
"""

import frida, sys, time, base64, json, os
from pathlib import Path

OUTPUT_DIR = Path(__file__).parent / "L3"
WVD_OUT    = Path(__file__).parent / "device.wvd"

private_key_der = None
client_id_blob  = None

JS_HOOK = r"""
'use strict';

// ── Helpers ──────────────────────────────────────────────────────────────────

function hexStr(ptr, len) {
    if (ptr.isNull()) return null;
    return ptr.readByteArray(len);
}

function findExport(modName, symName) {
    try {
        return Module.findExportByName(modName, symName);
    } catch(e) { return null; }
}

// ── Hook OpenSSL i2d_RSAPrivateKey ────────────────────────────────────────────
// When the CDM serialises the RSA key to DER we capture it.
(function hookRSA() {
    const sym = findExport(null, 'i2d_RSAPrivateKey');
    if (!sym) { console.log('[!] i2d_RSAPrivateKey not found'); return; }

    Interceptor.attach(sym, {
        onLeave: function(retval) {
            const len = retval.toInt32();
            if (len <= 0) return;
            // arg1 is RSA*, arg2 is unsigned char** (output buffer)
            const ppBuf = this.context.rsi;  // x86_64: 2nd arg in rsi
            if (ppBuf.isNull()) return;
            const pBuf = ppBuf.readPointer();
            if (pBuf.isNull()) return;
            // Walk back to start (i2d advances the pointer)
            const start = pBuf.sub(len);
            const der = start.readByteArray(len);
            send({ type: 'private_key', data: Array.from(new Uint8Array(der)) });
            console.log('[+] Captured RSA private key (' + len + ' bytes)');
        }
    });
    console.log('[*] Hooked i2d_RSAPrivateKey');
})();

// ── Hook RSA_sign (captures key at signing time) ──────────────────────────────
(function hookRSASign() {
    const sym = findExport(null, 'RSA_sign');
    if (!sym) { console.log('[!] RSA_sign not found'); return; }

    Interceptor.attach(sym, {
        onEnter: function(args) {
            // args[3] = RSA* rsa
            const rsa = args[3];
            if (rsa.isNull()) return;
            // Use RSA_size to get key size
            const rsaSizeSym = findExport(null, 'RSA_size');
            if (!rsaSizeSym) return;
            const rsaSize = new NativeFunction(rsaSizeSym, 'int', ['pointer']);
            const sz = rsaSize(rsa) * 8; // bits
            console.log('[*] RSA_sign called, key size ~' + sz + ' bits');
        }
    });
})();

// ── Hook EVP_PKEY_get1_RSA then serialize ─────────────────────────────────────
(function hookEVP() {
    const get1 = findExport(null, 'EVP_PKEY_get1_RSA');
    const i2d  = findExport(null, 'i2d_RSAPrivateKey');
    if (!get1 || !i2d) return;

    const i2dFn = new NativeFunction(i2d, 'int', ['pointer', 'pointer']);
    Interceptor.attach(get1, {
        onLeave: function(retval) {
            if (retval.isNull()) return;
            const rsa = retval;
            const outPtr = Memory.alloc(8);
            outPtr.writePointer(NULL);
            const len = i2dFn(rsa, outPtr);
            if (len <= 0) return;
            const buf = outPtr.readPointer().sub(len);
            const der = buf.readByteArray(len);
            send({ type: 'private_key', data: Array.from(new Uint8Array(der)) });
            console.log('[+] EVP_PKEY_get1_RSA -> captured RSA key (' + len + ' bytes)');
        }
    });
})();

// ── Scan for ClientIdentification in heap ─────────────────────────────────────
// Protobuf starts: 0x08 (field1 varint) then type byte, then 0x12 (field2 bytes)...
// We trigger a license request generation which causes the CDM to serialize ClientID.
// Hook wvcdm::ClientIdentification::SerializeToString or similar.
(function hookClientID() {
    // Find libwvdrmengine.so
    const mod = Process.findModuleByName('libwvdrmengine.so');
    if (!mod) { console.log('[!] libwvdrmengine.so not loaded yet'); return; }

    console.log('[*] libwvdrmengine.so base: ' + mod.base);

    // Enumerate exports for anything containing 'Serialize' or 'client'
    const exports = mod.enumerateExports();
    for (const exp of exports) {
        const n = exp.name.toLowerCase();
        if ((n.includes('clientidentif') || n.includes('client_id')) && n.includes('serial')) {
            console.log('[*] Candidate export: ' + exp.name + ' @ ' + exp.address);
        }
    }
})();

console.log('[*] Frida hooks installed, waiting for CDM activity...');
"""

def on_message(message, data):
    global private_key_der, client_id_blob
    if message['type'] == 'send':
        payload = message['payload']
        mtype = payload.get('type')
        raw = bytes(payload.get('data', []))

        if mtype == 'private_key' and len(raw) > 200:
            private_key_der = raw
            print(f'[+] Private key captured: {len(raw)} bytes')
            save_progress()

        elif mtype == 'client_id' and len(raw) > 100:
            client_id_blob = raw
            print(f'[+] ClientIdentification captured: {len(raw)} bytes')
            save_progress()

    elif message['type'] == 'error':
        print(f'[!] Frida error: {message["description"]}')
    else:
        print(message)

def save_progress():
    if private_key_der:
        out = OUTPUT_DIR / "private_key.der"
        out.write_bytes(private_key_der)
        print(f'    Saved: {out}')

    if client_id_blob:
        out = OUTPUT_DIR / "client_id.bin"
        out.write_bytes(client_id_blob)
        print(f'    Saved: {out}')

    if private_key_der and client_id_blob:
        build_wvd()

def build_wvd():
    import subprocess
    print('[*] Building .wvd file...')
    result = subprocess.run([
        'pywidevine', 'create-device',
        '-t', 'android',
        '-l', '3',
        '-k', str(OUTPUT_DIR / 'private_key.der'),
        '-c', str(OUTPUT_DIR / 'client_id.bin'),
        '-o', str(WVD_OUT.parent),
    ], capture_output=True, text=True)
    print(result.stdout)
    if result.returncode == 0:
        print(f'[+] WVD created: {WVD_OUT}')
    else:
        print(f'[!] WVD creation failed: {result.stderr}')


def main():
    print('[*] Connecting to Frida...')
    dev = frida.get_usb_device()

    # Attach to widevine process
    target = 'android.hardware.drm@1.4-service.widevine'
    print(f'[*] Attaching to {target}...')
    session = dev.attach(target)
    session.on('detached', lambda reason: print(f'[!] Session detached: {reason}'))

    script = session.create_script(JS_HOOK)
    script.on('message', on_message)
    script.load()
    print('[*] Script loaded. Now trigger DRM activity in the emulator browser.')
    print('    Navigate to https://bitmovin.com/demos/drm and play a Widevine video.')
    print('    Waiting... (Ctrl+C to stop)\n')

    try:
        while True:
            time.sleep(1)
            if private_key_der and client_id_blob:
                print('[+] All data captured!')
                break
    except KeyboardInterrupt:
        pass

    # If we got partial data, save it anyway
    save_progress()
    print('\n[*] Done.')
    print(f'    private_key: {"OK" if private_key_der else "MISSING"}')
    print(f'    client_id:   {"OK" if client_id_blob else "MISSING"}')


if __name__ == '__main__':
    main()