#!/usr/bin/env python3
"""Restore the fetched PI database in an isolated, disposable local PostgreSQL.

Run only through the attended custody helper. No cluster, production database,
network, ports, host data mounts, IAM login or token decryption is involved.
"""
import argparse
from datetime import datetime, timezone
import hashlib
import importlib.util
import json
import os
from pathlib import Path
import signal
import subprocess
import sys
import tempfile
import time
import uuid

ROOT=Path(__file__).resolve().parents[2]
spec=importlib.util.spec_from_file_location('custody',ROOT/'sso-mfa/bootstrap/sops-private.py')
custody=importlib.util.module_from_spec(spec);spec.loader.exec_module(custody)
RECIPIENT='age1zvryunvjhvpkmasskauga2heeg0ztnte9ymgppvjge36ekumk50syr3tsz'
IMAGE='postgres@sha256:ff23cdce56cac62ada6f66013e1a50864c0abbe79d132d40b6e05bd80f378a70'
OBJECT='r01-net-kingdom-pg/net-kingdom-pg-privacyidea_db-20261015T023000Z.dump.age'
SIZE=89129309


class Refused(Exception):pass


def run(args,**kwargs):
    result=subprocess.run(args,stdout=subprocess.PIPE,stderr=subprocess.DEVNULL,timeout=600,**kwargs)
    if result.returncode:raise Refused()
    return result.stdout


def transfer_input(receipt):
    r=json.loads(receipt.read_text())
    if (r.get('schema')!='netkingdom.privacyidea-offhost-transfer.v1' or
        r.get('status')!='ciphertext_fetched' or r.get('object')!=OBJECT or
        r.get('ciphertext_bytes')!=SIZE or r.get('object_unchanged_during_transfer') is not True):raise Refused()
    path=Path(r['ciphertext_path'])
    if not path.is_absolute() or path.parent.parent!=Path('/tmp') or not path.parent.name.startswith('nk49-offhost-'):raise Refused()
    if any(p.is_symlink() for p in (path,*path.parents)):raise Refused()
    st=path.stat()
    if st.st_uid!=os.getuid() or st.st_mode&0o077 or st.st_nlink!=1 or not path.is_file() or st.st_size!=SIZE:raise Refused()
    blob=path.read_bytes()
    if len(blob)!=SIZE or hashlib.sha256(blob).hexdigest()!=r['ciphertext_sha256']:raise Refused()
    return blob,r['ciphertext_sha256']


def restore(dump,container):
    # Require the exact image already present; do not pull through a mutable tag.
    image_id=run(['docker','image','inspect','--format','{{.Id}}',IMAGE]).decode().strip()
    run(['docker','run','--detach','--pull=never','--name',container,'--network','none',
         '--security-opt','no-new-privileges','--memory','2g','--cpus','2',
         '--tmpfs','/var/lib/postgresql/data:rw,nosuid,nodev,size=1073741824',
         '--env','POSTGRES_HOST_AUTH_METHOD=trust','--env','POSTGRES_DB=recovery',
         IMAGE,'postgres','-c','listen_addresses='])
    for _ in range(60):
        # The image starts a temporary server during initdb. Wait until the
        # entrypoint has exec'd the final server before attempting the restore.
        pid1=subprocess.run(['docker','exec',container,'cat','/proc/1/comm'],
                            stdout=subprocess.PIPE,stderr=subprocess.DEVNULL,timeout=10)
        if pid1.returncode or pid1.stdout.strip()!=b'postgres':
            time.sleep(1);continue
        ready=subprocess.run(['docker','exec',container,'pg_isready','-U','postgres','-d','recovery'],
                             stdout=subprocess.DEVNULL,stderr=subprocess.DEVNULL,timeout=10)
        if ready.returncode==0:break
        time.sleep(1)
    else:raise Refused()
    version=run(['docker','exec',container,'pg_restore','--version']).decode().strip()
    with dump.open('rb') as source:
        run(['docker','exec','-i',container,'pg_restore','-U','postgres','-d','recovery',
             '--no-owner','--no-privileges','--exit-on-error'],stdin=source)
    # Fixed aggregate queries only. No identity, resolver config or token values.
    counts={}
    for table in ('token','resolver','realm'):
        raw=run(['docker','exec',container,'psql','-X','-U','postgres','-d','recovery','-At',
                 '-v','ON_ERROR_STOP=1','-c',f'SELECT count(*) FROM public.{table};']).strip()
        if not raw.isdigit():raise Refused()
        counts[table]=int(raw)
    return {'image':IMAGE,'image_id':image_id,'pg_restore_version':version,'aggregate_row_counts':counts}


