#!/usr/bin/env python3
"""U30: real loopback OAuth2/OIDC AS + client + JWT/introspection resource servers."""
import argparse
import base64
import contextlib
from datetime import datetime
import hashlib
import http.cookiejar
import json
import os
from pathlib import Path
import platform
import re
import secrets
import signal
import socket
import subprocess
import tempfile
import time
import urllib.error
import urllib.parse as url
import urllib.request
import zipfile

ROOT = Path(__file__).resolve().parent
CACHE = Path.home() / '.cache/software-engineering-study/java-depth-impl/u30'
RESULTS, PROCESSES = [], []
CONTROL = secrets.token_hex(24)
LOGS = ENV = JAVA = None


def now(): return datetime.now().astimezone().isoformat()
def require(value, message):
    if not value: raise AssertionError(message)


def stop(p):
    if p.poll() is None:
        os.killpg(p.pid, signal.SIGTERM)
        try: p.wait(timeout=8)
        except subprocess.TimeoutExpired:
            os.killpg(p.pid, signal.SIGKILL)
            p.wait(timeout=5)
    return p.returncode


def command(name, args, timeout=240):
    log = LOGS / (name + '.log')
    with log.open('w') as f:
        p = subprocess.Popen([str(a) for a in args], cwd=ROOT, env=ENV, stdout=f, stderr=subprocess.STDOUT, start_new_session=True)
        try: code = p.wait(timeout=timeout)
        except BaseException:
            stop(p)
            raise
    RESULTS.append({'name': name, 'exitCode': code, 'log': str(log), 'logSha256': hashlib.sha256(log.read_bytes()).hexdigest()})
    require(code == 0, name + ' failed; see ' + str(log))
    return log.read_text()


@contextlib.contextmanager
def server(name, port, issuer, client, rs):
    log = LOGS / (name + '.log')
    profile = 'rs' if name.startswith('rs') else name
    args = [str(JAVA / 'bin/java'), '-jar', str(ROOT / 'target/oauth2-authorization-1.jar'),
            '--spring.profiles.active=' + profile, '--server.port=' + str(port),
            '--server.servlet.session.cookie.name=' + profile.upper() + '_SESSION',
            '--lab.issuer=' + issuer, '--lab.client=' + client, '--lab.rs=' + rs]
    with log.open('w') as f:
        p = subprocess.Popen(args, cwd=ROOT, env=ENV, stdout=f, stderr=subprocess.STDOUT, start_new_session=True)
        try:
            deadline = time.monotonic() + 45
            while time.monotonic() < deadline:
                found = re.search(r'U30BOUND (\d+)', log.read_text())
                if found: break
                require(p.poll() is None, name + ' stopped; see ' + str(log))
                time.sleep(.05)
            else: raise AssertionError(name + ' readiness timeout; see ' + str(log))
            yield 'http://127.0.0.1:' + found[1]
        finally:
            code = stop(p)
            PROCESSES.append({'name': name, 'pid': p.pid, 'exitCode': code, 'reaped': p.poll() is not None,
                              'log': str(log), 'logSha256': hashlib.sha256(log.read_bytes()).hexdigest()})


class NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl): return None


class Browser:
    def __init__(self):
        self.jar = http.cookiejar.CookieJar()
        self.opener = urllib.request.build_opener(urllib.request.ProxyHandler({}), urllib.request.HTTPCookieProcessor(self.jar), NoRedirect())

    def request(self, address, expected, label, fields=None, bearer=None, basic=False, control=False, headers=None):
        h = dict(headers or {})
        if bearer: h['Authorization'] = 'Bearer ' + bearer
        if basic: h['Authorization'] = 'Basic ' + base64.b64encode((basic if isinstance(basic,str) else 'lab-client:lab-secret').encode()).decode()
        if control: h['X-Lab-Control'] = CONTROL
        req = urllib.request.Request(address, data=url.urlencode(fields).encode() if fields is not None else None, headers=h)
        try: response = self.opener.open(req, timeout=12)
        except urllib.error.HTTPError as e: response = e
        with response:
            code, rh, raw = response.code, response.headers, response.read(1000000).decode()
        parsed = url.urlsplit(address)
        # No raw authorization/cookies, query values, response token/code/state/nonce/verifier or private keys.
        RESULTS.append({'name': label, 'method': 'POST' if fields is not None else 'GET',
                        'path': parsed.path, 'queryKeys': sorted(url.parse_qs(parsed.query)),
                        'formFields': sorted(fields or {}), 'bearer': 'present-redacted' if bearer else 'absent',
                        'status': code, 'expectedStatus': expected, 'observedAt': now(),
                        'locationPath': url.urlsplit(rh.get('Location', '')).path,
                        'challengeScheme': rh.get('WWW-Authenticate', '').split(' ')[0]})
        require(code in expected if isinstance(expected, list) else code == expected,
                f'{label}: expected {expected}, actual {code}; response fields omitted')
        return rh, raw


