#!/usr/bin/env python3
"""Strict data-bearing identity archive format, native checks and private sealing.

No custody key is used to create a bundle. Its manifest and plaintext digests
stay inside encryption. Public receipts expose only ciphertext provenance.
"""
import base64
from datetime import datetime,timezone
import hashlib
import importlib.util
import io
import json
import os
from pathlib import Path,PurePosixPath
import re
import sqlite3
import stat
import subprocess
import tarfile
import tempfile

ROOT=Path(__file__).resolve().parents[2]
PROFILE=ROOT/'tools/bootstrap-recovery/identity-backup-profile.json'
LIMIT=512*1024*1024
FILE_LIMIT=256*1024*1024
NAME=re.compile(r'[A-Za-z0-9_.-]{1,160}\Z')


class Refused(Exception):pass


def profile():return json.loads(PROFILE.read_text())


def private(path,directory=False,max_bytes=FILE_LIMIT):
    if any(p.is_symlink() for p in (path,*path.parents)):raise Refused()
    s=path.lstat()
    if s.st_uid!=os.getuid() or s.st_mode&0o077:raise Refused()
    if directory:
        if not stat.S_ISDIR(s.st_mode):raise Refused()
    elif not stat.S_ISREG(s.st_mode) or s.st_nlink!=1 or s.st_size>max_bytes:raise Refused()


def allowed(name,profile_data=None):
    p=PurePosixPath(name)
    if p.is_absolute() or any(not NAME.fullmatch(x) or x in {'.','..'} for x in p.parts):return False
    if len(p.parts)>=3 and p.parts[:2] in [('data','lldap'),('data','authelia'),('data','keycape'),('data','privacyidea')]:return True
    if name=='databases/privacyidea.dump':return True
    for c in (profile_data or profile())['components']:
        prefix=f"resources/{c['namespace']}/"
        if name==prefix+'deployment/'+c['deployment']+'.json':return True
        if any(name==prefix+'secret/'+n+'.json' for n in c['secret_refs']):return True
        if any(name==prefix+'configmap/'+n+'.json' for n in c['configmap_refs']):return True
    return False


def save(root,name,value,profile_data=None):
    if not allowed(name,profile_data) or len(value)>FILE_LIMIT or b'AGE-SECRET-KEY-' in value:raise Refused()
    target=root/name;target.parent.mkdir(mode=0o700,parents=True,exist_ok=True)
    fd=os.open(target,os.O_WRONLY|os.O_CREAT|os.O_EXCL|os.O_NOFOLLOW,0o600)
    with os.fdopen(fd,'wb') as out:out.write(value)


def import_files(raw,root,mode):
    names=set();total=0
    with tarfile.open(fileobj=io.BytesIO(raw),mode='r:') as archive:
        for member in archive:
            expected=('data/privacyidea/',) if mode=='mfa' else ('data/lldap/','data/authelia/','data/keycape/')
            if (not member.isfile() or not member.name.startswith(expected) or member.name in names
                or member.size<0 or member.size>FILE_LIMIT or len(names)>=2048):raise Refused()
            names.add(member.name);total+=member.size
            if total>LIMIT:raise Refused()
            data=archive.extractfile(member).read()
            if len(data)!=member.size:raise Refused()
            save(root,member.name,data)


def resource(root,ns,kind,name):
    p=root/f'resources/{ns}/{kind}/{name}.json'
    data=json.loads(p.read_bytes())
    if data['metadata']['name']!=name or data['metadata']['namespace']!=ns:raise Refused()
    return data


