"""Build an unsigned, logging-only diagnostic APK for the exact Quest Unity binary.

Preserves certificate decisions. Adds Android log probes after native X509
verification and after the Unity verification callback. Never edits the input.
Run zipalign and sign with the installed test application's key afterward.
"""
import argparse
import hashlib
import struct
import zipfile
from pathlib import Path

EXPECTED = 'd68dd33530dc9c6f963da17cad03d56fbaF51ac4bef297d05bf56684a44ee607'.lower()
CAVE = 0x125AE00
LIMIT = 0x125B100
LOG_PRINT = 0xBA9F0

def word(v):
    return struct.pack('<I', v)

def branch(pc, target, link=False):
    delta = target - pc
    assert delta % 4 == 0 and -(1 << 27) <= delta < (1 << 27)
    return word((0x94000000 if link else 0x14000000) | ((delta >> 2) & 0x3FFFFFF))

def adr(pc, target, reg):
    delta = target - pc
    assert -(1 << 20) <= delta < (1 << 20)
    return word(0x10000000 | ((delta & 3) << 29) | (((delta >> 2) & 0x7FFFF) << 5) | reg)

def mov(dst, src, wide=True):
    return word((0xAA0003E0 if wide else 0x2A0003E0) | (src << 16) | dst)

def pair(a, b, offset, load=False):
    return word((0xA9400000 if load else 0xA9000000) | ((offset // 8) << 15) | (b << 10) | (31 << 5) | a)

def patch(binary, radicade_system_trust=False):
    assert hashlib.sha256(binary).hexdigest() == EXPECTED, 'Unknown Unity binary'
    out = bytearray(binary)
    assert not any(out[CAVE:LIMIT]), 'Diagnostic space is occupied'
    # Existing RX PT_LOAD ends before alignment padding; extend only into that
    # verified empty padding, staying before the following RW PT_LOAD file range.
    ph = struct.unpack_from('<Q', out, 32)[0]
    assert struct.unpack_from('<II', out, ph) == (1, 5)
    assert struct.unpack_from('<QQ', out, ph + 32) == (0x125ADF4, 0x125ADF4)
    struct.pack_into('<QQ', out, ph + 32, LIMIT, LIMIT)
    tag = CAVE + 0x240
    fmt1 = tag + 16
    fmt2 = fmt1 + 96
    strings = {
        tag: b'RadicadeTLS\0',
        fmt1: b'x509 result=%d raw_flags=0x%x unity_flags=0x%x hostname=%s callback=%p\0',
        fmt2: b'callback result=0x%x callback=%p error_code=%u\0',
    }
    for addr, data in strings.items():
        assert addr + len(data) <= LIMIT
        out[addr:addr + len(data)] = data

    probes = [(0x75BA0C, CAVE, fmt1), (0x75F3B8, CAVE + 0x120, fmt2)]
    for hook, start, fmt in probes:
        code = bytearray(word(0xD10283FF))  # sub sp, sp, #160
        for reg in range(0, 18, 2):
            code += pair(reg, reg + 1, reg * 8)
        code += pair(18, 30, 144)
        if hook == 0x75BA0C:
            assert binary[hook:hook+4] == word(0x1284DFE8) # mov w8, #-0x2700
            code += mov(3, 21, False)
            code += word(0xB8400000 | ((-0x5C & 0x1FF) << 12) | (29 << 5) | 4)
            code += mov(5, 20, False) + mov(6, 28) + mov(7, 25)
        else:
            assert binary[hook:hook+4] == word(0xB90AC280) # str w0, [x20,#0xac0]
            code += mov(3, 0, False) + mov(4, 24)
            # Original errorstate at pre-probe SP+0x10; code is its +4 field.
            code += word(0xB9400000 | ((180 // 4) << 10) | (31 << 5) | 5)
        code += word(0x52800080) # mov w0, #4 (INFO)
        code += adr(start + len(code), tag, 1)
        code += adr(start + len(code), fmt, 2)
        code += branch(start + len(code), LOG_PRINT, True)
        for reg in range(0, 18, 2):
            code += pair(reg, reg + 1, reg * 8, True)
        code += pair(18, 30, 144, True)
        code += word(0x910283FF) # add sp, sp, #160
        code += binary[hook:hook+4]
        code += branch(start + len(code), hook + 4)
        assert len(code) <= 0x120
        out[start:start+len(code)] = code
        out[hook:hook+4] = branch(hook, start)
    if radicade_system_trust:
        # Route only this exact hostname through Unity's existing platform-root
        # verifier. It checks the complete chain AND hostname. All other hosts
        # retain the configured managed callback. No accept-all return patch.
        start = CAVE + 0x90
        host = CAVE + 0x2E0
        hook = 0x75F3B4
        assert binary[hook:hook+4] == word(0xD63F0300) # blr x24
        hostname = b'ns.radicade.app\0'
        out[host:host+len(hostname)] = hostname
        code = bytearray(word(0xD10083FF)) # sub sp, sp, #32
        code += pair(0, 1, 0) + pair(2, 30, 16)
        code += word(0xF9400000 | ((0xA90 // 8) << 10) | (20 << 5))
        # A null hostname must keep the original path.
        null_branch = len(code)
        code += word(0)
        code += adr(start + len(code), host, 1)
        code += branch(start + len(code), 0xBA7D0, True) # strcmp
        mismatch_branch = len(code)
        code += word(0)
        code += pair(0, 1, 0, True) + pair(2, 30, 16, True)
        code += word(0x910083FF) + mov(0, 20)
        code += branch(start + len(code), 0x75F5D0, True)
        code += branch(start + len(code), hook + 4)
        original_path = start + len(code)
        code += pair(0, 1, 0, True) + pair(2, 30, 16, True)
        code += word(0x910083FF) + binary[hook:hook+4]
        code += branch(start + len(code), hook + 4)
        struct.pack_into('<I', code, null_branch, 0xB4000000 | (((original_path - start - null_branch) // 4) << 5))
        struct.pack_into('<I', code, mismatch_branch, 0x35000000 | (((original_path - start - mismatch_branch) // 4) << 5))
        assert len(code) <= 0x90
        out[start:start+len(code)] = code
        out[hook:hook+4] = branch(hook, start)
    return bytes(out)

def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument('input', type=Path)
    p.add_argument('output', type=Path)
    p.add_argument('--radicade-system-trust', action='store_true', help='Use Unity platform certificate+hostname verification for the exact Radicade host')
    args = p.parse_args()
    if args.output.exists():
        raise SystemExit('Output already exists; choose a new file')
    args.output.parent.mkdir(parents=True, exist_ok=True)
    with zipfile.ZipFile(args.input) as source:
        original = source.read('lib/arm64-v8a/libunity.so')
        modified = patch(original, args.radicade_system_trust)
        with zipfile.ZipFile(args.output, 'w') as target:
            for entry in source.infolist():
                if entry.filename.upper().startswith('META-INF/') and entry.filename.upper().endswith(('.RSA', '.SF', '.MF', '.DSA', '.EC')):
                    continue
                data = modified if entry.filename == 'lib/arm64-v8a/libunity.so' else source.read(entry)
                target.writestr(entry, data)
    print('Unsigned diagnostic APK:', args.output)
    print('Patched Unity SHA256:', hashlib.sha256(modified).hexdigest())

if __name__ == '__main__':
    main()