#!/usr/bin/env python3
"""
Captures Widevine L3 device key material from the DRM HAL process at runtime.

Three complementary strategies:
  1. GetTokenFromKeybox → std::string* output (device token / derived key)
  2. wvcdm::_oecc104(buf, &len, level) → 4KB buffer filled with key material
  3. LoadCertificatePrivateKey.onEnter → dump 4KB of CryptoSession object;
     Python brute-forces every 16-byte chunk as AES device_key against
     wrapped_privkey.bin to find the right one without needing to know
     the keybox format at all.
"""
import frida, time
from pathlib import Path
from Crypto.Cipher import AES

OUT     = Path(__file__).parent / 'L3'
WRAPPED = OUT / 'wrapped_privkey.bin'
OUT.mkdir(exist_ok=True)

AY64_IV = bytes.fromhex('41a210cb7f65b9856a9b87090982a290')

# ── Symbol names ──────────────────────────────────────────────────────────────
SYM_LOAD_KEY  = '_ZN5wvcdm13CryptoSession25LoadCertificatePrivateKeyERKNS_16CryptoWrappedKeyE'
SYM_GEN_RSA   = '_ZN5wvcdm13CryptoSession20GenerateRsaSignatureERKNSt3__112basic_stringIcNS1_11char_traitsIcEENS1_9allocatorIcEEEEPS7_'
SYM_PREP_SIGN = '_ZN5wvcdm13CryptoSession28PrepareAndSignLicenseRequestERKNSt3__112basic_stringIcNS1_11char_traitsIcEENS1_9allocatorIcEEEEPS7_SA_'
SYM_FS_READ   = '_ZN6wvoec333OEMCrypto_Level3AndroidFileSystem4ReadEPKcPvm'
SYM_L3_SIG    = 'Level3_GenerateRSASignature'
SYM_GET_TOKEN = '_ZN5wvcdm13CryptoSession18GetTokenFromKeyboxEPNSt3__112basic_stringIcNS1_11char_traitsIcEENS1_9allocatorIcEEEE'
SYM_OECC104   = '_ZN5wvcdm8_oecc104EPhPmNS_13SecurityLevelE'
SYM_OECC107   = '_oecc107'   # actual RSA key decryptor (calls AES-CBC with device_key)
SYM_WPLLQDNZ  = 'wpllqdnz'  # real crypto impl: reads klsosnbn+0x3d0, LCG-derives device_key