def checks(root,profile_data=None):
    # Checks describe stored formats/integrity, not current login validity.
    for c in (profile_data or profile())['components']:
        deployment=resource(root,c['namespace'],'deployment',c['deployment'])
        if deployment.get('kind')!='Deployment' or deployment['metadata']['uid']!=c['deployment_uid']:raise Refused()
        for name in c['secret_refs']:
            data=resource(root,c['namespace'],'secret',name)
            if data.get('kind')!='Secret' or not isinstance(data.get('data'),dict) or not data['data']:raise Refused()
            for value in data['data'].values():
                if b'AGE-SECRET-KEY-' in base64.b64decode(value,validate=True):raise Refused()
        for name in c['configmap_refs']:
            if resource(root,c['namespace'],'configmap',name).get('kind')!='ConfigMap':raise Refused()
    for name,file in [('lldap','users.db'),('authelia','db.sqlite3')]:
        db=root/f'data/{name}/{file}'
        with sqlite3.connect(db.as_uri()+'?mode=ro',uri=True) as con:
            if con.execute('PRAGMA integrity_check').fetchall()!=[('ok',)]:raise Refused()
    with (root/'databases/privacyidea.dump').open('rb') as source:
        if source.read(5)!=b'PGDMP':raise Refused()
    spec=importlib.util.spec_from_file_location('pi_keys',ROOT/'sso-mfa/k8s/privacyidea/key-custody-remote.py')
    keys=importlib.util.module_from_spec(spec);spec.loader.exec_module(keys)
    keys.validate({name:(root/'data/privacyidea'/name).read_bytes() for name in ('enckey','private.pem','public.pem')})
    from cryptography.hazmat.primitives import serialization
    from cryptography.hazmat.primitives.asymmetric import rsa
    config=resource(root,'sso','secret','keycape-config')['data']
    if not base64.b64decode(config['config.yaml'],validate=True):raise Refused()
    signing=serialization.load_pem_private_key(base64.b64decode(config['key.pem'],validate=True),password=None)
    if not isinstance(signing,rsa.RSAPrivateKey) or signing.key_size<2048:raise Refused()
    # Policy must exist, even for an empty set of receipts; do not silently
    # restore an issuer without its current authentication policy.
    policy=json.loads((root/'data/keycape/policy.json').read_bytes())
    if not isinstance(policy,dict):raise Refused()
    return {'sqlite_integrity_passed':True,'privacyidea_key_format_and_pair_passed':True,
            'keycape_signing_key_format_passed':True,'policy_json_present':True,
            'database_custom_format_present':True,'native_identity_recovery_proven':False}


def inventory(root):
    private(root,True);entries=[];total=0
    for path in sorted(root.rglob('*')):
        if path.is_dir():private(path,True);continue
        private(path);name=path.relative_to(root).as_posix()
        if not allowed(name):raise Refused()
        with path.open('rb') as source:value=source.read(FILE_LIMIT+1)
        if len(value)>FILE_LIMIT or b'AGE-SECRET-KEY-' in value:raise Refused()
        total+=len(value)
        if total>LIMIT or len(entries)>=4096:raise Refused()
        entries.append({'path':name,'bytes':len(value),'sha256':hashlib.sha256(value).hexdigest()})
    if not entries:raise Refused()
    return entries


