"""Key transport against real filesystem/crypto and a synthetic Kubernetes API."""
import base64
import importlib.util
import json
import os
from pathlib import Path
import subprocess

import pytest
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa

ROOT = Path(__file__).resolve().parents[1]
PI = ROOT / 'sso-mfa/k8s/privacyidea'


def module(name, file):
    spec = importlib.util.spec_from_file_location(name, file)
    result = importlib.util.module_from_spec(spec); spec.loader.exec_module(result)
    return result


@pytest.fixture
def keys():
    private = rsa.generate_private_key(public_exponent=65537, key_size=2048)
    return {'enckey': os.urandom(96),
            'private.pem': private.private_bytes(serialization.Encoding.PEM, serialization.PrivateFormat.TraditionalOpenSSL, serialization.NoEncryption()),
            'public.pem': private.public_key().public_bytes(serialization.Encoding.PEM, serialization.PublicFormat.SubjectPublicKeyInfo)}


@pytest.fixture
def case(tmp_path, keys):
    root = tmp_path / 'private'; root.mkdir(mode=0o700)
    (root / 'privacyidea').mkdir(mode=0o700)
    volume = tmp_path / 'volume'; volume.mkdir(mode=0o750)
    for name, value in keys.items():
        (volume / name).write_bytes(value); (volume / name).chmod(0o600)
    binpath = tmp_path / 'bin'; binpath.mkdir()
    metadata = lambda name, uid: {'name': name, 'namespace': 'mfa', 'uid': uid, 'resourceVersion': '1'}
    mount = {'name': 'data', 'mountPath': '/etc/privacyidea'}
    spec = {'securityContext': {'fsGroup': 999}, 'volumes': [{'name': 'data', 'persistentVolumeClaim': {'claimName': 'privacyidea-data'}}],
            'containers': [{'name': 'privacyidea', 'volumeMounts': [mount]}]}
    deployment = {'apiVersion': 'apps/v1', 'kind': 'Deployment', 'metadata': {**metadata('privacyidea', 'deployment-uid'), 'generation': 1},
                  'spec': {'replicas': 1, 'template': {'spec': spec}}, 'status': {'replicas': 1, 'observedGeneration': 1}}
    pod = {'apiVersion': 'v1', 'kind': 'Pod', 'metadata': {**metadata('privacyidea-source', 'pod-uid'),
           'ownerReferences': [{'kind': 'ReplicaSet', 'name': 'privacyidea-rs', 'uid': 'rs-uid', 'controller': True}]},
           'spec': spec, 'status': {'phase': 'Running', 'containerStatuses': [{'name': 'privacyidea', 'ready': True}]}}
    state = {'deployment': deployment, 'pvc': {'apiVersion': 'v1', 'kind': 'PersistentVolumeClaim', 'metadata': metadata('privacyidea-data', 'pvc-uid'), 'status': {'phase': 'Bound'}},
             'pods': {'apiVersion': 'v1', 'kind': 'PodList', 'items': [pod]},
             'replicaset': {'apiVersion': 'apps/v1', 'kind': 'ReplicaSet', 'metadata': {**metadata('privacyidea-rs', 'rs-uid'),
                 'ownerReferences': [{'kind': 'Deployment', 'uid': 'deployment-uid', 'controller': True}]}}, 'secrets': {}}
    statepath = tmp_path / 'state.json'; statepath.write_text(json.dumps(state))
    binary = binpath / 'kubectl'
    binary.write_text('''#!/usr/bin/python3
import base64,importlib.util,json,os,sys
from pathlib import Path
args=sys.argv[1:];p=Path(os.environ['STATE']);state=json.loads(p.read_text())
raw=sys.stdin.buffer.read() if args[0] in ('exec','create','patch') else b''
with open(os.environ['CALLS'],'a') as f:f.write(json.dumps({'args':args,'stdin':raw.decode()})+'\\n')
if os.environ.get('FAIL')==args[0]:
 print('SYNTHETIC_SECRET_DIAGNOSTIC');print('SYNTHETIC_SECRET_DIAGNOSTIC',file=sys.stderr);sys.exit(1)
if args[0]=='get':
 if args[1]=='secret':
  obj=state['secrets'].get(args[2]);print(json.dumps(obj) if obj is not None else '');sys.exit(0)
 if args[1]=='pod':print('privacyidea-source');sys.exit(0)
 print(json.dumps(state[args[1]]));sys.exit(0)
if args[0]=='exec':
 if args[-2]=='cat':sys.stdout.buffer.write((Path(os.environ['VOLUME'])/Path(args[-1]).name).read_bytes());sys.exit(0)
 if 'test' in args:sys.exit(0)
 spec=importlib.util.spec_from_file_location('remote',Path(os.environ['SOURCE'])/'key-custody-remote.py')
 remote=importlib.util.module_from_spec(spec);spec.loader.exec_module(remote)
 assert args[-1]==(Path(os.environ['SOURCE'])/'key-custody-remote.py').read_text()
 request=json.loads(raw)
 result=remote.operate(request,Path(os.environ['VOLUME']),uid=os.getuid(),gid=os.getgid())
 if os.environ.get('RACE')=='pod':
  state['pods']['items'][0]['metadata']['uid']='changed-pod';p.write_text(json.dumps(state))
 print(json.dumps(result));sys.exit(0)
if args[0]=='create':
 if 'generic' in args:print('{}');sys.exit(0)
 obj=json.loads(raw);name=obj['metadata']['name']
 if name in state['secrets'] or os.environ.get('RACE')=='create':sys.exit(1)
 obj['metadata'].update(uid='secret-'+name,resourceVersion='1')
elif args[0]=='patch':
 name=args[2];obj=state['secrets'][name]
 if os.environ.get('RACE')=='patch':sys.exit(1)
 for operation in json.loads(raw):
  parts=[v.replace('~1','/').replace('~0','~') for v in operation['path'].split('/')[1:]];parent=obj
  for part in parts[:-1]:parent=parent[part]
  if operation['op']=='test':assert parent[parts[-1]]==operation['value']
  elif operation['op']=='remove':del parent[parts[-1]]
  else:sys.exit(2)
 obj['metadata']['resourceVersion']='2'
elif args[0]=='apply':sys.exit(0)
else:sys.exit(2)
state['secrets'][name]=obj;p.write_text(json.dumps(state))
if os.environ.get('BAD_RESPONSE'):obj['data']['unexpected']='eA=='
print(json.dumps(obj))
'''); binary.chmod(0o700)
    env = {'PATH': str(binpath) + ':' + os.environ['PATH'], 'HOME': str(tmp_path), 'STATE': str(statepath),
           'CALLS': str(tmp_path/'calls'), 'VOLUME': str(volume), 'SOURCE': str(PI)}
    return {'root': root, 'volume': volume, 'keys': keys, 'env': env, 'statepath': statepath}