JS = r"""
'use strict';
var libwv = Process.findModuleByName('libwvhidl.so');
if (!libwv) { send({type:'log', msg:'[!] libwvhidl.so not loaded'}); }
else { send({type:'log', msg:'libwvhidl.so @ ' + libwv.base + '  size=' + libwv.size}); }

function sym(name) {
    if (!libwv) return null;
    try { var a = libwv.findExportByName(name); return (a && !a.isNull()) ? a : null; }
    catch(e) { return null; }
}

var addrLoad   = sym('""" + SYM_LOAD_KEY + r"""');
var addrGen    = sym('""" + SYM_GEN_RSA + r"""');
var addrPrep   = sym('""" + SYM_PREP_SIGN + r"""');
var addrFsRd   = sym('""" + SYM_FS_READ + r"""');
var addrL3     = sym('""" + SYM_L3_SIG + r"""');
var addrTok    = sym('""" + SYM_GET_TOKEN + r"""');
var addrOecc   = sym('""" + SYM_OECC104 + r"""');
var addrOecc107 = sym('""" + SYM_OECC107 + r"""');

// ── klsosnbn global BSS (module_base + 0x2d41d8) — RSA key stored here post-load ──
// Read only within the mapped range to avoid crossing page boundary
var addrKlsosnbn = libwv ? libwv.base.add(0x2d41d8) : null;
var klsosnbnReadSize = 0;
if (addrKlsosnbn) {
    // findRangeByAddress is broken in Frida 17; enumerate instead
    var klsosnbnN = parseInt(addrKlsosnbn.toString(), 16);
    Process.enumerateRanges({protection:'rw-', coalesce:false}).forEach(function(r) {
        if (klsosnbnReadSize > 0) return;
        var baseN = parseInt(r.base.toString(), 16);
        if (klsosnbnN >= baseN && klsosnbnN < baseN + r.size) {
            klsosnbnReadSize = Math.min(r.size - (klsosnbnN - baseN), 8192);
        }
    });
    send({type:'log', msg:'klsosnbn @ ' + addrKlsosnbn + '  readable=' + klsosnbnReadSize + ' B'});
}

send({type:'addrs',
    LoadCertificatePrivateKey:    String(addrLoad),
    GetTokenFromKeybox:           String(addrTok),
    _oecc104:                     String(addrOecc),
    _oecc107:                     String(addrOecc107),
    FS_Read:                      String(addrFsRd),
    GenerateRsaSignature:         String(addrGen),
    Level3_GenerateRSASignature:  String(addrL3),
    klsosnbn:                     String(addrKlsosnbn)
});

function hex(p, n) {
    try { return Array.from(new Uint8Array(p.readByteArray(n)))
              .map(function(b){return ('0'+b.toString(16)).slice(-2);}).join(''); }
    catch(e) { return '(err)'; }
}

// ── RSA serialisation via libcrypto.so ────────────────────────────────────────
// libwvhidl.so dynamically links libcrypto.so (BoringSSL).
// Module.findExportByName(null,...) is broken in Frida 17; use per-module lookup.
var fnI2D = null;
(function() {
    var lc = Process.findModuleByName('libcrypto.so');
    if (!lc) { send({type:'log', msg:'[!] libcrypto.so not found in process'}); return; }
    send({type:'log', msg:'libcrypto.so @ ' + lc.base});

    // Prefer BoringSSL-native RSA_private_key_to_bytes(out**, out_len*, rsa*)
    var rtkAddr = lc.findExportByName('RSA_private_key_to_bytes');
    if (rtkAddr && !rtkAddr.isNull()) {
        var _rtb = new NativeFunction(rtkAddr, 'int', ['pointer', 'pointer', 'pointer']);
        fnI2D = function(rsa, pp) {
            var pLen = Memory.alloc(8); pLen.writeU64(0);
            pp.writePointer(ptr(0));
            if (_rtb(pp, pLen, rsa) !== 1) return -1;
            return pLen.readU64().toNumber();
        };
        send({type:'log', msg:'fnI2D = RSA_private_key_to_bytes @ ' + rtkAddr});
        return;
    }
    // Fallback: i2d_RSAPrivateKey(rsa*, pp**)  [OpenSSL compat]
    var i2dAddr = lc.findExportByName('i2d_RSAPrivateKey');
    if (i2dAddr && !i2dAddr.isNull()) {
        var _i2d = new NativeFunction(i2dAddr, 'int', ['pointer', 'pointer']);
        fnI2D = function(rsa, pp) { return _i2d(rsa, pp); };
        send({type:'log', msg:'fnI2D = i2d_RSAPrivateKey @ ' + i2dAddr});
        return;
    }
    send({type:'log', msg:'[!] No RSA serialiser in libcrypto.so — struct scan only'});
})();

// ── Validate a pointer as BoringSSL RSA-2048 private key struct ───────────────
// RSA struct layout (BoringSSL, Android 10-13):
//   +0x00 BIGNUM *n  +0x08 BIGNUM *e  +0x10 BIGNUM *d  +0x18 BIGNUM *p  ...
// BIGNUM struct: +0x00 BN_ULONG *d  +0x08 uint32 top  +0x0c uint32 dmax
//                +0x10 uint32 neg   +0x14 uint32 flags
function isPrivateRSA2048(rsa) {
    if (!rsa || rsa.isNull()) return false;
    try {
        var n = rsa.readPointer();           // BIGNUM *n
        if (n.isNull()) return false;
        var n_top = n.add(8).readU32();      // n->top (32 for RSA-2048)
        if (n_top < 16 || n_top > 64) return false;  // accept 1024/2048/4096
        if (n.add(0x10).readU32() !== 0) return false; // n->neg must be 0
        var e = rsa.add(8).readPointer();    // BIGNUM *e
        if (e.isNull()) return false;
        var e_top = e.add(8).readU32();
        if (e_top < 1 || e_top > 4) return false;
        var d = rsa.add(0x10).readPointer(); // BIGNUM *d — NULL for public keys
        if (d.isNull()) return false;
        var d_top = d.add(8).readU32();
        if (d_top < 16 || d_top > 64) return false;
        return true;
    } catch(ex) { return false; }
}

var rsaAlreadySent = false;
function trySerializeRSA(rsa, src) {
    if (rsaAlreadySent || !rsa || rsa.isNull()) return false;
    if (!isPrivateRSA2048(rsa)) return false;
    if (!fnI2D) {
        send({type:'log', msg:'isPrivateRSA2048 PASS @ ' + rsa + ' [' + src + '] — no serialiser'});
        // Fall back: dump 2048 bytes around the RSA struct for offline analysis
        try {
            send({type:'rsa_struct_dump', src: src, rsa: rsa.toString(),
                  data: Array.from(new Uint8Array(rsa.readByteArray(2048)))});
        } catch(ex) {}
        return false;
    }
    try {
        var pp = Memory.alloc(8); pp.writePointer(ptr(0));
        var len = fnI2D(rsa, pp);
        if (len <= 100) return false;
        var buf = pp.readPointer();
        if (buf.isNull()) return false;
        rsaAlreadySent = true;
        send({type:'rsa_i2d_hit', src: src, len: len,
              data: Array.from(new Uint8Array(buf.readByteArray(len)))});
        return true;
    } catch(ex) {
        send({type:'log', msg:'trySerializeRSA exception [' + src + ']: ' + ex});
        return false;
    }
}

// ── Scan all rw-/rwx memory for PKCS#1 or PKCS#8 RSA DER ────────────────────
function scanForRSADER(tag) {
    if (rsaAlreadySent) return;
    var patterns = [
        '30 82 ?? ?? 02 01 00 02 82 01 01',  // PKCS#1 RSA-2048 n=257B
        '30 82 ?? ?? 02 01 00 02 82 01 00',  // PKCS#1 RSA-2048 n=256B
        '30 82 ?? ?? 02 01 00 30 0d 06 09',  // PKCS#8
    ];
    var found = 0;
    ['rw-', 'rwx'].forEach(function(prot) {
        if (found > 0) return;
        Process.enumerateRanges({protection: prot, coalesce: true}).forEach(function(r) {
            if (r.size < 300 || found > 0) return;
            patterns.forEach(function(pat) {
                if (found > 0) return;
                try {
                    Memory.scan(r.base, r.size, pat, {
                        onMatch: function(addr) {
                            try {
                                var inner = (addr.add(2).readU8() << 8) | addr.add(3).readU8();
                                var total = 4 + inner;
                                if (total < 400 || total > 6000) return;
                                found++;
                                rsaAlreadySent = true;
                                send({type:'rsa_scan_hit', tag: tag, addr: addr.toString(),
                                      data: Array.from(new Uint8Array(addr.readByteArray(total)))});
                                return 'stop';
                            } catch(e) {}
                        },
                        onError: function() {},
                        onComplete: function() {}
                    });
                } catch(e) {}
            });
        });
    });
    send({type:'log', msg: tag + ' DER scan: ' + (found > 0 ? 'HIT' : 'no match')});
}

// ── Scan all rw- memory for a valid BoringSSL private RSA-2048 struct ─────────
function scanForRSAStruct(tag) {
    if (rsaAlreadySent) return;
    var found = 0;
    Process.enumerateRanges({protection: 'rw-', coalesce: true}).forEach(function(r) {
        if (r.size < 64 || found > 0) return;
        // Step by 8 bytes (pointer-aligned)
        var N = Math.min(r.size - 64, 65536);  // limit per region to avoid timeout
        try {
            for (var i = 0; i < N && found === 0; i += 8) {
                var candidate = r.base.add(i);
                if (isPrivateRSA2048(candidate)) {
                    send({type:'log', msg: tag + ' RSA struct candidate @ ' + candidate});
                    if (trySerializeRSA(candidate, tag + '_struct_scan')) {
                        found++;
                    }
                }
            }
        } catch(e) {}
    });
    send({type:'log', msg: tag + ' struct scan: ' + (found > 0 ? 'HIT' : 'no match')});
}

// ── Hook RSA_free — intercept PRIVATE key just before it's freed ──────────────
(function() {
    var lc = Process.findModuleByName('libcrypto.so');
    if (!lc) return;
    var a = lc.findExportByName('RSA_free');
    if (!a || a.isNull()) { send({type:'log', msg:'RSA_free not found in libcrypto.so'}); return; }
    send({type:'log', msg:'Hooking RSA_free @ ' + a});
    Interceptor.attach(a, {
        onEnter: function(args) {
            if (rsaAlreadySent) return;
            trySerializeRSA(args[0], 'RSA_free');
        }
    });
})();

// ── Saved session state (shared across hooks) ─────────────────────────────────
var savedSessImpl = null;  // set in _oecc107[+0xa6]
var savedCsess    = null;  // set in LoadCertPrivKey.onEnter

// ── 1. Hook LoadCertificatePrivateKey ─────────────────────────────────────────
if (addrLoad) {
    Interceptor.attach(addrLoad, {
        onEnter: function(args) {
            var a1       = args[1];
            var dataSize = a1.add(0x10).readU64().toNumber();
            var dataPtr  = a1.add(0x18).readPointer();
            this.csess   = args[0];
            savedCsess   = args[0];

            send({type:'log', msg:
                'LoadCertPrivKey ENTER  this=' + args[0] +
                '  data_size=' + dataSize + '  data_ptr=' + dataPtr});

            // Save wrapped key blob
            if (dataSize > 0 && dataSize < 10000 && !dataPtr.isNull()) {
                try {
                    send({type:'log', msg: '  blob[0..31]: ' + hex(dataPtr, 32)});
                    send({type:'wrapped_key',
                          data: Array.from(new Uint8Array(dataPtr.readByteArray(dataSize)))});
                } catch(e) {}
            }

            // Dump CryptoSession object (4096 bytes) → Python brute-forces device_key
            try {
                var csData = args[0].readByteArray(4096);
                send({type:'csess_dump', data: Array.from(new Uint8Array(csData))});
            } catch(e) {
                send({type:'log', msg: '  csess dump failed: ' + e});
            }
        },
        onLeave: function(ret) {
            send({type:'log', msg: 'LoadCertPrivKey LEAVE ret=' + ret});
            if (ret.toInt32() !== 0) return;

            send({type:'log', msg:'LoadCertPrivKey success — dumping key stores...'});

            // ── Quick targeted dumps (run synchronously — fast) ──────────────
            // 1a. Dump klsosnbn — RSA key or RSA* lives here post-load
            if (addrKlsosnbn && klsosnbnReadSize > 0) {
                try {
                    send({type:'klsosnbn_after',
                          data: Array.from(new Uint8Array(addrKlsosnbn.readByteArray(klsosnbnReadSize)))});
                } catch(e) { send({type:'log', msg:'klsosnbn dump failed: ' + e}); }
            }

            // 1b. Dump CryptoSession AFTER load
            if (this.csess && !this.csess.isNull()) {
                try {
                    send({type:'csess_after',
                          data: Array.from(new Uint8Array(this.csess.readByteArray(4096)))});
                } catch(e) {}
            }

            // 1c. Dump Level3 session_impl AFTER load
            if (savedSessImpl && !savedSessImpl.isNull()) {
                try {
                    send({type:'oecc107_sess_after',
                          data: Array.from(new Uint8Array(savedSessImpl.readByteArray(4096)))});
                } catch(e) {}
            }

            // ── Defer slow scans — run AFTER we return so the DRM thread is
            //    unblocked and GenerateRsaSignature can fire normally ──────────
            var csessCapture = this.csess;
            setTimeout(function() {
                if (rsaAlreadySent) return;
                send({type:'log', msg:'[deferred] running DER + struct scans...'});
                scanForRSADER('load_deferred');
                if (!rsaAlreadySent) scanForRSAStruct('load_deferred');
            }, 3000);  // 3 s delay — lets GenerateRsaSignature fire first

            send({type:'log', msg:'LoadCertPrivKey done — returning quickly, waiting for GenerateRsaSignature...'});
        }
    });
    send({type:'log', msg:'Hooked LoadCertificatePrivateKey @ ' + addrLoad});
}

// ── 2. Hook _oecc104 → captures key material buffer ──────────────────────────
// wvcdm::_oecc104(uint8_t* buf, size_t* len, SecurityLevel)
// Called at the start of LoadCertPrivKey; fills buf with device cert/key data.
if (addrOecc) {
    Interceptor.attach(addrOecc, {
        onEnter: function(args) {
            this.buf    = args[0];
            this.lenPtr = args[1];
        },
        onLeave: function(ret) {
            var rc  = ret.toInt32();
            var len = 0;
            try { len = this.lenPtr.readU64().toNumber(); } catch(e) {}
            send({type:'log', msg:'_oecc104: ret=' + rc + '  len=' + len +
                 (len > 0 ? '  data[0..15]=' + hex(this.buf, 16) : '')});
            if (rc === 0 && len > 0 && len < 65536 && !this.buf.isNull()) {
                try {
                    send({type:'oecc104_data',
                          data: Array.from(new Uint8Array(this.buf.readByteArray(len)))});
                } catch(e) {}
            }
        }
    });
    send({type:'log', msg:'Hooked _oecc104 @ ' + addrOecc});
}

// ── 3. Hook GetTokenFromKeybox ────────────────────────────────────────────────
// After return, args[1] is a filled std::string* with the device token.
// We dump ALL 24 raw bytes of the string object (Python determines the layout).
if (addrTok) {
    Interceptor.attach(addrTok, {
        onEnter: function(args) { this.out = args[1]; },
        onLeave: function(ret) {
            try {
                var raw24 = Array.from(new Uint8Array(this.out.readByteArray(24)));
                send({type:'token_raw', data: raw24,
                      msg: 'GetTokenFromKeybox ret=' + ret +
                           '  raw24=' + raw24.map(b=>('0'+b.toString(16)).slice(-2)).join('')});
            } catch(e) {
                send({type:'log', msg:'GetTokenFromKeybox → read failed: ' + e});
            }
        }
    });
    send({type:'log', msg:'Hooked GetTokenFromKeybox @ ' + addrTok});
}

// ── 4. Hook FS::Read (save ALL files) ────────────────────────────────────────
if (addrFsRd) {
    Interceptor.attach(addrFsRd, {
        onEnter: function(args) {
            this.path = args[1].readCString();
            this.buf  = args[2];
        },
        onLeave: function(ret) {
            var n = ret.toInt32();
            send({type:'log', msg:'FS::Read "' + this.path + '"  got=' + n +
                 '  data[0..15]=' + hex(this.buf, Math.min(16, n))});
            if (n > 0 && n < 65536) {
                try {
                    send({type:'fs_read', path: this.path,
                          data: Array.from(new Uint8Array(this.buf.readByteArray(n)))});
                } catch(e) {}
            }
        }
    });
    send({type:'log', msg:'Hooked FS::Read @ ' + addrFsRd});
}

// ── 5. _oecc107: actual RSA private key decryptor ────────────────────────────
// Signature (from disassembly of LoadCertificatePrivateKey → calls _oecc107):
//   _oecc107(session_handle:int, key_type:int, wrapped_key_ptr:ptr, wrapped_key_size:int)
//   key_type=0 → dispatches via session_impl->field_0x110(session_handle, wrapped_key_ptr, size)
//   key_type=1 → dispatches via session_impl->field_0x118(session_handle, key_type, wrapped_key_ptr, size)
//
// Offsets inside _oecc107 (base = _oecc107 export):
//   +0xa6 (0x119696): rbx = OEMCrypto Level3 session implementation pointer (non-null confirmed)
//   +0x103 (0x1196f3): about to call vtable via session_impl->field_0x110 (when key_type==0)
//
// Strategy:
//   A) Hook at +0xa6 to dump 4096 bytes of the OEMCrypto session struct → brute-force device_key
//   B) Hook at +0x103 to capture vtable function address and attach to it dynamically
//   C) In vtable function, dump args and scan memory for device_key context

var hookedVtableFuncs = {};

if (addrOecc107) {
    // ── A: Outer hook on _oecc107 ─────────────────────────────────────────────
    Interceptor.attach(addrOecc107, {
        onEnter: function(args) {
            send({type:'log', msg:
                '_oecc107 ENTER  sess=' + args[0].toInt32() +
                '  key_type=' + args[1].toInt32() +
                '  wrapped_ptr=' + args[2] +
                '  wrapped_sz=' + args[3].toInt32()});
        },
        onLeave: function(ret) {
            send({type:'log', msg:'_oecc107 LEAVE ret=' + ret.toInt32()});
            if (ret.toInt32() === 0 && savedSessImpl && !savedSessImpl.isNull()) {
                // Re-dump session_impl post-success — RSA* now stored in here
                try {
                    send({type:'oecc107_sess_after',
                          data: Array.from(new Uint8Array(savedSessImpl.readByteArray(4096)))});
                } catch(e) {}
            }
        }
    });
    send({type:'log', msg:'Hooked _oecc107 @ ' + addrOecc107});

    // ── B: Internal hook at +0xa6 — rbx = OEMCrypto Level3 session impl ───────
    // At this offset: rbx has been confirmed non-null (test rbx,rbx; jne here)
    var addrInner_A = addrOecc107.add(0xa6);
    Interceptor.attach(addrInner_A, {
        onEnter: function() {
            var rbx = this.context.rbx;
            savedSessImpl = rbx;  // save for _oecc107.onLeave
            send({type:'log', msg:'_oecc107[+0xa6]: session_impl=' + rbx});
            if (rbx && !rbx.isNull()) {
                try {
                    // First 8 bytes are likely a vtable or struct pointer
                    send({type:'log', msg:'  rbx[0..31]=' + hex(rbx, 32)});
                    send({type:'oecc107_sess',
                          data: Array.from(new Uint8Array(rbx.readByteArray(4096)))});
                } catch(e) {
                    send({type:'log', msg:'  rbx dump failed: ' + e});
                }
            }
        }
    });

    // ── C: Internal hook at +0xb2 — capture vtable func for path 2 (func_0x118 non-null)
    // At this offset: rax = rbx->field_0x118 (non-null), rbx = session impl
    var addrInner_B = addrOecc107.add(0xb2);
    Interceptor.attach(addrInner_B, {
        onEnter: function() {
            var rax = this.context.rax;  // = rbx->field_0x118 = vtable func
            var rbx = this.context.rbx;
            var key = rax.toString();
            send({type:'log', msg:'_oecc107[+0xb2]: vtfunc_0x118=' + rax + ' sess_impl=' + rbx});
            if (!hookedVtableFuncs[key] && !rax.isNull()) {
                hookedVtableFuncs[key] = true;
                attachVtableDecrypt(rax, 'path2_0x118');
            }
        }
    });

    // ── D: Internal hook at +0x103 — capture vtable func for path 1 (func_0x118 null, key_type==0)
    var addrInner_C = addrOecc107.add(0x103);
    Interceptor.attach(addrInner_C, {
        onEnter: function() {
            var rbx = this.context.rbx;
            if (!rbx || rbx.isNull()) return;
            try {
                var vtfunc = rbx.add(0x110).readPointer();
                var key = vtfunc.toString();
                send({type:'log', msg:'_oecc107[+0x103]: vtfunc_0x110=' + vtfunc});
                if (!hookedVtableFuncs[key] && !vtfunc.isNull()) {
                    hookedVtableFuncs[key] = true;
                    attachVtableDecrypt(vtfunc, 'path1_0x110');
                }
            } catch(e) {
                send({type:'log', msg:'_oecc107[+0x103] error: ' + e});
            }
        }
    });
}

function attachVtableDecrypt(vtfunc, tag) {
    Interceptor.attach(vtfunc, {
        onEnter: function(a) {
            this.sessH = a[0].toInt32();
            this.wkSz  = a[2].toInt32();
            send({type:'log', msg:tag + ' vtfunc_decrypt ENTER sess=' + this.sessH + ' wkSz=' + this.wkSz});
            try {
                var sp = this.context.sp;
                send({type:'vtfunc_stack',
                      data: Array.from(new Uint8Array(sp.readByteArray(2048)))});
            } catch(e) {}
        },
        onLeave: function(ret) {
            send({type:'log', msg:tag + ' vtfunc_decrypt LEAVE ret=' + ret.toInt32()});
            if (ret.toInt32() !== 0) return;
            try {
                var sp = this.context.sp;
                send({type:'vtfunc_stack_leave',
                      data: Array.from(new Uint8Array(sp.readByteArray(2048)))});
            } catch(e) {}
            // Scan for RSA DER right here — earlier timing than LoadCertPrivKey.onLeave
            // The JIT code may not have freed the DER buffer yet at this point.
            if (!rsaAlreadySent) scanForRSADER(tag + '_vtfunc_leave');
        }
    });
}

// ── 6. Signing hooks ──────────────────────────────────────────────────────────
if (addrGen) {
    Interceptor.attach(addrGen, {
        onEnter: function(a) {
            send({type:'log', msg:'GenerateRsaSignature CALLED this=' + a[0]});
            var thisCapture = a[0];

            // Quick dumps only — don't block the signing call
            if (addrKlsosnbn && klsosnbnReadSize > 0) {
                try {
                    send({type:'klsosnbn_genrsa',
                          data: Array.from(new Uint8Array(addrKlsosnbn.readByteArray(klsosnbnReadSize)))});
                } catch(e) {}
            }
            if (!thisCapture.isNull()) {
                try {
                    send({type:'csess_genrsa',
                          data: Array.from(new Uint8Array(thisCapture.readByteArray(512)))});
                } catch(e) {}
            }

            // Defer slow scans so signing completes normally
            setTimeout(function() {
                if (rsaAlreadySent) return;
                send({type:'log', msg:'[genrsa deferred] running DER + struct scans...'});
                scanForRSADER('genrsa_deferred');
                if (!rsaAlreadySent) scanForRSAStruct('genrsa_deferred');
            }, 100);
        }
    });
}
if (addrPrep) {
    Interceptor.attach(addrPrep, {
        onEnter: function(a) { send({type:'log', msg:'PrepareAndSignLicenseRequest CALLED'}); }
    });
}
if (addrL3) {
    Interceptor.attach(addrL3, {
        onEnter: function(a) { send({type:'log', msg:'Level3_GenerateRSASignature CALLED'}); },
        onLeave: function(r) { send({type:'log', msg:'Level3_GenerateRSASignature returned=' + r.toInt32()}); }
    });
}

send({type:'ready'});
"""


