#!/usr/bin/env python3
"""Owner-run native privacyIDEA existing-token MFA check on the restored database.

Restored LLDAP, the restored database and a privacyIDEA worker share one network-free
namespace. The owner names their login and types a token's PIN+OTP privately. The call
is POST /validate/check by user and realm, as KeyCape makes it, through the restored LDAP
resolver: first with a wrong encryption key, then with a wrong code, then as typed.
Codes never enter receipts, argv or environment; the login is stored only as a hash.
"""
import argparse
import hashlib
import importlib.util
import json
import os
from pathlib import Path
import re
import signal
import subprocess
import sys
import tempfile
import uuid
from datetime import datetime, timezone

HERE = Path(__file__).resolve().parent
# Reused unedited: container, isolation and cleanup helpers of the accepted stage.
spec = importlib.util.spec_from_file_location('isolated_privacyidea', HERE / 'isolated-privacyidea.py')
pi = importlib.util.module_from_spec(spec)
spec.loader.exec_module(pi)
local, bundle, run = pi.local, pi.bundle, pi.run
LOGIN = re.compile(r'[A-Za-z0-9_.@-]{1,128}')
REALM = re.compile(r'[A-Za-z0-9_.-]{1,64}')
CLUSTER_LDAP = 'ldap://lldap.sso.svc.cluster.local:3890'
LOOPBACK_LDAP = 'ldap://127.0.0.1:3890'

REALMS = r"""
import contextlib,io,json,logging
logging.disable(logging.CRITICAL)
result={'status':'failed'}
try:
    with contextlib.redirect_stdout(io.StringIO()),contextlib.redirect_stderr(io.StringIO()):
        from privacyidea.app import create_app
        from privacyidea.lib.realm import get_realms,get_default_realm
        app=create_app(config_name='testing',config_file='/recovery/pi.cfg',silent=True)
        with app.app_context():
            names=sorted(get_realms());default=get_default_realm()
    result={'status':'ok','realms':names,'default':default}
except Exception as error:result={'status':'failed','failure_type':type(error).__name__}
print(json.dumps(result))
"""

# Looks the user up through the restored resolver and the restored directory.
RESOLVE = r"""
import contextlib,io,json,logging,sys
logging.disable(logging.CRITICAL)
v=json.load(sys.stdin)
result={'status':'failed'}
stage='native_application_configuration'
try:
    with contextlib.redirect_stdout(io.StringIO()),contextlib.redirect_stderr(io.StringIO()):
        from privacyidea.app import create_app
        app=create_app(config_name='testing',config_file='/recovery/pi.cfg',silent=True)
        with app.app_context():
            from privacyidea.lib.user import User
            from privacyidea.lib.token import get_tokens
            stage='user_lookup_through_restored_resolver'
            user=User(v['login'],v['realm'])
            found=bool(user.exist())
            tokens=[]
            if found:
                stage='user_token_listing'
                tokens=[{'type':t.type,'active':bool(t.is_active()),'locked':bool(t.token.locked),
                         'revoked':bool(t.token.revoked)} for t in get_tokens(user=user)]
    result={'status':'ok','user_found':found,'tokens':tokens}
except Exception as error:result={'status':'failed','failed_stage':stage,'failure_type':type(error).__name__}
print(json.dumps(result))
"""

# One check per process: privacyIDEA loads its encryption module once per app.
# The call is the endpoint KeyCape uses, so stored policies and the resolver both apply.
CHECK = r"""
import contextlib,io,json,logging,os,sys
logging.disable(logging.CRITICAL)
v=json.load(sys.stdin)
result={'status':'failed'}
stage='network_isolation'
try:
    with contextlib.redirect_stdout(io.StringIO()),contextlib.redirect_stderr(io.StringIO()):
        from pathlib import Path
        assert {p.name for p in Path('/sys/class/net').iterdir()}=={'lo'}
        assert len(Path('/proc/net/route').read_text().splitlines())==1
        config='/recovery/pi.cfg'
        if v['mode']=='wrong_key':
            stage='wrong_key_configuration'
            original=Path(config).read_text()
            assert '/recovery/enckey' in original
            Path('/recovery/wrong-enckey').write_bytes(os.urandom(96))
            Path('/recovery/pi-wrongkey.cfg').write_text(original.replace('/recovery/enckey','/recovery/wrong-enckey'))
            config='/recovery/pi-wrongkey.cfg'
        stage='native_application_configuration'
        # create_app applies PRIVACYIDEA_CONFIGFILE after config_file, so it must name the file in use.
        os.environ['PRIVACYIDEA_CONFIGFILE']=config
        from privacyidea.app import create_app
        app=create_app(config_name='testing',config_file=config,silent=True)
        with app.app_context():
            # A wrong-key run that silently used the original key would be vacuous.
            assert app.config['PI_ENCFILE']==('/recovery/wrong-enckey' if v['mode']=='wrong_key' else '/recovery/enckey')
            stage='native_validation'
            accepted=False;refusal='none';code=None
            try:
                response=app.test_client().post('/validate/check',data={'user':v['login'],'realm':v['realm'],'pass':v['passcode']})
                body=response.get_json(silent=True) or {}
                outcome=body.get('result') or {}
                accepted=response.status_code==200 and outcome.get('status') is True and outcome.get('value') is True
                code=(outcome.get('error') or {}).get('code')
                message=str((body.get('detail') or {}).get('message','')).lower()
                if accepted:refusal='none'
                elif code is not None:refusal='error_'+str(code)
                else:refusal='wrong_pin' if 'pin' in message else 'wrong_value' if ('otp' in message or 'wrong' in message) else 'other_refusal'
            except Exception as error:
                accepted=False;refusal='exception_'+type(error).__name__[:40]
    result={'status':'completed','accepted':accepted,'refusal':refusal,
            'only_loopback_interface_verified':True,'external_routes_absent':True}
except Exception as error:result={'status':'failed','failed_stage':stage,'failure_type':type(error).__name__}
print(json.dumps(result))
"""