def state(case): return json.loads(case['statepath'].read_text())
def save(case, value): case['statepath'].write_text(json.dumps(value))
def calls(case):
    p=Path(case['env']['CALLS']);return [json.loads(line) for line in p.read_text().splitlines()] if p.exists() else []


def invoke(case, action='export', *args, **extra):
    script = 'enckey-bootstrap.sh' if action=='export' else 'enckey-restore.sh'
    return subprocess.run(['bash','-x',str(PI/script),str(case['root']),*args],env={**case['env'],**extra},capture_output=True,text=True,timeout=30)


def quiet(result, case):
    assert 'SYNTHETIC_SECRET_DIAGNOSTIC' not in result.stdout+result.stderr
    assert 'BEGIN RSA PRIVATE KEY' not in result.stdout+result.stderr
    assert base64.b64encode(case['keys']['enckey']).decode() not in result.stdout+result.stderr
    assert 'Traceback' not in result.stdout+result.stderr


def test_export_roundtrip_and_preserves_matching_copies(case):
    result=invoke(case);assert result.returncode==0,result.stderr;quiet(result,case)
    for name,value in case['keys'].items():
        p=case['root']/'privacyidea'/('pi.enc' if name=='enckey' else name)
        assert p.read_bytes()==value and p.stat().st_mode & 0o077==0
    expected=state(case);count=len(calls(case));assert invoke(case).returncode==0
    assert state(case)==expected and all(c['args'][0] in ('get','exec') for c in calls(case)[count:])
    argv=json.dumps([c['args'] for c in calls(case)])
    assert base64.b64encode(case['keys']['enckey']).decode() not in argv
    assert all(c['args'][0]!='apply' for c in calls(case))