def lcg_derive_key(raw16: bytes, seed: int = 0x18c) -> tuple[bytes, bytes]:
    """Reproduce wpllqdnz LCG loop offline.
    Returns (derived_key, lcg_high_bytes).
    derived_key[i] = raw16[i] XOR (lcg_state & 0xFF)  → lands at rsp+0x10
    lcg_high[i]    = (lcg_state >> 8) & 0xFF           → lands at rsp+0x00
    LCG: state = state * 0x19660d + 0x3c6ef35f  (mod 2^64)
    """
    state = seed & 0xFFFFFFFFFFFFFFFF
    derived, lcg_high = [], []
    for i in range(0, 16, 2):
        state = (state * 0x19660d + 0x3c6ef35f) & 0xFFFFFFFFFFFFFFFF
        derived.append(raw16[i]   ^ (state & 0xFF))
        lcg_high.append((state >> 8) & 0xFF)
        state = (state * 0x19660d + 0x3c6ef35f) & 0xFFFFFFFFFFFFFFFF
        derived.append(raw16[i+1] ^ (state & 0xFF))
        lcg_high.append((state >> 8) & 0xFF)
    return bytes(derived), bytes(lcg_high)


def check_der(data: bytes) -> str | None:
    if len(data) < 8: return None
    if data[0:2] == b'\x30\x82' and data[4:7] == b'\x02\x01\x00' and data[7:9] == b'\x02\x82':
        return 'pkcs1'
    if data[0:2] == b'\x30\x82' and data[4:7] == b'\x02\x01\x00' and data[7:8] == b'\x30':
        return 'pkcs8'
    return None