def obj(result): return json.loads(result[1])
def location(result):
    require(bool(result[0].get('Location')), 'redirect Location required')
    return result[0]['Location']

def control(as_url, **fields):
    return Browser().request(as_url + '/lab/control', 200, 'set synthetic fault/rotation control', fields=fields, control=True)

def mint(as_url, kind='valid'):
    return obj(Browser().request(as_url + '/lab/mint', 200, 'mint synthetic ' + kind, fields={'kind': kind}, control=True))['token']

def pkce():
    verifier = secrets.token_urlsafe(48)
    return verifier, base64.urlsafe_b64encode(hashlib.sha256(verifier.encode()).digest()).rstrip(b'=').decode()


def login(browser, as_url):
    csrf = obj(browser.request(as_url + '/lab/csrf', 200, 'fetch AS login CSRF'))
    browser.request(as_url + '/login', 302, 'real AS form authentication', fields={'username': 'alice', 'password': 'lab-password', csrf['parameterName']: csrf['token']})


def authorize(as_url, browser=None, challenge=True, method='S256'):
    browser = browser or Browser()
    login(browser, as_url)
    verifier, hashed = pkce()
    q = {'response_type': 'code', 'client_id': 'lab-client', 'redirect_uri': 'http://127.0.0.1/callback',
         'scope': 'openid profile orders.read orders.write', 'state': secrets.token_urlsafe(24), 'nonce': secrets.token_urlsafe(24)}
    if challenge: q.update(code_challenge=hashed, code_challenge_method=method)
    result = browser.request(as_url + '/oauth2/authorize?' + url.urlencode(q), 302, 'authorization request ' + ('S256' if challenge else 'missing PKCE'))
    params = url.parse_qs(url.urlsplit(location(result)).query)
    require(params.get('state', [''])[0] == q['state'], 'AS must return state')
    return browser, verifier, params


def exchange(as_url, code, verifier, expected=200):
    return Browser().request(as_url + '/oauth2/token', expected, 'authorization code exchange', basic=True,
        fields={'grant_type': 'authorization_code', 'code': code, 'redirect_uri': 'http://127.0.0.1/callback', 'code_verifier': verifier})


def roles(as_url, rs, client):
    _, _, missing = authorize(as_url, challenge=False)
    require(missing.get('error') == ['invalid_request'], 'PKCE-required client rejects missing challenge')
    _, _, plain = authorize(as_url, method='plain')
    require(plain.get('error') == ['invalid_request'], 'S256 policy rejects plain')
    _, verifier, params = authorize(as_url)
    code = params['code'][0]
    require(obj(exchange(as_url, code, 'wrong-verifier', 400))['error'] == 'invalid_grant', 'wrong verifier denied')
    # A failed exchange can invalidate the code in this implementation. Use a fresh authorization.
    _, verifier, params = authorize(as_url)
    tokens = obj(exchange(as_url, params['code'][0], verifier))
    require('id_token' in tokens and 'refresh_token' in tokens and tokens['token_type'] == 'Bearer', 'OIDC + access + refresh roles')
    require(obj(exchange(as_url, params['code'][0], verifier, 400))['error'] == 'invalid_grant', 'single-use code')
    # A replay invalidates associated tokens; the positive resource flow uses a fresh grant.
    _, verifier, params = authorize(as_url)
    tokens = obj(exchange(as_url, params['code'][0], verifier))
    b = Browser()
    require(obj(b.request(rs + '/jwt/read', 200, 'AS-issued access token at JWT RS', bearer=tokens['access_token']))['name'] == 'alice', 'resource identity')
    b.request(rs + '/jwt/read', 401, 'ID token cannot replace access token', bearer=tokens['id_token'])
    id_state = obj(b.request(as_url + '/oauth2/introspect', 200, 'ID token active and wrong resource audience confirmed at AS', basic=True, fields={'token': tokens['id_token']}))
    require(id_state.get('active') is True and 'orders-api' not in id_state.get('aud', []), 'opaque aud negative must reach an active token')
    RESULTS.append({'name': 'opaque audience negative precondition', 'activeAtAS': True, 'audienceMatchesResource': False})
    b.request(rs + '/opaque/read', 401, 'active ID token audience rejected by opaque resource policy', bearer=tokens['id_token'])
    b.request(rs + '/opaque/read', 200, 'AS-issued token via real introspection', bearer=tokens['access_token'])
    print('PASS roles: S256 / real code exchange / replay / ID-token separation', flush=True)