@pytest.mark.parametrize('defect',['different','symlink','hardlink','public','fifo'])
def test_export_refuses_bad_local_copy_before_secret_writes(case,defect):
    p=case['root']/'privacyidea/pi.enc'
    if defect=='fifo':os.mkfifo(p,0o600)
    elif defect=='symlink':p.symlink_to(case['volume']/'enckey')
    else:
        p.write_bytes(b'old-copy' if defect=='different' else case['keys']['enckey']);p.chmod(0o644 if defect=='public' else 0o600)
        if defect=='hardlink':os.link(p,p.with_name('alias'))
    result=invoke(case);assert result.returncode!=0;quiet(result,case)
    assert not state(case)['secrets'] and not (p.parent/'private.pem').exists()
    if defect=='different':assert p.read_bytes()==b'old-copy'


@pytest.mark.parametrize('defect',['missing','short','link','mismatched-rsa'])
def test_export_refuses_invalid_remote_keys_without_generation(case,defect,keys):
    p=case['volume']/'enckey'
    if defect=='missing':p.unlink()
    elif defect=='short':p.write_bytes(b'short')
    elif defect=='link':os.link(p,p.with_name('alias'))
    else:
        other=rsa.generate_private_key(public_exponent=65537,key_size=2048)
        (case['volume']/'public.pem').write_bytes(other.public_key().public_bytes(serialization.Encoding.PEM,serialization.PublicFormat.SubjectPublicKeyInfo))
    result=invoke(case);assert result.returncode!=0;quiet(result,case)
    assert not list((case['root']/'privacyidea').iterdir()) and not state(case)['secrets']
    assert not any('pi-manage' in c['args'] for c in calls(case))


@pytest.mark.parametrize('defect',['two-pods','not-ready','wrong-owner','wrong-pvc','pod-race'])
def test_export_binds_ready_owned_pod_and_pvc(case,defect):
    obj=state(case);pod=obj['pods']['items'][0]
    if defect=='two-pods':obj['pods']['items'].append(pod)
    elif defect=='not-ready':pod['status']['containerStatuses'][0]['ready']=False
    elif defect=='wrong-owner':obj['replicaset']['metadata']['ownerReferences'][0]['uid']='other'
    elif defect=='wrong-pvc':pod['spec']['volumes'][0]['persistentVolumeClaim']['claimName']='other'
    save(case,obj)
    result=invoke(case,RACE='pod' if defect=='pod-race' else '');assert result.returncode!=0;quiet(result,case)
    assert not list((case['root']/'privacyidea').iterdir()) and not state(case)['secrets']


def test_export_failure_retains_local_recovery_copies(case):
    result=invoke(case,FAIL='create');assert result.returncode!=0;quiet(result,case)
    assert (case['root']/'privacyidea/pi.enc').read_bytes()==case['keys']['enckey']
    assert not list((case['root']/'privacyidea').glob('.key-export-*'))


@pytest.mark.parametrize('defect',['different','legacy-field','owner','immutable','create-race','bad-response'])
def test_export_existing_backup_and_race_refusal(case,defect):
    if defect not in ('create-race','bad-response'):
        assert invoke(case).returncode==0
        obj=state(case);secret=obj['secrets']['privacyidea-enckey']
        if defect=='different':secret['data']['enckey']=base64.b64encode(b'other').decode()
        elif defect=='legacy-field':secret['data']['PI_ENCFILE']=secret['data'].pop('enckey')
        elif defect=='owner':secret['metadata']['ownerReferences']=[{'uid':'other-owner'}]
        else:secret['immutable']=True
        save(case,obj)
    before=state(case)
    result=invoke(case,RACE='create' if defect=='create-race' else '',BAD_RESPONSE='yes' if defect=='bad-response' else '')
    assert result.returncode!=0;quiet(result,case)
    if defect!='bad-response':assert state(case)==before