def try_device_key(candidate: bytes, wrapped: bytes, label: str = '') -> bool:
    """Try a 16-byte candidate as AES-128-CBC device_key against wrapped_privkey.bin."""
    if len(candidate) != 16 or len(set(candidate)) < 4:
        return False
    ivs = [b'\x00' * 16, AY64_IV, wrapped[:16]]
    cts = [wrapped, wrapped[16:]]
    for iv in ivs:
        for ct in cts:
            if len(ct) % 16 != 0 or len(ct) == 0:
                continue
            try:
                plain = AES.new(candidate, AES.MODE_CBC, iv).decrypt(ct)
                kt = check_der(plain)
                if kt:
                    # Strip PKCS#7 padding
                    pad = plain[-1]
                    if 1 <= pad <= 16 and plain[-pad:] == bytes([pad]) * pad:
                        plain = plain[:-pad]
                    if plain[0:2] == b'\x30\x82':
                        inner = (plain[2] << 8) | plain[3]
                        plain = plain[:4 + inner]
                    out = OUT / 'private_key.der'
                    out.write_bytes(plain)
                    print(f'\n[+] PRIVATE KEY FOUND! ({kt}, {len(plain)} bytes)')
                    print(f'    device_key: {candidate.hex()}')
                    print(f'    iv:         {iv.hex()}')
                    print(f'    ct_len:     {len(ct)}')
                    if label:
                        print(f'    source:     {label}')
                    print(f'    Saved → {out}')
                    return True
            except Exception:
                pass
    return False