def seal(root,output,capture):
    output=output.absolute();private(output.parent,True)
    if output.exists() or output.is_symlink():raise Refused()
    checked=checks(root);entries=inventory(root)
    manifest={'schema':'netkingdom.identity-archive.v1','profile_sha256':hashlib.sha256(PROFILE.read_bytes()).hexdigest(),
              'profile':profile(),
              'captured_at':datetime.now(timezone.utc).isoformat(),'capture':capture,'files':entries,'checks':checked}
    with tempfile.TemporaryDirectory(prefix='.identity-seal-',dir=output.parent) as temporary:
        ciphertext=Path(temporary)/'ciphertext'
        fd=os.open(ciphertext,os.O_WRONLY|os.O_CREAT|os.O_EXCL,0o600)
        proc=subprocess.Popen(['age','-r',profile()['public_recipient']],stdin=subprocess.PIPE,stdout=fd,stderr=subprocess.DEVNULL)
        os.close(fd)
        try:
            with tarfile.open(fileobj=proc.stdin,mode='w|') as archive:
                raw=json.dumps(manifest,sort_keys=True).encode();m=tarfile.TarInfo('manifest.json');m.size=len(raw);m.mode=0o600
                archive.addfile(m,io.BytesIO(raw))
                for entry in entries:
                    with (root/entry['path']).open('rb') as source:
                        value=source.read(FILE_LIMIT+1)
                    if hashlib.sha256(value).hexdigest()!=entry['sha256']:raise Refused()
                    m=tarfile.TarInfo(entry['path']);m.size=len(value);m.mode=0o600
                    archive.addfile(m,io.BytesIO(value))
            proc.stdin.close()
            if proc.wait(timeout=120) or not 0<ciphertext.stat().st_size<=LIMIT:raise Refused()
            # Check source completeness after encryption too; changed/new files
            # refuse publication. This is not same-UID adversary isolation.
            if inventory(root)!=entries:raise Refused()
            with ciphertext.open('rb') as source:
                os.fsync(source.fileno());digest=hashlib.file_digest(source,'sha256').hexdigest()
            os.link(ciphertext,output)
        finally:
            if proc.poll() is None:proc.kill();proc.wait()
    return {'schema':'netkingdom.identity-capture.v1','status':'encrypted_bundle_created',
            'ciphertext_sha256':digest,'ciphertext_bytes':output.stat().st_size,'ciphertext_path':str(output),
            'profile_sha256':manifest['profile_sha256'],'file_count':len(entries),
            'checks':checked,'offsite_recovery_proven':False,'secret_values_exported':False}


def restore(archive_path,root):
    private(root,True)
    with tarfile.open(archive_path,mode='r:') as archive:
        first=archive.next()
        if first is None or first.name!='manifest.json' or not first.isfile() or not 0<first.size<2*1024*1024:raise Refused()
        manifest=json.load(archive.extractfile(first))
        if manifest.get('schema')!='netkingdom.identity-archive.v1':raise Refused()
        captured=manifest['profile']
        if captured.get('schema')!='netkingdom.identity-backup-profile.v1':raise Refused()
        if len(captured['components'])!=4:raise Refused()
        if {(c['namespace'],c['deployment']) for c in captured['components']}!={('sso','lldap'),('sso','authelia'),('sso','keycape'),('mfa','privacyidea')}:raise Refused()
        for c in captured['components']:
            for key in ('secret_refs','configmap_refs'):
                if len(set(c[key]))!=len(c[key]) or any(not NAME.fullmatch(n) or n in {'.','..'} for n in c[key]):raise Refused()
        entries=manifest['files'];expected={e['path']:e for e in entries}
        if len(expected)!=len(entries) or len(entries)>4096:raise Refused()
        seen=set();total=0
        # Iteration starts with the already-read first member; explicitly skip
        # only that member. A later duplicate manifest is never accepted.
        for index,member in enumerate(archive):
            if index==0 and member is first:continue
            name=member.name
            if (not member.isfile() or name in seen or name not in expected or not allowed(name,captured)
                or not 0<=member.size<=FILE_LIMIT or member.size!=expected[name]['bytes']):raise Refused()
            seen.add(name);total+=member.size
            if total>LIMIT:raise Refused()
            raw=archive.extractfile(member).read()
            if len(raw)!=member.size or hashlib.sha256(raw).hexdigest()!=expected[name]['sha256']:raise Refused()
            save(root,name,raw,captured)
        if seen!=set(expected):raise Refused()
    return manifest,checks(root,captured)


def require_runtime_host():
    import socket
    host=profile()['runtime_host']
    if (os.geteuid()!=0 or socket.gethostname()!=host['hostname']
        or hashlib.sha256(Path('/etc/machine-id').read_bytes()).hexdigest()!=host['machine_id_sha256']):raise Refused()