def test_export_removes_annotation_conditionally_and_preserves_extra_data(case):
    assert invoke(case).returncode==0;obj=state(case)
    for secret in obj['secrets'].values():
        secret['metadata']['annotations']={'kubectl.kubernetes.io/last-applied-configuration':'synthetic-copy','owner':'keep'}
        secret['data']['extra']='eA=='
    save(case,obj);assert invoke(case).returncode==0
    for secret in state(case)['secrets'].values():
        assert secret['metadata']['annotations']=={'owner':'keep'} and secret['data']['extra']=='eA=='
    patches=[json.loads(c['stdin']) for c in calls(case) if c['args'][0]=='patch']
    assert len(patches)==2 and all({op['path'] for op in patch if op['op']=='test'}=={'/metadata/uid','/metadata/resourceVersion','/data'} for patch in patches)


def test_explicit_first_bootstrap_creates_only_missing_audit_pair(case):
    for name in ('private.pem','public.pem'):(case['volume']/name).unlink()
    assert invoke(case).returncode!=0
    result=invoke(case,'export','--initialize-audit-keys');assert result.returncode==0,result.stderr;quiet(result,case)
    assert (case['volume']/'enckey').read_bytes()==case['keys']['enckey']
    assert invoke(case,'export','--initialize-audit-keys').returncode!=0
    assert invoke(case).returncode==0


def recovery(case):
    assert invoke(case).returncode==0
    host=module('host',PI/'key-custody.py')
    pod=host.render('ghcr.io/gpappsoft/privacyidea-docker@sha256:'+'a'*64,'pvc-uid')
    pod['metadata'].update(uid='restore-pod-uid',resourceVersion='1')
    pod['status']={'phase':'Running','containerStatuses':[{'name':'recovery','ready':True}]}
    obj=state(case);obj['pods']['items']=[pod];obj['deployment']['spec']['replicas']=0;obj['deployment']['status']['replicas']=0;save(case,obj)
    for name in case['keys']:(case['volume']/name).unlink()


def test_restore_installs_on_actual_volume_and_retry_preserves_files(case):
    recovery(case);result=invoke(case,'restore','--pvc-uid','pvc-uid');assert result.returncode==0,result.stderr;quiet(result,case)
    for name,value in case['keys'].items():assert (case['volume']/name).read_bytes()==value
    before={p.name:p.stat().st_ino for p in case['volume'].iterdir()}
    assert invoke(case,'restore','--pvc-uid','pvc-uid').returncode==0
    assert before=={p.name:p.stat().st_ino for p in case['volume'].iterdir()}


@pytest.mark.parametrize('defect',['deployment-running','pod-terminating','other-pod','wrong-pvc-uid','mutable-image','no-mount','app-entrypoint','token-mount'])
def test_restore_rejects_unstopped_or_wrong_target_before_key_transfer(case,defect):
    recovery(case);obj=state(case);pod=obj['pods']['items'][0]
    if defect=='deployment-running':obj['deployment']['spec']['replicas']=1
    elif defect=='pod-terminating':pod['metadata']['deletionTimestamp']='now'
    elif defect=='other-pod':obj['pods']['items'].append(pod)
    elif defect=='wrong-pvc-uid':obj['pvc']['metadata']['uid']='replacement-pvc'
    elif defect=='mutable-image':pod['spec']['containers'][0]['image']='image:latest'
    elif defect=='no-mount':pod['spec']['containers'][0]['volumeMounts']=[]
    elif defect=='app-entrypoint':pod['spec']['containers'][0]['command']=['pi-manage','run']
    else:pod['spec']['automountServiceAccountToken']=True
    save(case,obj);count=len(calls(case));result=invoke(case,'restore','--pvc-uid','pvc-uid')
    assert result.returncode!=0;quiet(result,case)
    assert not list(case['volume'].iterdir()) and all(c['args'][0]=='get' for c in calls(case)[count:])