WAIT_LDAP = r"""
import socket,sys,time
for _ in range(240):
    try:
        with socket.create_connection(('127.0.0.1',3890),timeout=1):sys.exit(0)
    except OSError:time.sleep(.25)
sys.exit(1)
"""


def preflight():
    """The launcher calls this before any key prompt."""
    return pi.preflight()


def worker_json(worker, script, values=None, timeout=300):
    args = ['docker', 'exec'] + (['-i'] if values is not None else []) + [worker, 'python3', '-c', script]
    raw = run(args, input=None if values is None else json.dumps(values).encode(), timeout=timeout)
    return json.loads(raw)


def different_code(passcode):
    """Same shape, different final digit: still a well-formed but wrong code."""
    last = passcode[-1]
    return passcode[:-1] + ('0' if last != '0' else '1')


def psql(database, statement):
    out = run(['docker', 'exec', database, 'psql', '-U', 'postgres', '-d', 'recovery', '-v', 'ON_ERROR_STOP=1',
               '-At', '-c', statement])
    return out.decode().strip()


def rewrite_resolvers(database):
    """Point archived LDAP resolvers at the restored directory, in the disposable copy only.

    Only the archived in-cluster address is accepted; anything else is an owner decision.
    The encrypted bind password is left untouched so the original key must decrypt it.
    """
    rows = psql(database, 'SELECT c."Value" FROM resolverconfig c JOIN resolver r ON r.id=c.resolver_id '
                          "WHERE r.rtype='ldapresolver' AND c.\"Key\"='LDAPURI'").splitlines()
    if not rows or any(value != CLUSTER_LDAP for value in rows):
        raise bundle.Refused()
    changed = psql(database, 'UPDATE resolverconfig SET "Value"=\'' + LOOPBACK_LDAP + '\' WHERE "Key"=\'LDAPURI\' '
                             "AND resolver_id IN (SELECT id FROM resolver WHERE rtype='ldapresolver') "
                             "AND \"Value\"='" + CLUSTER_LDAP + "'")
    if changed != 'UPDATE ' + str(len(rows)):
        raise bundle.Refused()
    return len(rows)


def start_directory(name, database, root):
    """Restored LLDAP joins the database's network-free namespace (loopback only)."""
    config, jwt, admin_password, _ = local.configuration(root)
    database_id = run(['docker', 'container', 'inspect', database, '--format', '{{.Id}}']).decode().strip()
    run(['docker', 'run', '--detach', '--pull=never', '--name', name,
         '--network', 'container:' + database_id, '--read-only', '--cap-drop', 'ALL',
         '--security-opt', 'no-new-privileges', '--user', '1000:1000',
         '--memory', '256m', '--cpus', '1', '--pids-limit', '128',
         '--tmpfs', '/data:rw,nosuid,nodev,noexec,size=268435456,uid=1000,gid=1000,mode=0700',
         '--tmpfs', '/run/recovery:rw,nosuid,nodev,noexec,size=1048576,uid=1000,gid=1000,mode=0700',
         '--tmpfs', '/tmp:rw,nosuid,nodev,noexec,size=16777216,uid=1000,gid=1000',
         '--env', 'LLDAP_JWT_SECRET_FILE=/run/recovery/jwt',
         '--env', 'LLDAP_LDAP_USER_PASS_FILE=/run/recovery/admin-password',
         '--entrypoint', '/bin/sh', local.IMAGE, '-c', 'sleep 1800'])
    pi.isolation(name, 'container:' + database_id, '1000:1000')
    payload = {p.relative_to(root / 'data/lldap').as_posix(): p.read_bytes()
               for p in (root / 'data/lldap').rglob('*') if p.is_file()}
    run(['docker', 'exec', '-i', name, 'tar', '-xf', '-', '-C', '/data'], input=local.archive_files(payload))
    run(['docker', 'exec', '-i', name, 'tar', '-xf', '-', '-C', '/run/recovery'],
        input=local.archive_files({'lldap.toml': config, 'jwt': jwt, 'admin-password': admin_password}))
    run(['docker', 'exec', '--detach', name, '/bin/sh', '-c',
         'exec /app/lldap run --config-file /run/recovery/lldap.toml > /dev/null 2>&1'])