def claims(as_url, rs, client):
    b = Browser()
    b.request(rs + '/jwt/read', 401, 'missing bearer')
    b.request(rs + '/jwt/read', 401, 'malformed bearer', bearer='not-a-jwt')
    for kind in ['algorithm', 'signature', 'issuer', 'audience', 'expired', 'no-exp', 'future', 'unknown-kid', 'no-tenant']:
        b.request(rs + '/jwt/read', 401, kind + ' denied', bearer=mint(as_url, kind))
        b.request(rs + '/jwt/read', 200, kind + ' corrected token recovers', bearer=mint(as_url))
    # Same worker-pool/context must not grant the next unauthenticated request.
    b.request(rs + '/jwt/read', 401, 'identity not inherited by next request')
    print('PASS claims: alg/signature/iss/aud/exp/nbf/kid/tenant + recovery', flush=True)


def authorization(as_url, rs, client):
    b, token = Browser(), mint(as_url)
    b.request(rs + '/jwt/read', 403, 'valid signature but missing scope', bearer=mint(as_url, 'scope'))
    b.request(rs + '/jwt/admin', 403, 'request permission denies reader', bearer=token)
    b.request(rs + '/jwt/method', 403, 'method proxy denies missing write scope', bearer=mint(as_url, 'scope'))
    b.request(rs + '/jwt/method', 200, 'method proxy permits write scope', bearer=token)
    for path in ['/jwt/orders/other?tenant=t1', '/jwt/orders/foreign?tenant=t1', '/jwt/orders/foreign?tenant=t2', '/jwt/orders/own?tenant=t2']:
        b.request(rs + path, 403, 'object/tenant boundary ' + path, bearer=token)
    b.request(rs + '/jwt/orders/own?tenant=t1', 200, 'trusted tenant and owner succeed', bearer=token)
    # Valid t2 identity can read its t2 record; changing a client parameter alone cannot do this.
    b.request(rs + '/jwt/orders/foreign?tenant=t2', 200, 'trusted t2 token recovers legitimate t2 access', bearer=mint(as_url, 'tenant'))
    b.request(rs + '/jwt/orders/own?tenant=t1', 403, 't2 identity denied t1 record', bearer=mint(as_url, 'tenant'))
    print('PASS authorization: request / method / object / trusted tenant', flush=True)


