#!/usr/bin/env python3
"""
Offline Widevine L3 private key extraction — revised.

Validation strategy: HMAC-SHA256(oem_key, ay64_ciphertext) == ay64.dat6
(No 'KBOX' magic exists in this binary; protobuf-format keybox is used instead.)

Steps:
  1. Brute-force the OEM key in libwvhidl.so using HMAC-SHA256 validation.
  2. Use the OEM key (= HMAC key) to AES-128-CBC decrypt ay64.dat → keybox.
  3. Extract device_key from decrypted keybox.
  4. Try various IV strategies to decrypt wrapped_privkey.bin → RSA DER key.
"""
import hmac as hmacmod, hashlib, struct, sys
from pathlib import Path
from Crypto.Cipher import AES

BASE        = Path(__file__).parent
L3          = BASE / 'L3'
LIBWV       = Path('/tmp/libwvhidl.so')

AY64_ENC    = L3 / 'fs_read_128b.bin'      # 128-byte encrypted keybox
AY64_IV     = bytes.fromhex('41a210cb7f65b9856a9b87090982a290')   # from ay64.dat5
AY64_DAT6   = bytes.fromhex('04ef4d3dbf0663b5e794520e51f6d4e275946d857b1cc703acdbe4b3583c1d36')  # HMAC-SHA256
WRAPPED_KEY = L3 / 'wrapped_privkey.bin'

def hmac_sha256(key: bytes, data: bytes) -> bytes:
    return hmacmod.new(key, data, hashlib.sha256).digest()

def aes_cbc_decrypt(key: bytes, iv: bytes, data: bytes) -> bytes:
    return AES.new(key, AES.MODE_CBC, iv).decrypt(data)


def find_hmac_key(lib: bytes, ciphertext: bytes, expected_mac: bytes):
    """Try every 16-byte-aligned chunk as HMAC-SHA256 key for validation."""

    # Strategy A: HMAC-SHA256(key, ciphertext) == expected
    print(f'[*] Strategy A: HMAC-SHA256(key16, ct128) == dat6  ({len(lib)//16} candidates)...')
    for off in range(0, len(lib) - 16, 16):
        k = lib[off:off + 16]
        if len(set(k)) < 5:
            continue
        if hmac_sha256(k, ciphertext) == expected_mac:
            print(f'    [+] Found (strategy A, off=0x{off:x}): {k.hex()}')
            return k

    # Strategy B: HMAC-SHA256(key, iv || ciphertext) == expected
    print('[*] Strategy B: HMAC-SHA256(key16, iv||ct) == dat6...')
    msg = AY64_IV + ciphertext
    for off in range(0, len(lib) - 16, 16):
        k = lib[off:off + 16]
        if len(set(k)) < 5:
            continue
        if hmac_sha256(k, msg) == expected_mac:
            print(f'    [+] Found (strategy B, off=0x{off:x}): {k.hex()}')
            return k

    # Strategy C: 32-byte candidates
    print('[*] Strategy C: HMAC-SHA256(key32, ct128) == dat6...')
    for off in range(0, len(lib) - 32, 16):
        k = lib[off:off + 32]
        if len(set(k)) < 8:
            continue
        if hmac_sha256(k, ciphertext) == expected_mac:
            print(f'    [+] Found (strategy C 32-byte, off=0x{off:x}): {k.hex()}')
            return k

    # Strategy D: HMAC of plaintext (try AES-128-CBC first with each candidate as BOTH key)
    print('[*] Strategy D: AES-CBC decrypt then HMAC-SHA256(same_key, plaintext) == dat6...')
    for off in range(0, len(lib) - 16, 16):
        k = lib[off:off + 16]
        if len(set(k)) < 5:
            continue
        plain = aes_cbc_decrypt(k, AY64_IV, ciphertext)
        if hmac_sha256(k, plain) == expected_mac:
            print(f'    [+] Found (strategy D, off=0x{off:x}): {k.hex()}')
            return k

    return None


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 strip_der(data: bytes) -> bytes:
    """Trim to actual DER length (strip PKCS#7 padding if present)."""
    pad = data[-1]
    if 1 <= pad <= 16 and data[-pad:] == bytes([pad]) * pad:
        data = data[:-pad]
    if data[0:2] == b'\x30\x82':
        inner = (data[2] << 8) | data[3]
        return data[:4 + inner]
    return data


# ── Main ──────────────────────────────────────────────────────────────────────
print('[*] Widevine L3 offline key extractor (HMAC-SHA256 validation)')
print()

