#!/usr/bin/env python3
"""
Capture the Widevine ClientIdentification by getting a license challenge
from MediaDrm and parsing out the embedded client_id.

Run while the emulator is running with frida-server active.
Saves L3/client_id.bin on success.
"""
import frida, sys, time, os
from pathlib import Path

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

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

JS = r"""
'use strict';
send({type:'log', msg:'Script loaded in PID ' + Process.id});

var MAX_ATTEMPTS = 30;
var attempt = 0;

function tryJava() {
    attempt++;
    var avail = (typeof Java !== 'undefined') && Java.available;
    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);

            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:'Calling getKeyRequest...'});
            var req = drm.getKeyRequest(session, pssh, 'video/mp4', 1, null);
            var challenge = req.getData();
            send({type:'log', msg:'Challenge: ' + challenge.length + ' bytes'});

            // Send full bytes
            var arr = Array.from(challenge);
            send({type:'challenge', data: arr});

            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':
            raw = bytes(p['data'])
            print(f'[+] Challenge: {len(raw)} bytes  head={raw[:16].hex()}')
            # Save raw challenge
            (OUT / 'challenge.bin').write_bytes(raw)
            print(f'    Saved challenge → L3/challenge.bin')
            # Parse to extract ClientIdentification
            try:
                from pywidevine import license_protocol_pb2 as pb
                sm = pb.SignedMessage()
                sm.MergeFromString(raw)
                print(f'    SignedMessage.type={sm.type}  msg={len(sm.msg)} bytes  sig={len(sm.signature)} bytes')
                lr = pb.LicenseRequest()
                lr.MergeFromString(sm.msg)
                ci = lr.client_id
                ci_bytes = ci.SerializeToString()
                print(f'    ClientIdentification: type={ci.type}  token={len(ci.token)} bytes')
                print(f'    ClientIdentification total: {len(ci_bytes)} bytes  head={ci_bytes[:16].hex()}')
                if len(ci_bytes) > 50:
                    (OUT / 'client_id.bin').write_bytes(ci_bytes)
                    print(f'[+] Saved → L3/client_id.bin')
                else:
                    print(f'[!] ClientIdentification too small ({len(ci_bytes)} bytes) — may be empty')
                    print(f'    Full LicenseRequest:\n{lr}')
            except Exception as e:
                print(f'[!] Parse error: {e}')
        elif t == 'done':
            print('[*] MediaDrm complete.')
            sys.exit(0)
        elif t == 'error':
            print(f'[!] Error: {p["msg"]}')
    elif msg['type'] == 'error':
        print(f'[!] Frida JS error: {msg["description"]}')

dev = frida.get_usb_device()

PREFERRED = [
    'com.google.android.gms.persistent',
    'com.google.android.gms',
    'com.google.android.gms.unstable',
]

print('[*] Finding GMS process...')
procs = {p.name: p.pid for p in dev.enumerate_processes()}
session = None
for name in PREFERRED:
    pid = procs.get(name)
    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}')

if session is None:
    # Fallback: try any process with Java
    print('[*] GMS not found by preferred name. Trying other processes...')
    for name, pid in procs.items():
        if name in ('frida-server', 'adbd', 'zygote', 'zygote64', 'init', 'surfaceflinger'):
            continue
        if not any(k in name for k in ('google', 'android', 'gms')):
            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('[*] Waiting for license challenge...\n')

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