def restore_and_check(root, receipt, choose_user, choose_code):
    """choose_user(realms, default) -> (login, realm); choose_code(summary) -> PIN+code."""
    stem = 'nk49-pi-mfa-' + uuid.uuid4().hex
    database, directory, worker = stem + '-db', stem + '-dir', stem + '-worker'
    stage = 'database_restore'
    try:
        pi.start_database(database)
        with (root / 'databases/privacyidea.dump').open('rb') as dump:
            run(['docker', 'exec', '-i', database, 'pg_restore', '-U', 'postgres', '-d', 'recovery',
                 '--no-owner', '--no-privileges', '--exit-on-error'], stdin=dump, timeout=600)
        stage = 'resolver_rewrite'
        receipt['ldap_resolvers_rewritten_to_loopback'] = rewrite_resolvers(database)
        stage = 'directory_launch'
        start_directory(directory, database, root)
        stage = 'worker_launch'
        pi.start_worker(worker, database, root)
        receipt['container_isolation_verified_before_private_import'] = True
        stage = 'directory_startup'
        run(['docker', 'exec', worker, 'python3', '-c', WAIT_LDAP], timeout=90)
        stage = 'realm_listing'
        realms = worker_json(worker, REALMS)
        if realms.get('status') != 'ok' or not realms['realms']:
            raise bundle.Refused()
        stage = 'owner_user_selection'
        login, realm = choose_user(realms['realms'], realms.get('default'))
        if not LOGIN.fullmatch(login) or realm not in realms['realms']:
            raise bundle.Refused()
        stage = 'user_lookup_through_restored_resolver'
        user = worker_json(worker, RESOLVE, {'login': login, 'realm': realm})
        if user.get('status') != 'ok':
            receipt['failed_stage'] = user.get('failed_stage', stage)
            raise bundle.Refused()
        receipt['user_resolved_through_restored_resolver_and_directory'] = bool(user['user_found'])
        if not user['user_found']:
            raise bundle.Refused()
        usable = [t for t in user['tokens'] if t['active'] and not t['locked'] and not t['revoked']]
        receipt.update(selected_user_sha256=hashlib.sha256(login.encode()).hexdigest(),
                       user_token_count=len(user['tokens']), user_usable_token_count=len(usable),
                       user_usable_token_types=sorted({t['type'] for t in usable}))
        stage = 'owner_code_entry'
        if not usable:
            raise bundle.Refused()
        passcode = choose_code({'tokens': user['tokens'], 'usable': len(usable)})
        if not passcode or len(passcode) > 256:
            raise bundle.Refused()
        outcomes = {}
        for mode, code in (('wrong_key', passcode), ('wrong_code', different_code(passcode)), ('correct', passcode)):
            stage = 'check_' + mode
            outcome = worker_json(worker, CHECK, {'mode': mode, 'login': login, 'realm': realm, 'passcode': code})
            if outcome.get('status') != 'completed':
                receipt['failed_stage'] = outcome.get('failed_stage', stage)
                raise bundle.Refused()
            outcomes[mode] = outcome
        receipt.update(
            wrong_key_refused=not outcomes['wrong_key']['accepted'], wrong_key_refusal=outcomes['wrong_key']['refusal'],
            wrong_code_refused=not outcomes['wrong_code']['accepted'], wrong_code_refusal=outcomes['wrong_code']['refusal'],
            existing_token_native_check_accepted=outcomes['correct']['accepted'],
            only_loopback_interface_verified=all(o['only_loopback_interface_verified'] for o in outcomes.values()),
            external_routes_absent=all(o['external_routes_absent'] for o in outcomes.values()))
        # Only the last step can spend a counter or replay-protected code.
        proven = (receipt['wrong_key_refused'] and receipt['wrong_code_refused']
                  and receipt['existing_token_native_check_accepted'])
        receipt['existing_token_mfa_proven'] = proven
        receipt['user_to_token_resolver_proven'] = proven
        if not proven:
            stage = 'acceptance'
            raise bundle.Refused()
    except Exception:
        receipt.setdefault('failed_stage', stage)
        raise
    finally:
        pi.cleanup([worker, directory, database])
        receipt['container_cleanup_verified'] = True