for p in (LIBWV, AY64_ENC, WRAPPED_KEY):
    if not p.exists():
        print(f'[!] Missing: {p}')
        if p == LIBWV:
            print('    adb pull /vendor/lib64/libwvhidl.so /tmp/libwvhidl.so')
        sys.exit(1)
    print(f'[OK] {p}  ({p.stat().st_size:,} bytes)')
print(f'[OK] ay64.dat6 (hardcoded): {AY64_DAT6.hex()}')
print()

lib_bytes   = LIBWV.read_bytes()
ay64_cipher = AY64_ENC.read_bytes()
wrapped     = WRAPPED_KEY.read_bytes()

# ── Step 1: Find OEM key / HMAC key ──────────────────────────────────────────
hmac_key = find_hmac_key(lib_bytes, ay64_cipher, AY64_DAT6)

if hmac_key is None:
    print('\n[!] HMAC-SHA256 key not found in libwvhidl.so.')
    print('[!] The key may be derived at runtime from device-specific data.')
    print('[!] Falling back to Frida runtime hook (GetTokenFromKeybox).')
    print('    Run: python3 dump_wvd2.py  then trigger DRM in browser')
    sys.exit(1)

print(f'\n[+] HMAC key:       {hmac_key.hex()}')

# ── Step 2: Decrypt keybox ────────────────────────────────────────────────────
# Use the HMAC key as the AES key too (common in Widevine L3 implementations)
keybox = aes_cbc_decrypt(hmac_key, AY64_IV, ay64_cipher)
print(f'[+] Decrypted keybox (128 bytes): {keybox.hex()}')

# Parse protobuf keybox: find a 16-byte AES key embedded in it
# Protobuf field 2 (wire type 2 = LEN) likely contains the device key
# Also try raw struct layout just in case
possible_device_keys = set()

# Try every 16-byte-aligned block in the decrypted keybox as a device key candidate
for off in range(0, 113, 16):
    possible_device_keys.add(keybox[off:off + 16])

# Protobuf parse: look for 2-byte length-prefixed blobs of length 16
for i in range(len(keybox) - 2):
    tag = keybox[i]
    wire = tag & 0x07
    if wire == 2:  # length-delimited
        length = keybox[i + 1]
        if length == 16 and i + 2 + 16 <= len(keybox):
            possible_device_keys.add(keybox[i + 2:i + 18])

print(f'[*] {len(possible_device_keys)} device_key candidates from decrypted keybox')

# ── Step 3: Decrypt wrapped_privkey.bin ───────────────────────────────────────
print(f'\n[*] Trying to decrypt wrapped_privkey.bin ({len(wrapped)} bytes)...')

ivs = [
    ('IV=0x00*16',       b'\x00' * 16),
    ('IV=wrapped[0:16]', wrapped[:16]),
    ('IV=ay64.dat5',     AY64_IV),
    ('IV=hmac_key',      hmac_key if len(hmac_key) == 16 else hmac_key[:16]),
]
cts = [
    ('ct=all',           wrapped),
    ('ct=wrapped[16:]',  wrapped[16:]),
]

found = False
for device_key in possible_device_keys:
    for iv_name, iv in ivs:
        for ct_name, ct in cts:
            if len(ct) % 16 != 0 or len(ct) == 0:
                continue
            plain = aes_cbc_decrypt(device_key, iv, ct)
            kt = check_der(plain)
            if kt:
                print(f'\n[+] RSA KEY FOUND!')
                print(f'    device_key: {device_key.hex()}')
                print(f'    {iv_name}, {ct_name}')
                key_bytes = strip_der(plain)
                out = L3 / 'private_key.der'
                out.write_bytes(key_bytes)
                print(f'    Saved: {out}  ({len(key_bytes)} bytes)')
                found = True
                break
        if found:
            break
    if found:
        break

if not found:
    print('\n[!] No RSA key found. Plaintext samples (first 8 bytes per combo):')
    first_key = next(iter(possible_device_keys))
    for iv_name, iv in ivs:
        for ct_name, ct in cts:
            if len(ct) % 16 != 0:
                continue
            plain = aes_cbc_decrypt(first_key, iv, ct)
            print(f'    [{iv_name} {ct_name}] {plain[:16].hex()}')
    print()
    print('[*] Next step: use Frida GetTokenFromKeybox hook to get device_key at runtime')
    print('    python3 dump_wvd2.py')

if found:
    print()
    print('[*] Next steps:')
    print('  1. Get client_id.bin — run dump_wvd2.py while triggering DRM,')
    print('     then capture from GenerateRsaSignature output')
    print('  2. pywidevine create-device -t android -l 3 \\')
    print(f'       -k {L3}/private_key.der -c L3/client_id.bin -o L3/')