#!/usr/bin/env python3
"""
Trigger a Widevine MediaDrm key request from within a GMS Java process.
This causes the Android DRM HAL to call Level3_GenerateRSASignature, which
dump_wvd2.py intercepts to extract the RSA private key via memory scan.

Usage:
  Terminal A: python3 dump_wvd2.py          (hook DRM HAL)
  Terminal B: python3 trigger_mediadrm.py   (trigger the DRM call)
"""
import frida, sys, time

PSSH_B64 = 'AAAAOnBzc2gAAAAA7e+LqXnWSs6jyCfc1R0h7QAAABoIARIQ1DG5TL5NQ/2zoimqi0ixiUjj3JWbBg=='

JS = r"""
'use strict';

// Log immediately so we know the script loaded
send({type:'log', msg: 'Script loaded in PID ' + Process.id});
send({type:'log', msg: 'typeof Java = ' + typeof Java});

// Retry loop: wait up to 5 seconds for Java to become available
var MAX_ATTEMPTS = 25;
var attempt = 0;

function tryJava() {
    attempt++;
    var avail = (typeof Java !== 'undefined') && Java.available;
    send({type:'log', msg: 'attempt ' + attempt + ': Java.available = ' + avail});

    if (avail) {
        runMediaDrm();
    } else if (attempt < MAX_ATTEMPTS) {
        setTimeout(tryJava, 200);
    } else {
        send({type:'error', msg: 'Java not available after ' + MAX_ATTEMPTS + ' attempts'});
    }
}

function runMediaDrm() {
    Java.perform(function() {
        try {
            var UUID     = Java.use('java.util.UUID');
            var MediaDrm = Java.use('android.media.MediaDrm');
            var Base64   = Java.use('android.util.Base64');

            var wvUUID = UUID.fromString('edef8ba9-79d6-4ace-a3c9-27dcd51d21ed');
            send({type:'log', msg: 'Creating MediaDrm...'});
            var drm = MediaDrm.$new(wvUUID);

            send({type:'log', msg: 'Opening session...'});
            var session = drm.openSession();
            send({type:'log', msg: 'Session opened: ' + session.length + ' bytes'});

            var pssh = Base64.decode('""" + PSSH_B64 + r"""', 0);
            send({type:'log', msg: 'PSSH: ' + pssh.length + ' bytes'});

            send({type:'log', msg: 'Calling getKeyRequest → triggers Level3_GenerateRSASignature...'});
            var req = drm.getKeyRequest(session, pssh, 'video/mp4', 1 /*STREAMING*/, null);
            var challenge = req.getData();

            function toHex(arr) {
                var h = '';
                for (var i = 0; i < Math.min(arr.length, 32); i++)
                    h += ('0' + (arr[i] & 0xff).toString(16)).slice(-2);
                return h;
            }

            send({type:'challenge', len: challenge.length, hex: toHex(challenge)});
            drm.closeSession(session);
            drm.release();
            send({type:'done'});

        } catch(e) {
            send({type:'error', msg: e.toString()});
        }
    });
}

setTimeout(tryJava, 100);
"""

def on_message(msg, data):
    if msg['type'] == 'send':
        p = msg['payload']
        t = p.get('type')
        if t == 'log':
            print(f'    {p["msg"]}')
        elif t == 'challenge':
            print(f'\n[+] License challenge: {p["len"]} bytes')
            print(f'    First 32 bytes: {p["hex"]}')
            if p['len'] > 100:
                print('[+] Challenge looks valid! Check dump_wvd2.py output for the RSA key.')
            else:
                print('[!] Challenge too small — CDM might not be provisioned')
        elif t == 'done':
            print('[*] MediaDrm call complete.')
            sys.exit(0)
        elif t == 'error':
            print(f'[!] Error: {p["msg"]}')
            sys.exit(1)
    elif msg['type'] == 'error':
        print(f'[!] Frida JS error: {msg["description"]}')

dev = frida.get_usb_device()

# GMS persistent process — always running, has Java VM, runs as system user
PREFERRED = [
    ('com.google.android.gms.persistent', None),
    ('com.google.android.gms',            None),
    ('com.google.android.gms.unstable',   None),
]

print('[*] Finding target process...')
procs = {p.name: p.pid for p in dev.enumerate_processes()}
for name, _ in PREFERRED:
    short = name.split('.')[-1] if '.' in name else name
    # match full name or short name
    pid = procs.get(name) or procs.get(short)
    if pid:
        print(f'[*] Found: {name} PID={pid}')
        try:
            session = dev.attach(pid)
            print(f'[*] Attached.')
            break
        except Exception as e:
            print(f'    Attach failed: {e}')
            session = None
else:
    # Last resort: try everything
    print('[*] GMS not found by name. Scanning all processes...')
    session = None
    for name, pid in procs.items():
        if name in ('frida-server', 'adbd', 'zygote', 'zygote64', 'init'):
            continue
        try:
            session = dev.attach(pid)
            print(f'[*] Attached to {name} (PID {pid})')
            break
        except Exception:
            pass

if session is None:
    print('[!] Could not attach to any process. Is frida-server running?')
    sys.exit(1)

script = session.create_script(JS)
script.on('message', on_message)
script.load()
print('[*] Script loaded. Waiting for MediaDrm call...\n')

try:
    while True:
        time.sleep(0.5)
except KeyboardInterrupt:
    print('\n[!] Stopped.')