#!/usr/bin/env python3
"""Private in-helper PVC capture. Stdout is a secret-bearing tar: parent MUST capture.

Sources are mounted read-only; SQLite uses its online backup API, never raw copy
or immutable=1. No production file is generated, replaced, chmod'd or removed.
"""
import hashlib
import io
import os
from pathlib import Path
import signal
import sqlite3
import stat
import sys
import tarfile
import tempfile

MAX_FILE=256*1024*1024
MAX_TOTAL=512*1024*1024
MAX_COUNT=2048


class Refused(Exception):pass


def read(path):
    if any(p.is_symlink() for p in (path,*path.parents)):raise Refused()
    fd=os.open(path,os.O_RDONLY|os.O_NOFOLLOW|os.O_NONBLOCK)
    try:
        before=os.fstat(fd)
        if not stat.S_ISREG(before.st_mode) or before.st_nlink!=1 or before.st_size>MAX_FILE:raise Refused()
        with os.fdopen(os.dup(fd),'rb') as stream:data=stream.read(MAX_FILE+1)
        after=os.stat(path,follow_symlinks=False)
        if (len(data)>MAX_FILE or b'AGE-SECRET-KEY-' in data or
            any(getattr(before,k)!=getattr(after,k) for k in ('st_dev','st_ino','st_size','st_mtime_ns','st_ctime_ns'))):raise Refused()
        return data
    finally:os.close(fd)


def tree(root,database=None,omit_notification=False,omit_backups=True):
    result={}
    for directory,dirs,files in os.walk(root,followlinks=False):
        path=Path(directory)
        if path.is_symlink():raise Refused()
        # Historical local snapshots are redundant; do not recursively include
        # accumulated backups. WAL/SHM are represented by the native snapshot.
        if omit_backups:dirs[:]=[name for name in dirs if name!='backups']
        for name in dirs:
            if (path/name).is_symlink():raise Refused()
        for name in files:
            relative=(path/name).relative_to(root).as_posix()
            if database and relative in {database,database+'-wal',database+'-shm',database+'-journal'}:continue
            if omit_notification and relative=='notification.txt':continue
            result[relative]=read(path/name)
            if len(result)>MAX_COUNT or sum(map(len,result.values()))>MAX_TOTAL:raise Refused()
    return result


def snapshot(source,target):
    if source.is_symlink() or not source.is_file():raise Refused()
    # A read-only URI intentionally fails if the engine needs a write to its
    # production WAL/SHM. Never fall back to a raw copy or mutate the mount.
    with sqlite3.connect(source.as_uri()+'?mode=ro',uri=True,timeout=10) as origin:
        with sqlite3.connect(target) as dest:
            origin.backup(dest,pages=128,sleep=0.05)
            if dest.execute('PRAGMA integrity_check').fetchall()!=[('ok',)]:raise Refused()
    return read(target)


def capture(mode,root=Path('/source'),scratch=Path('/scratch')):
    values={}
    with tempfile.TemporaryDirectory(prefix='snapshot-',dir=scratch) as temporary:
        if mode=='sso':
            for name,database in [('lldap','users.db'),('authelia','db.sqlite3'),('keycape',None)]:
                source=root/name
                if source.is_symlink() or not source.is_dir():raise Refused()
                first=tree(source,database,name=='authelia',name!='keycape')
                if database:first[database]=snapshot(source/database,Path(temporary)/(name+'.sqlite'))
                second=tree(source,database,name=='authelia',name!='keycape')
                if {k:v for k,v in first.items() if k!=database}!=second:raise Refused()
                values.update({f'data/{name}/{k}':v for k,v in first.items()})
        elif mode=='mfa':
            first=tree(root/'privacyidea')
            if not {'enckey','private.pem','public.pem'}<=set(first) or first!=tree(root/'privacyidea'):raise Refused()
            values={f'data/privacyidea/{name}':value for name,value in first.items()}
        else:raise Refused()
        if len(values)>MAX_COUNT or sum(map(len,values.values()))>MAX_TOTAL:raise Refused()
        return values


def emit(values,destination):
    with tarfile.open(fileobj=destination,mode='w|') as archive:
        for name,value in sorted(values.items()):
            member=tarfile.TarInfo(name);member.size=len(value);member.mode=0o600
            archive.addfile(member,io.BytesIO(value))


def interrupted(*args):raise Refused()


if __name__=='__main__':
    for sig in (signal.SIGTERM,signal.SIGHUP,signal.SIGINT):signal.signal(sig,interrupted)
    try:
        os.umask(0o077)
        if len(sys.argv)!=2:raise Refused()
        emit(capture(sys.argv[1]),sys.stdout.buffer)
    except Exception:
        print('ERROR: private file capture refused; diagnostics suppressed.',file=sys.stderr)
        raise SystemExit(1)