def main():
    p=argparse.ArgumentParser(description=__doc__)
    p.add_argument('--transfer-receipt',required=True);p.add_argument('--receipt',required=True)
    a=p.parse_args();os.umask(0o077)
    # Reserve the receipt before any recovery action; never overwrite evidence.
    fd=os.open(a.receipt,os.O_WRONLY|os.O_CREAT|os.O_EXCL|os.O_NOFOLLOW,0o600)
    r={'schema':'netkingdom.privacyidea-database-drill.v1','status':'failed',
       'started_at':datetime.now(timezone.utc).isoformat(),'source_sha256':hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
       'production_contacted':False,'secret_values_exported':False,'native_identity_recovery_proven':False,
       'token_decryption_proven':False,'audit_signature_continuity_proven':False,
       'full_identity_file_coverage_proven':False,'container_cleanup_verified':False,
       'temporary_plaintext_cleanup_verified':False}
    container='nk49-pi-recovery-'+uuid.uuid4().hex
    keyfd=None;work=None
    try:
        blob,digest=transfer_input(Path(a.transfer_receipt));r['ciphertext_sha256']=digest
        keyfd=custody.private_key_fd()
        derived=run(['age-keygen','-y',f'/proc/self/fd/{keyfd}'],pass_fds=(keyfd,)).strip()
        if derived!=RECIPIENT.encode():raise Refused()
        r['registered_recipient_match']=True
        with tempfile.TemporaryDirectory(prefix='nk49-pi-drill-',dir='/tmp') as temporary:
            work=Path(temporary)
            env={'PATH':os.environ.get('PATH',os.defpath),'HOME':temporary,'XDG_CONFIG_HOME':temporary}
            wrong=work/'wrong.key'
            run(['age-keygen','-o',str(wrong)],env=env)
            negative=subprocess.run(['age','-d','-i',str(wrong)],input=blob,stdout=subprocess.DEVNULL,
                                    stderr=subprocess.DEVNULL,timeout=60,env=env)
            if negative.returncode==0:raise Refused()
            r['wrong_key_refused']=True
            dump=work/'privacyidea.dump'
            with dump.open('xb') as output:
                result=subprocess.run(['age','-d','-i',f'/proc/self/fd/{keyfd}'],input=blob,
                    stdout=output,stderr=subprocess.DEVNULL,timeout=120,pass_fds=(keyfd,),env=env)
            if result.returncode or not 5<dump.stat().st_size<SIZE:raise Refused()
            with dump.open('rb') as source:
                if source.read(5)!=b'PGDMP':raise Refused()
            r['offhost_ciphertext_decryption_verified']=True
            try:
                r.update(restore(dump,container));r['database_restore_and_readback_passed']=True
            finally:
                removed=subprocess.run(['docker','rm','--force','--volumes',container],stdout=subprocess.DEVNULL,
                                       stderr=subprocess.DEVNULL,timeout=60)
                # A removal error is a failed ceremony even if the restore passed.
                if removed.returncode:raise Refused()
                absent=subprocess.run(['docker','container','inspect',container],stdout=subprocess.DEVNULL,
                                      stderr=subprocess.DEVNULL,timeout=15)
                # Confirm the daemon remains reachable, so inspect failure means absent.
                run(['docker','info','--format','{{.ServerVersion}}'])
                if absent.returncode==0:raise Refused()
                r['container_cleanup_verified']=True
        if work.exists():raise Refused()
        r['temporary_plaintext_cleanup_verified']=True;r['status']='isolated_database_drill_passed'
    except Exception:
        r['error']='database_drill_failed_diagnostics_suppressed'
    finally:
        if keyfd is not None:os.close(keyfd)
        if work is not None and not work.exists():r['temporary_plaintext_cleanup_verified']=True
        r['finished_at']=datetime.now(timezone.utc).isoformat()
        with os.fdopen(fd,'w') as out:json.dump(r,out,indent=2);out.write('\n')
    if r['status']=='isolated_database_drill_passed':
        print('OK: off-host database decrypted, isolated restore/readback passed, disposable data removed. Full identity recovery remains open.')
        return 0
    print('ERROR: database drill failed; private diagnostics suppressed. Retain the safe receipt for review.',file=sys.stderr)
    return 1


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