def lifecycle(as_url, rs, client):
    _, verifier, p = authorize(as_url)
    t = obj(exchange(as_url, p['code'][0], verifier))
    b = Browser()
    b.request(rs + '/opaque/read', 200, 'active before revocation', bearer=t['access_token'])
    control(as_url, introspectionDown='true')
    try:
        b.request(as_url + '/oauth2/introspect', 503, 'direct introspection proves fault injection', basic=True, fields={'token': t['access_token']})
        b.request(rs + '/opaque/read', 503, 'introspection outage never grants identity', bearer=t['access_token'])
    finally: control(as_url, introspectionDown='false')
    b.request(as_url + '/oauth2/introspect', 200, 'direct introspection proves recovery', basic=True, fields={'token': t['access_token']})
    b.request(rs + '/opaque/read', 200, 'introspection recovery', bearer=t['access_token'])
    b.request(as_url + '/oauth2/revoke', 200, 'actual AS access revocation', basic=True, fields={'token': t['access_token'], 'token_type_hint': 'access_token'})
    b.request(rs + '/opaque/read', 401, 'introspection rejects revoked token', bearer=t['access_token'])
    b.request(rs + '/jwt/read', 200, 'local JWT remains valid until local policy expiry', bearer=t['access_token'])
    b.request(as_url + '/oauth2/revoke', 200, 'unknown token revocation also returns success', basic=True, fields={'token': 'synthetic-unknown-token'})
    # Refresh rotation and revocation are independent lifecycle controls.
    _, verifier, p = authorize(as_url)
    t = obj(exchange(as_url, p['code'][0], verifier))
    rotated = obj(b.request(as_url + '/oauth2/token', 200, 'refresh rotates credentials', basic=True, fields={'grant_type': 'refresh_token', 'refresh_token': t['refresh_token']}))
    require(rotated['refresh_token'] != t['refresh_token'], 'reuseRefreshTokens=false rotates value')
    b.request(as_url + '/oauth2/token', 400, 'old refresh credential denied', basic=True, fields={'grant_type': 'refresh_token', 'refresh_token': t['refresh_token']})
    b.request(as_url + '/oauth2/revoke', 200, 'actual AS refresh revocation', basic=True, fields={'token': rotated['refresh_token'], 'token_type_hint': 'refresh_token'})
    b.request(as_url + '/oauth2/token', 400, 'revoked refresh cannot mint', basic=True, fields={'grant_type': 'refresh_token', 'refresh_token': rotated['refresh_token']})
    print('PASS lifecycle: introspection outage/recovery/revocation vs local JWT + refresh rotation', flush=True)


def oidc(as_url, rs, client):
    b = Browser()
    start = b.request(client + '/oauth2/authorization/lab', 302, 'real Spring OAuth2 client initiation')
    uri = location(start)
    q = url.parse_qs(url.urlsplit(uri).query)
    require(q.get('code_challenge_method') == ['S256'] and all(q.get(k) for k in ['state','nonce','code_challenge']), 'client emits PKCE + state + nonce')
    login(b, as_url)
    callback = location(b.request(uri, 302, 'OIDC AS authorization'))
    b.request(callback, 302, 'real client code exchange and ID-token verification')
    require(obj(b.request(client + '/client/me', 200, 'OIDC identity saved across HTTP requests'))['name'] == 'alice', 'OIDC subject')
    b.request(client + '/client/resource', 200, 'OAuth2 RestClient sends access token to RS')
    # A separate attempt with the right cookie but altered state must fail.
    attack = Browser()
    uri = location(attack.request(client + '/oauth2/authorization/lab', 302, 'state negative initiation'))
    login(attack, as_url)
    cb = location(attack.request(uri, 302, 'state negative authorization'))
    parsed = url.urlsplit(cb)
    values = url.parse_qs(parsed.query)
    values['state'] = ['forged-state']
    altered = url.urlunsplit(parsed._replace(query=url.urlencode({k:v[0] for k,v in values.items()})))
    attack.request(altered, 401, 'mismatched state rejected')
    attack.request(client + '/client/me', 302, 'state failure does not save identity')
    # Real AS signs a deliberately wrong nonce in the ID token, not a mock client decoder.
    control(as_url, badNonce='true')
    try:
        attack = Browser()
        uri = location(attack.request(client + '/oauth2/authorization/lab', 302, 'nonce negative initiation'))
        login(attack, as_url)
        cb = location(attack.request(uri, 302, 'nonce negative authorization'))
        attack.request(cb, 401, 'signed ID token with wrong nonce rejected')
        attack.request(client + '/client/me', 302, 'nonce failure does not save identity')
    finally: control(as_url, badNonce='false')
    fresh = Browser()
    uri = location(fresh.request(client + '/oauth2/authorization/lab', 302, 'OIDC recovery initiation'))
    login(fresh, as_url)
    fresh.request(location(fresh.request(uri, 302, 'OIDC recovery authorization')), 302, 'OIDC recovery callback')
    fresh.request(client + '/client/me', 200, 'OIDC recovery identity')
    print('PASS OIDC client: real login / RestClient / state / nonce / recovery', flush=True)