def brute_region(raw: bytes, wrapped: bytes, label: str) -> bool:
    """Try every byte offset in raw as a device_key candidate (unaligned search)."""
    for off in range(len(raw) - 15):
        if try_device_key(raw[off:off + 16], wrapped, f'{label}@{off}'):
            return True
    return False


def decode_std_string(raw24: list[int]) -> bytes | None:
    """Decode LLVM libc++ 24-byte string object (multiple layout variants)."""
    b = bytes(raw24)
    # Variant A: flag at bit 63 of first qword (byte[7] bit 7)
    if b[7] & 0x80:
        import struct
        size = struct.unpack_from('<Q', b, 8)[0]
        ptr_val = struct.unpack_from('<Q', b, 16)[0]
        if 0 < size < 65536 and ptr_val > 0x1000:
            try:
                import frida as _frida_ref
                # Return raw24 with metadata for Python-side resolution
                return None  # can't read pointers here; fall through
            except Exception:
                pass
    # Variant B: flag at bit 0 of byte[0] (short mode)
    if not (b[0] & 1):
        size = (b[0] >> 1) & 0x7f
        if 0 < size <= 22:
            return b[1:1 + size]
    # Variant C: flag at byte[23] bit 0 (old layout)
    if b[23] & 1:
        pass  # need pointer resolution
    return None