def owner_user(realms, default):
    """Terminal-only prompts; the login is hashed in receipts."""
    if not (sys.stdin.isatty() and sys.stderr.isatty()):
        raise bundle.Refused()
    print('Realms in the restored database: ' + ', '.join(realms) + ('  (default: %s)' % default if default else ''))
    login = input('Your login name (the user your token belongs to): ').strip()
    realm = default or realms[0]
    if len(realms) > 1:
        realm = input('Realm [%s]: ' % realm).strip() or realm
    return login, realm


def owner_code(summary):
    print('Found %d token(s) for that user, %d usable. Any one of them may be used.' % (len(summary['tokens']), summary['usable']))
    return local.offline.private_prompt('Token PIN + current code, as typed at login (hidden): ')


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--transfer-receipt', required=True)
    parser.add_argument('--receipt', required=True)
    args = parser.parse_args()
    os.umask(0o077)
    fd = os.open(args.receipt, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW, 0o600)
    result = {'schema': 'netkingdom.isolated-privacyidea-mfa.v1', 'status': 'failed',
              'production_contacted': False, 'secret_values_exported': False,
              'existing_token_mfa_proven': False, 'user_to_token_resolver_proven': False,
              'container_cleanup_verified': False, 'temporary_plaintext_cleanup_verified': False,
              'source_sha256': hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
              'reused_source_sha256': {name: hashlib.sha256((HERE / name).read_bytes()).hexdigest()
                                       for name in ('isolated-privacyidea.py', 'isolated-lldap.py',
                                                    'identity-bundle.py', 'verify-identity.py',
                                                    'retrieve-identity-offline.py')}}
    keyfd = None
    work = None
    stage = 'target_preflight'
    try:
        result['target'] = preflight()
        stage = 'transfer_validation'
        transfer = json.loads(Path(args.transfer_receipt).read_text())
        local.offline.validate(transfer)
        if transfer.get('independent_owner_retrieval_verified') is not True:
            raise bundle.Refused()
        path = Path(transfer['fetched_path'])
        bundle.private(path, max_bytes=bundle.LIMIT)
        if path.stat().st_size != transfer['ciphertext_bytes'] or local.offline.offsite.digest(path) != transfer['ciphertext_sha256']:
            raise bundle.Refused()
        keyfd = local.verify.custody.private_key_fd()
        if run(['age-keygen', '-y', f'/proc/self/fd/{keyfd}'], pass_fds=(keyfd,)).strip() != bundle.profile()['public_recipient'].encode():
            raise bundle.Refused()
        with tempfile.TemporaryDirectory(prefix='nk49-pi-recovery-', dir='/dev/shm') as temporary:
            work = Path(temporary)
            archive = work / 'archive.tar'
            stage = 'archive_decryption'
            with archive.open('xb') as output:
                p = subprocess.run(['age', '-d', '-i', f'/proc/self/fd/{keyfd}', str(path)],
                                   stdout=output, stderr=subprocess.DEVNULL, pass_fds=(keyfd,), timeout=180)
            if p.returncode or archive.stat().st_size > bundle.LIMIT + 4 * 1024 * 1024:
                raise bundle.Refused()
            restored = work / 'restored'
            restored.mkdir(mode=0o700)
            stage = 'archive_validation'
            manifest, _ = bundle.restore(archive, restored)
            result.update(ciphertext_sha256=transfer['ciphertext_sha256'],
                          manifest_file_count=len(manifest['files']), all_manifest_hashes_verified=True)
            stage = 'native_mfa_check'
            restore_and_check(restored, result, owner_user, owner_code)
        if work.exists():
            raise bundle.Refused()
        result.update(status='isolated_privacyidea_mfa_verified', temporary_plaintext_cleanup_verified=True)
    except (Exception, KeyboardInterrupt):
        result.update(status='failed', error='isolated_privacyidea_mfa_failed_private_diagnostics_suppressed')
        result.setdefault('failed_stage', stage)
    finally:
        if keyfd is not None:
            os.close(keyfd)
        if work is not None and not work.exists():
            result['temporary_plaintext_cleanup_verified'] = True
        result['finished_at'] = datetime.now(timezone.utc).isoformat()
        with os.fdopen(fd, 'w') as output:
            json.dump(result, output, indent=2)
            output.write('\n')
    ok = result['status'] == 'isolated_privacyidea_mfa_verified'
    print('OK: restored token decrypted and passed the native MFA check; wrong key and wrong code were refused; cleanup passed.'
          if ok else 'ERROR: isolated privacyIDEA MFA check refused; review the safe receipt. No production change.')
    return 0 if ok else 1


if __name__ == '__main__':
    for sig in (signal.SIGINT, signal.SIGTERM, signal.SIGHUP):
        signal.signal(sig, local.offline.interrupted)
    raise SystemExit(main())