def rotation(as_url, rs, client):
    b = Browser()
    a = mint(as_url)
    b.request(rs + '/jwt/read', 200, 'A key initially cached', bearer=a)
    control(as_url, publish='AB', active='B')
    token_b = mint(as_url)
    b.request(rs + '/jwt/read', 200, 'unknown B kid refreshes JWKS during overlap', bearer=token_b)
    b.request(rs + '/jwt/read', 200, 'A token remains valid in overlap', bearer=a)
    control(as_url, jwksDown='true')
    try:
        b.request(as_url + '/oauth2/jwks', 503, 'direct JWKS proves fault injection reached endpoint')
        b.request(rs + '/jwt/read', 200, 'known cached B works during JWKS outage', bearer=token_b)
        b.request(rs + '/jwt/read', 503, 'unknown kid needs remote JWKS and fails closed during outage', bearer=mint(as_url, 'unknown-kid'))
    finally: control(as_url, jwksDown='false')
    b.request(as_url + '/oauth2/jwks', 200, 'direct JWKS proves endpoint recovery')
    b.request(rs + '/jwt/read', 200, 'JWKS recovery for valid B', bearer=token_b)
    control(as_url, publish='B')
    b.request(rs + '/jwt/read', 200, 'A still accepted from warm cache immediately after publication retirement', bearer=a)
    # Force a cache refresh with an unknown kid; retirement alone does not push cache invalidation.
    b.request(rs + '/jwt/read', 401, 'unknown kid forces refresh after A retirement', bearer=mint(as_url, 'unknown-kid'))
    b.request(rs + '/jwt/read', 401, 'A rejected after observed cache refresh', bearer=a)
    b.request(rs + '/jwt/read', 200, 'B remains accepted after retirement', bearer=token_b)
    print('PASS key rotation: A -> AB -> B; cached outage / refresh / retirement', flush=True)


def reference(as_url, rs, client):
    b = Browser()
    token = obj(b.request(as_url + '/oauth2/token', 200, 'real reference-format access grant',
        basic='reference-client:reference-secret', fields={'grant_type': 'client_credentials', 'scope': 'orders.read'}))['access_token']
    require('.' not in token, 'reference token is not a JWT serialization')
    b.request(rs + '/opaque/read', 200, 'reference token introspected through actual AS', bearer=token)
    b.request(rs + '/jwt/read', 401, 'reference token cannot be locally JWT-decoded', bearer=token)
    b.request(as_url + '/oauth2/revoke', 200, 'reference token revoked at AS', basic='reference-client:reference-secret', fields={'token': token})
    b.request(rs + '/opaque/read', 401, 'revoked reference token denied', bearer=token)
    print('PASS reference: real opaque format / introspection / revocation', flush=True)


CASES = {'reference': reference, 'roles': roles, 'claims': claims, 'authorization': authorization, 'lifecycle': lifecycle, 'oidc': oidc, 'rotation': rotation}