def test_restore_refuses_partial_mismatching_state_without_overwrite(case):
    recovery(case);p=case['volume']/'private.pem';p.write_bytes(b'other');p.chmod(0o600)
    result=invoke(case,'restore','--pvc-uid','pvc-uid');assert result.returncode!=0;quiet(result,case)
    assert p.read_bytes()==b'other' and not (case['volume']/'enckey').exists()


def test_restored_key_decrypts_existing_native_privacyidea_values(case):
    pytest.importorskip('privacyidea')
    from privacyidea.lib.security.default import DefaultSecurityModule
    source=DefaultSecurityModule({'file':str(case['volume']/'enckey')})
    ciphertext=source.encrypt_pin('synthetic-token-pin')
    password=source.encrypt_password('synthetic-configuration-value')
    from privacyidea.lib.crypto import Sign
    signer=Sign((case['volume']/'private.pem').read_bytes(),(case['volume']/'public.pem').read_bytes())
    signature=signer.sign('synthetic-audit-record')
    recovery(case);assert invoke(case,'restore','--pvc-uid','pvc-uid').returncode==0
    restored=DefaultSecurityModule({'file':str(case['volume']/'enckey')})
    assert restored.decrypt_pin(ciphertext)=='synthetic-token-pin'
    assert restored.decrypt_password(password)=='synthetic-configuration-value'
    recovered_signer=Sign((case['volume']/'private.pem').read_bytes(),(case['volume']/'public.pem').read_bytes())
    assert recovered_signer.verify('synthetic-audit-record',signature)
    assert signer.verify('synthetic-new-audit-record',recovered_signer.sign('synthetic-new-audit-record'))


def test_remote_partial_restore_keeps_keys_and_retries_without_replacement(tmp_path, keys, monkeypatch):
    remote=module('partial_remote',PI/'key-custody-remote.py')
    real_link=remote.os.link
    def interrupted(source,target,**kwargs):
        if target=='public.pem':raise OSError('synthetic-write-failure')
        return real_link(source,target,**kwargs)
    request={'action':'restore','keys':{k:base64.b64encode(v).decode() for k,v in keys.items()}}
    monkeypatch.setattr(remote.os,'link',interrupted)
    with pytest.raises(OSError):remote.operate(request,tmp_path,uid=os.getuid(),gid=os.getgid())
    assert (tmp_path/'enckey').read_bytes()==keys['enckey']
    assert (tmp_path/'private.pem').read_bytes()==keys['private.pem']
    assert not list(tmp_path.glob('.key-custody-*'))
    inodes={p.name:p.stat().st_ino for p in tmp_path.iterdir()}
    monkeypatch.setattr(remote.os,'link',real_link)
    assert remote.operate(request,tmp_path,uid=os.getuid(),gid=os.getgid())=={'result':'restored'}
    assert all((tmp_path/name).stat().st_ino==inode for name,inode in inodes.items())


def test_remote_partial_audit_initialization_cannot_generate_another_pair(tmp_path, keys, monkeypatch):
    remote=module('init_remote',PI/'key-custody-remote.py')
    p=tmp_path/'enckey';p.write_bytes(keys['enckey']);p.chmod(0o600)
    real_link=remote.os.link
    def interrupted(source,target,**kwargs):
        if target=='public.pem':raise OSError('synthetic-write-failure')
        return real_link(source,target,**kwargs)
    monkeypatch.setattr(remote.os,'link',interrupted)
    with pytest.raises(OSError):remote.operate({'action':'initialize-audit'},tmp_path,uid=os.getuid(),gid=os.getgid())
    before=(tmp_path/'private.pem').read_bytes()
    monkeypatch.setattr(remote.os,'link',real_link)
    with pytest.raises(remote.Refused):remote.operate({'action':'initialize-audit'},tmp_path,uid=os.getuid(),gid=os.getgid())
    assert (tmp_path/'private.pem').read_bytes()==before


def test_unknown_argument_is_not_echoed(case):
    result=invoke(case,'export','--SYNTHETIC_SECRET_DIAGNOSTIC')
    assert result.returncode!=0;quiet(result,case);assert not calls(case)


def test_kubectl_generic_list_output_is_accepted(case):
    obj=state(case);obj['pods']['kind']='List';save(case,obj)
    assert invoke(case).returncode==0