#!/usr/bin/env python3
"""
Self-trigger approach: attach to the DRM HAL process, hook RSA_private_key_from_bytes
(fires when the RSA key is LOADED from the keybox), and also call Level3_GenerateRSASignature
directly with dummy data to force a signing operation without needing any external trigger.
"""
import frida, time, sys
from pathlib import Path

OUT = Path(__file__).parent / "L3"

JS = r"""
'use strict';

var libcrypto = Process.findModuleByName('libcrypto.so');
var libwv      = Process.findModuleByName('libwvhidl.so');

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

var rsaToBytesAddr      = findFn(libcrypto, 'RSA_private_key_to_bytes');
var i2dRsaAddr          = findFn(libcrypto, 'i2d_RSAPrivateKey');
var rsaSignAddr         = findFn(libcrypto, 'RSA_sign');
var rsaSignRawAddr      = findFn(libcrypto, 'RSA_sign_raw');
var rsaPrivEncAddr      = findFn(libcrypto, 'RSA_private_encrypt');
var rsaFromBytesAddr    = findFn(libcrypto, 'RSA_private_key_from_bytes');
var l3SigAddr           = findFn(libwv,     'Level3_GenerateRSASignature');

send({type:'addrs',
    rsaToBytes:   String(rsaToBytesAddr),
    i2dRsa:       String(i2dRsaAddr),
    rsaSign:      String(rsaSignAddr),
    rsaSignRaw:   String(rsaSignRawAddr),
    rsaPrivEnc:   String(rsaPrivEncAddr),
    rsaFromBytes: String(rsaFromBytesAddr),
    l3Sig:        String(l3SigAddr)
});

var alreadySent = false;

// ── Serializer ────────────────────────────────────────────────────────────────
var dumpRSAKey;
if (rsaToBytesAddr) {
    var rsaToBytes = new NativeFunction(rsaToBytesAddr, 'int', ['pointer', 'pointer', 'pointer']);
    dumpRSAKey = function(rsa, label) {
        if (alreadySent || !rsa || rsa.isNull()) return;
        var pBuf = Memory.alloc(8), pLen = Memory.alloc(8);
        pBuf.writePointer(ptr('0x0')); pLen.writeU64(0);
        var ok = rsaToBytes(pBuf, pLen, rsa);
        if (ok !== 1) { send({type:'log', msg:'RSA_private_key_to_bytes failed: '+ok}); 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 (i2dRsaAddr) {
    var i2dFn = new NativeFunction(i2dRsaAddr, 'int', ['pointer', 'pointer']);
    dumpRSAKey = function(rsa, label) {
        if (alreadySent || !rsa || rsa.isNull()) return;
        var ppBuf = Memory.alloc(8); ppBuf.writePointer(ptr('0x0'));
        var len = i2dFn(rsa, ppBuf);
        if (len <= 0) return;
        var buf = ppBuf.readPointer(); if (buf.isNull()) return;
        alreadySent = true;
        send({type:'private_key', src:label, data: Array.from(new Uint8Array(buf.readByteArray(len)))});
    };
} else {
    dumpRSAKey = function(rsa, label) { send({type:'log', msg: label + ' fired - no serializer'}); };
}

// ── Hook: RSA_private_key_from_bytes(in, in_len) → RSA*
// Fires when the DRM service LOADS the private key from the decrypted keybox.
// We capture the raw DER bytes directly from the argument.
if (rsaFromBytesAddr) {
    var rsaFromBytes = new NativeFunction(rsaFromBytesAddr, 'pointer', ['pointer', 'pointer']);
    Interceptor.attach(rsaFromBytesAddr, {
        onEnter: function(args) {
            this.inPtr = args[0];
            // args[1] is pointer-to-length (CBS struct in BoringSSL)
            // Try to read the length from the CBS struct (first field is data pointer, second is len)
            try {
                this.inLen = args[1].readU64().toNumber();
            } catch(e) { this.inLen = 0; }
        },
        onLeave: function(retval) {
            if (retval.isNull()) return;
            send({type:'log', msg:'RSA_private_key_from_bytes fired! RSA*='+retval+' inLen='+this.inLen});
            dumpRSAKey(retval, 'RSA_private_key_from_bytes');
        }
    });
    send({type:'log', msg:'Hooked RSA_private_key_from_bytes'});
}

// ── Hook: RSA_sign(hash_nid, in, in_len, out, out_len, RSA*)
if (rsaSignAddr) {
    Interceptor.attach(rsaSignAddr, {
        onEnter: function(a) { this.rsa = a[5]; },
        onLeave: function(r) { if (r.toInt32()===1) dumpRSAKey(this.rsa,'RSA_sign'); }
    });
    send({type:'log', msg:'Hooked RSA_sign'});
}

// ── Hook: RSA_sign_raw(RSA*, out_len*, out, max_out, in, in_len, padding)
if (rsaSignRawAddr) {
    Interceptor.attach(rsaSignRawAddr, {
        onEnter: function(a) { this.rsa = a[0]; },
        onLeave: function(r) { if (r.toInt32()===1) dumpRSAKey(this.rsa,'RSA_sign_raw'); }
    });
    send({type:'log', msg:'Hooked RSA_sign_raw'});
}

// ── Hook: RSA_private_encrypt
if (rsaPrivEncAddr) {
    Interceptor.attach(rsaPrivEncAddr, {
        onEnter: function(a) { this.rsa = a[3]; },
        onLeave: function(r) { if (r.toInt32()>0) dumpRSAKey(this.rsa,'RSA_private_encrypt'); }
    });
    send({type:'log', msg:'Hooked RSA_private_encrypt'});
}

// ── Hook: Level3_GenerateRSASignature — log args when it fires naturally
if (l3SigAddr) {
    Interceptor.attach(l3SigAddr, {
        onEnter: function(args) {
            send({type:'log', msg:'Level3_GenerateRSASignature CALLED!'});
            for (var i = 0; i < 6; i++) {
                var v = args[i];
                var preview = '';
                try { preview = ' -> ' + hexify(v.readByteArray(8)); } catch(e) {}
                send({type:'log', msg:'  arg['+i+']='+v+preview});
            }
        },
        onLeave: function(r) {
            send({type:'log', msg:'Level3_GenerateRSASignature returned: ' + r});
        }
    });
    send({type:'log', msg:'Hooked Level3_GenerateRSASignature'});

    // ── Self-call: try to invoke Level3_GenerateRSASignature with dummy data
    // We don't know the exact signature, so try two common patterns after a short delay.
    setTimeout(function() {
        send({type:'log', msg:'Attempting self-call of Level3_GenerateRSASignature...'});
        try {
            // Pattern 1: (msg_ptr, msg_len, sig_ptr, sig_len_ptr)
            var msg    = Memory.alloc(32);
            var sigBuf = Memory.alloc(512);
            var sigLen = Memory.alloc(8);
            msg.writeByteArray([0x00,0x01,0x02,0x03,0x04,0x05,0x06,0x07,
                                 0x08,0x09,0x0a,0x0b,0x0c,0x0d,0x0e,0x0f,
                                 0x10,0x11,0x12,0x13,0x14,0x15,0x16,0x17,
                                 0x18,0x19,0x1a,0x1b,0x1c,0x1d,0x1e,0x1f]);
            sigLen.writeU64(512);
            var fn1 = new NativeFunction(l3SigAddr, 'int',
                ['pointer','uint64','pointer','pointer']);
            var r1 = fn1(msg, 32, sigBuf, sigLen);
            send({type:'log', msg:'Pattern 1 returned: ' + r1 + ' sigLen=' + sigLen.readU64()});
        } catch(e) {
            send({type:'log', msg:'Pattern 1 threw: ' + e});
        }

        try {
            // Pattern 2: (msg_ptr, msg_len, key_id_ptr, key_id_len, sig_ptr, sig_len_ptr)
            var msg2    = Memory.alloc(32);
            var keyId   = Memory.alloc(16);
            var sigBuf2 = Memory.alloc(512);
            var sigLen2 = Memory.alloc(8);
            sigLen2.writeU64(512);
            var fn2 = new NativeFunction(l3SigAddr, 'int',
                ['pointer','uint64','pointer','uint64','pointer','pointer']);
            var r2 = fn2(msg2, 32, keyId, 16, sigBuf2, sigLen2);
            send({type:'log', msg:'Pattern 2 returned: ' + r2});
        } catch(e) {
            send({type:'log', msg:'Pattern 2 threw: ' + e});
        }
    }, 1000);
}

send({type:'ready'});

function hexify(arr) {
    return Array.from(new Uint8Array(arr))
        .map(function(b){ return ('0'+b.toString(16)).slice(-2); }).join('');
}
"""

captured = [False]

def on_message(msg, data):
    if msg['type'] == 'send':
        p = msg['payload']
        t = p.get('type')
        if t == 'addrs':
            print('[*] Addresses:')
            for k, v in p.items():
                if k != 'type':
                    print(f'    {k:18s} = {v}')
        elif t == 'ready':
            print('[*] All hooks installed. Waiting for RSA activity...')
            print('    (also attempting self-call in 1 second)\n')
        elif t == 'log':
            print(f'    {p["msg"]}')
        elif t == 'private_key' and not captured[0]:
            raw = bytes(p['data'])
            if len(raw) > 100:
                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: {p.get("src")}')
                print(f'    Header: {raw[:8].hex()}')
    elif msg['type'] == 'error':
        print(f'[!] JS Error: {msg["description"]}')
        if 'stack' in msg:
            print(f'    {msg["stack"]}')

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('[*] Waiting... (Ctrl+C to quit)\n')
try:
    while not captured[0]:
        time.sleep(1)
    print('[+] Done!')
except KeyboardInterrupt:
    print('\n[!] Stopped.')