def main():
    global LOGS, ENV, JAVA
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--case', choices=list(CASES))
    parser.add_argument('--label', default='acceptance')
    args = parser.parse_args()
    CACHE.mkdir(parents=True, exist_ok=True)
    LOGS = Path(tempfile.mkdtemp(prefix='run-', dir=CACHE))
    print('LOGS', LOGS, flush=True)
    evidence = {'ticket':'U30 #67','label':args.label,'startedAt':now(),'rawLogs':str(LOGS),
                'input':['python3',str(ROOT/'run.py'),*os.sys.argv[1:]],'results':RESULTS,'processes':PROCESSES}
    code=1
    try:
        selected=os.environ.get('JAVA_HOME')
        if not selected and Path('/usr/libexec/java_home').exists(): selected=subprocess.check_output(['/usr/libexec/java_home','-v','21'],text=True).strip()
        require(selected,'Set JAVA_HOME=JDK21')
        JAVA=Path(selected)
        ENV={k:v for k,v in os.environ.items() if not k.startswith(('SPRING_','SERVER_','LAB_','LOGGING_','JAVA_TOOL_OPTIONS','_JAVA_OPTIONS','JDK_JAVA_OPTIONS','MAVEN_ARGS','MAVEN_OPTS'))}
        ENV.update(JAVA_HOME=str(JAVA),PATH=str(JAVA/'bin')+os.pathsep+os.environ.get('PATH',''),LAB_CONTROL=CONTROL)
        version=command('jdk',[JAVA/'bin/java','-version'])
        require(re.search(r'version "21[.]',version),'JDK21 required')
        evidence['environment']={'os':platform.platform(),'arch':platform.machine(),'jdk':version.strip(),'python':platform.python_version(),'boot':'4.0.9','security':'7.0.8','framework':'7.0.10'}
        command('maven',['mvn','-version'])
        command('package',['mvn','-q','-B','-ntp','package'])
        command('dependency-tree',['mvn','-B','-ntp','dependency:tree'])
        jar=ROOT/'target/oauth2-authorization-1.jar'
        evidence['jarSha256']=hashlib.sha256(jar.read_bytes()).hexdigest()
        with zipfile.ZipFile(jar) as z:
            libs=sorted(n.split('/')[-1] for n in z.namelist() if n.startswith('BOOT-INF/lib/') and n.endswith('.jar'))
        evidence['packagedLibraries']=libs
        for name in ['spring-boot-4.0.9.jar','spring-security-core-7.0.8.jar','spring-security-oauth2-authorization-server-7.0.8.jar','spring-security-oauth2-client-7.0.8.jar','spring-security-oauth2-resource-server-7.0.8.jar','spring-security-oauth2-jose-7.0.8.jar','spring-webmvc-7.0.10.jar']:
            require(name in libs,'BOM library mismatch '+name)
        evidence['sourceSha256']={str(p.relative_to(ROOT)):hashlib.sha256(p.read_bytes()).hexdigest() for p in sorted(ROOT.rglob('*')) if p.is_file() and not {'target','evidence','__pycache__'}.intersection(p.relative_to(ROOT).parts)}
        # Reserve three distinct loopback ports, then release for Boot. Any intervening bind race fails visibly.
        sockets=[]
        for _ in range(3):
            s=socket.socket();s.bind(('127.0.0.1',0));sockets.append(s)
        ports=[s.getsockname()[1] for s in sockets]
        for s in sockets:s.close()
        as_url, rs, client=['http://127.0.0.1:'+str(p) for p in ports]
        evidence['roles']={'authorizationServer':as_url,'resourceServer':rs,'oidcClient':client}
        with server('as',ports[0],as_url,client,rs), server('rs',ports[1],as_url,client,rs), server('client',ports[2],as_url,client,rs):
            for name in [args.case] if args.case else CASES:
                CASES[name](as_url,rs,client)
                RESULTS.append({'name':name,'status':'PASS','finishedAt':now()})
        code=0
    except Exception as e:
        evidence['failure']={'type':type(e).__name__,'message':str(e)}
        print('FAIL',type(e).__name__,str(e),flush=True)
    finally:
        evidence.update(exitCode=code,finishedAt=now(),httpRequests=sum('expectedStatus' in r for r in RESULTS),allProcessesReaped=all(p['reaped'] for p in PROCESSES))
        evidence['limits']=['loopback HTTP synthetic lab; no production TLS/IdP/database/cluster','real OAuth2/OIDC, actual reference-format access token and JWT access token both remotely introspected; reference client uses machine identity','expiry/nbf inputs signed with real temporary RSA; clocks not mocked','JWKS refresh observed after explicit unknown-kid trigger; retirement alone does not invalidate every cache']
        out=LOGS/'result.json';out.write_text(json.dumps(evidence,ensure_ascii=False,indent=2)+'\n')
        if code==0 and not args.case:
            target=ROOT/'evidence/u30-2026-10-26.json';target.write_text(out.read_text())
        print('EVIDENCE',out,flush=True)
    return code

if __name__=='__main__':raise SystemExit(main())