# Python-side state
wrapped_key  = [None]
found_key    = [False]
fs_reads     = {}

def on_message(msg, data):
    global wrapped_key, found_key
    if found_key[0]:
        return
    if msg['type'] == 'send':
        p = msg['payload']
        t = p.get('type', '')

        if t == 'addrs':
            print('[*] Hook addresses:')
            for k, v in p.items():
                if k == 'type': continue
                ok = '!!' if v == 'null' else 'OK'
                print(f'    [{ok}] {k} = {v}')

        elif t == 'ready':
            print('\n[*] Hooks active — trigger DRM in the browser now\n')

        elif t == 'log':
            print(f'    {p["msg"]}')

        elif t == 'wrapped_key':
            raw = bytes(p['data'])
            wrapped_key[0] = raw
            WRAPPED.write_bytes(raw)
            print(f'\n[+] Wrapped key: {len(raw)} bytes  header={raw[:16].hex()}')

        elif t == 'csess_dump':
            raw = bytes(p['data'])
            print(f'\n[*] CryptoSession dump: {len(raw)} bytes — brute-forcing device_key...')
            out_path = OUT / 'csess_dump.bin'
            out_path.write_bytes(raw)
            wk = wrapped_key[0] or (WRAPPED.read_bytes() if WRAPPED.exists() else None)
            if wk:
                if brute_region(raw, wk, 'CryptoSession'):
                    found_key[0] = True
            else:
                print('    (no wrapped key yet — will retry when it arrives)')

        elif t == 'oecc104_data':
            raw = bytes(p['data'])
            print(f'\n[+] _oecc104 output: {len(raw)} bytes  header={raw[:16].hex()}')
            (OUT / 'oecc104_data.bin').write_bytes(raw)
            wk = wrapped_key[0] or (WRAPPED.read_bytes() if WRAPPED.exists() else None)
            if wk and not found_key[0]:
                if brute_region(raw, wk, '_oecc104_output'):
                    found_key[0] = True

        elif t == 'token_raw':
            raw = bytes(p['data'])
            print(f'\n[+] {p.get("msg", "token")}')
            (OUT / 'token_raw24.bin').write_bytes(raw)
            # Interpret std::string raw24: try to find inline data (short mode)
            # Short mode (LE LLVM libc++): byte[0] = (size<<1)|0, data at [1..]
            b0 = raw[0]
            if not (b0 & 1):
                sz = (b0 >> 1) & 0x7f
                if 0 < sz <= 22:
                    tok = raw[1:1 + sz]
                    print(f'    Decoded (short): {tok.hex()}')
                    (OUT / 'device_token.bin').write_bytes(tok)
                    wk = wrapped_key[0] or (WRAPPED.read_bytes() if WRAPPED.exists() else None)
                    if wk and not found_key[0]:
                        if try_device_key(tok[:16], wk, 'GetTokenFromKeybox'):
                            found_key[0] = True
            print(f'    raw24 hex: {raw.hex()}')

        elif t == 'oecc107_sess':
            raw = bytes(p['data'])
            print(f'\n[+] _oecc107 session_impl dump: {len(raw)} bytes  head={raw[:32].hex()}')
            (OUT / 'oecc107_sess.bin').write_bytes(raw)
            wk = wrapped_key[0] or (WRAPPED.read_bytes() if WRAPPED.exists() else None)
            if wk and not found_key[0]:
                print('    Brute-forcing device_key (all byte offsets)...')
                if brute_region(raw, wk, 'oecc107_session'):
                    found_key[0] = True

        elif t in ('vtfunc_stack', 'vtfunc_stack_leave'):
            raw = bytes(p['data'])
            label = 'vtfunc_stack' if t == 'vtfunc_stack' else 'vtfunc_stack_LEAVE'
            print(f'\n[+] {label}: {len(raw)} bytes  head={raw[:16].hex()}')
            fname = 'vtfunc_stack_leave.bin' if t == 'vtfunc_stack_leave' else 'vtfunc_stack.bin'
            (OUT / fname).write_bytes(raw)
            wk = wrapped_key[0] or (WRAPPED.read_bytes() if WRAPPED.exists() else None)
            if wk and not found_key[0]:
                print(f'    Brute-forcing device_key from {label}...')
                if brute_region(raw, wk, label):
                    found_key[0] = True

        elif t == 'rsa_i2d_hit':
            raw = bytes(p['data'])
            src = p.get('src', '?')
            print(f'\n[!!!] RSA private key serialised via i2d_RSAPrivateKey! ({len(raw)} bytes) src={src}')
            print(f'    head: {raw[:16].hex()}')
            out = OUT / 'private_key.der'
            out.write_bytes(raw)
            print(f'    Saved → {out}')
            found_key[0] = True

        elif t == 'rsa_struct_dump':
            raw = bytes(p['data'])
            rsa = p.get('rsa', '?')
            src = p.get('src', '?')
            print(f'\n[+] RSA struct (no serialiser): RSA*={rsa} src={src} ({len(raw)} bytes)')
            print(f'    head: {raw[:32].hex()}')
            fname = f'rsa_struct_{src.replace("/","_")}.bin'
            (OUT / fname).write_bytes(raw)
            print(f'    Saved → L3/{fname}')

        elif t == 'rsa_scan_hit':
            raw  = bytes(p['data'])
            addr = p.get('addr', '?')
            tag  = p.get('tag', '?')
            print(f'\n[!!!] RSA DER found in memory scan [{tag}] @ {addr} ({len(raw)} bytes)')
            print(f'    head: {raw[:16].hex()}')
            out = OUT / 'private_key.der'
            out.write_bytes(raw)
            print(f'    Saved → {out}')
            found_key[0] = True

        elif t in ('klsosnbn_after', 'klsosnbn_genrsa'):
            raw  = bytes(p['data'])
            label = 'klsosnbn_after_load' if t == 'klsosnbn_after' else 'klsosnbn_genrsa'
            fname = f'{label}.bin'
            (OUT / fname).write_bytes(raw)
            print(f'\n[+] {label}: {len(raw)} bytes  head={raw[:32].hex()}')
            # Check for RSA DER at various offsets (key might be stored in BSS)
            for off in range(0, min(len(raw) - 4, 3000), 4):
                if raw[off] == 0x30 and raw[off+1] == 0x82:
                    inner = (raw[off+2] << 8) | raw[off+3]
                    if 400 < inner < 5000 and off + 4 + inner <= len(raw):
                        chunk = raw[off:off+4+inner]
                        if len(chunk) > 7 and chunk[4:7] == b'\x02\x01\x00':
                            print(f'    [!!!] RSA DER candidate at klsosnbn+0x{off:x}!')
                            print(f'          head: {chunk[:16].hex()}')
                            (OUT / 'private_key.der').write_bytes(chunk)
                            found_key[0] = True
            wk = wrapped_key[0] or (WRAPPED.read_bytes() if WRAPPED.exists() else None)
            if wk and not found_key[0]:
                if brute_region(raw, wk, label):
                    found_key[0] = True

        elif t in ('csess_after', 'csess_genrsa'):
            raw   = bytes(p['data'])
            label = 'csess_after_load' if t == 'csess_after' else 'csess_genrsa'
            fname = f'{label}.bin'
            (OUT / fname).write_bytes(raw)
            print(f'\n[+] {label}: {len(raw)} bytes  head={raw[:16].hex()}')
            # Compare with the pre-load csess dump to find new RSA* pointer
            pre_path = OUT / 'csess_dump.bin'
            if pre_path.exists() and t == 'csess_after':
                pre = pre_path.read_bytes()
                import struct
                changed = []
                for off in range(0, min(len(pre), len(raw)) - 7, 8):
                    before = struct.unpack_from('<Q', pre, off)[0]
                    after  = struct.unpack_from('<Q', raw, off)[0]
                    if before == 0 and after != 0 and after > 0x7c0000000000:
                        changed.append((off, after))
                if changed:
                    print(f'    New non-null pointers in csess AFTER load (RSA* candidates):')
                    for off, val in changed:
                        print(f'      csess+0x{off:03x} → 0x{val:016x}')
                else:
                    print('    No new pointers vs pre-load dump')
            wk = wrapped_key[0] or (WRAPPED.read_bytes() if WRAPPED.exists() else None)
            if wk and not found_key[0]:
                if brute_region(raw, wk, label):
                    found_key[0] = True

        elif t == 'oecc107_sess_after':
            raw = bytes(p['data'])
            (OUT / 'oecc107_sess_after.bin').write_bytes(raw)
            print(f'\n[+] oecc107_sess AFTER vtable: {len(raw)} bytes  head={raw[:32].hex()}')
            # Compare with BEFORE dump
            pre_path = OUT / 'oecc107_sess.bin'
            if pre_path.exists():
                pre = pre_path.read_bytes()
                import struct
                changed = []
                for off in range(0, min(len(pre), len(raw)) - 7, 8):
                    before = struct.unpack_from('<Q', pre, off)[0]
                    after  = struct.unpack_from('<Q', raw, off)[0]
                    if before != after:
                        changed.append((off, before, after))
                if changed:
                    print(f'    Changed qwords (likely RSA* or related):')
                    for off, b4, af in changed[:20]:
                        print(f'      sess+0x{off:03x}: 0x{b4:016x} → 0x{af:016x}')
                else:
                    print('    No changes vs pre-load session dump')
            wk = wrapped_key[0] or (WRAPPED.read_bytes() if WRAPPED.exists() else None)
            if wk and not found_key[0]:
                if brute_region(raw, wk, 'oecc107_sess_after'):
                    found_key[0] = True

        elif t == 'wpll_raw16':
            raw  = bytes(p['data'])
            seed = p.get('seed', 0x18c)
            print(f'\n[+] wpllqdnz klsosnbn+0x3d0 (raw16): {raw.hex()}  seed=0x{seed:x}')
            (OUT / 'wpll_raw16.bin').write_bytes(raw)
            derived, lcg_high = lcg_derive_key(raw, seed)
            print(f'    LCG-derived key (rsp+0x10): {derived.hex()}')
            print(f'    LCG high bytes  (rsp+0x00): {lcg_high.hex()}')
            wk = wrapped_key[0] or (WRAPPED.read_bytes() if WRAPPED.exists() else None)
            if wk and not found_key[0]:
                for k, kn in [(derived, 'lcg_derived'), (lcg_high, 'lcg_high'), (raw, 'raw16')]:
                    if try_device_key(k, wk, f'wpll_{kn}'):
                        found_key[0] = True
                        break

        elif t == 'wpll_rbx256':
            raw = bytes(p['data'])
            print(f'\n[+] wpllqdnz rbx(klsosnbn+0x338) 256B: head={raw[:16].hex()}')
            (OUT / 'wpll_rbx256.bin').write_bytes(raw)
            wk = wrapped_key[0] or (WRAPPED.read_bytes() if WRAPPED.exists() else None)
            if wk and not found_key[0]:
                if brute_region(raw, wk, 'wpll_rbx256'):
                    found_key[0] = True

        elif t == 'wpll_lcg_stack':
            raw      = bytes(p['data'])
            lcg_high = raw[0:16]
            derived  = raw[16:32]
            print(f'\n[+] wpllqdnz LCG stack (post-loop):')
            print(f'    rsp+0x00 (lcg_high): {lcg_high.hex()}')
            print(f'    rsp+0x10 (derived):  {derived.hex()}')
            (OUT / 'wpll_lcg_stack.bin').write_bytes(raw)
            wk = wrapped_key[0] or (WRAPPED.read_bytes() if WRAPPED.exists() else None)
            if wk and not found_key[0]:
                for k, kn in [(derived, 'lcg_derived_stack'), (lcg_high, 'lcg_high_stack')]:
                    if try_device_key(k, wk, kn):
                        found_key[0] = True
                        break

        elif t == 'wpll_output':
            raw = bytes(p['data'])
            print(f'\n[+] wpllqdnz output buffer (r13 post-decrypt): head={raw[:32].hex()}')
            (OUT / 'wpll_output.bin').write_bytes(raw)
            kt = check_der(raw)
            if kt:
                print(f'    [!!!] RSA DER found directly in r13! ({kt})')
                pad = raw[-1]
                if 1 <= pad <= 16 and raw[-pad:] == bytes([pad]) * pad:
                    raw = raw[:-pad]
                if raw[0:2] == b'\x30\x82':
                    inner = (raw[2] << 8) | raw[3]
                    raw   = raw[:4 + inner]
                (OUT / 'private_key.der').write_bytes(raw)
                print(f'    Saved → L3/private_key.der  ({len(raw)} bytes)')
                found_key[0] = True
            else:
                wk = wrapped_key[0] or (WRAPPED.read_bytes() if WRAPPED.exists() else None)
                if wk and not found_key[0]:
                    if brute_region(raw, wk, 'wpll_output'):
                        found_key[0] = True

        elif t == 'fs_read':
            raw  = bytes(p['data'])
            path = p.get('path', 'unknown')
            name = (path.split('/')[-1] or path).replace('/', '_')
            (OUT / f'keybox_{name}').write_bytes(raw)
            if path not in fs_reads:
                fs_reads[path] = raw
                print(f'\n[+] FS read "{name}": {len(raw)} bytes  {raw.hex()}')

    elif msg['type'] == 'error':
        print(f'[!] JS Error: {msg["description"]}')
        if 'stack' in msg:
            print(f'    {msg["stack"][:300]}')


dev = frida.get_usb_device()
print('[*] Attaching to android.hardware.drm@1.4-service.widevine...')
session = dev.attach('android.hardware.drm@1.4-service.widevine')
script  = session.create_script(JS)
script.on('message', on_message)
script.load()

print('[*] Monitoring (Ctrl+C to stop)...\n')
try:
    while not found_key[0]:
        time.sleep(1)
    print('\n[*] Done! Check L3/private_key.der')
    time.sleep(2)
except KeyboardInterrupt:
    print('\n[!] Stopped.')

if not found_key[0] and WRAPPED.exists():
    print('\n[!] Device key not found automatically.')
    print('    Check L3/oecc104_data.bin and L3/token_raw24.bin for manual analysis.')