coulomb-social/scripts/restore-paired-snapshot.py

#!/usr/bin/env python3
"""Restore a private snapshot into isolated Docker PostgreSQL; never production.

No published ports; an internal network blocks external integrations. Compare
all restored COPY rows and sequence state, and all archived file hashes. A
successful empty-content restore alone does not prove real corpus recovery.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import os
import re
import subprocess
import sys
import tarfile
import tempfile
import time
from pathlib import Path


class Refused(RuntimeError):
    pass


def run(args: list[str], *, payload: bytes | None = None) -> bytes:
    p = subprocess.run(args, input=payload, capture_output=True, timeout=180)
    if p.returncode:
        raise Refused('command_failed')
    return p.stdout


def data_fingerprint(sql: bytes) -> dict:
    """Ignore SQL boilerplate, compare COPY data and explicit sequence state."""
    tables = {}
    sequences = []
    current = None
    rows = []
    for line in sql.splitlines():
        if current is not None:
            if line == b'\\.':
                tables[current] = {
                    'rows': len(rows),
                    'sha256': hashlib.sha256(b'\n'.join(sorted(rows))).hexdigest(),
                }
                current, rows = None, []
            else:
                rows.append(line)
        elif line.startswith(b'COPY ') and line.endswith(b' FROM stdin;'):
            current = line.decode('utf-8')
        elif line.startswith(b'SELECT pg_catalog.setval('):
            sequences.append(line.decode('utf-8'))
    if current is not None:
        raise Refused('unterminated_copy')
    if not tables:
        raise Refused('no_tables_in_dump')
    return {'tables': tables, 'sequences': sorted(sequences)}


def content_fingerprint(archive: Path, destination: Path) -> dict[str, str]:
    hashes = {}
    with tarfile.open(archive, 'r:gz') as t:
        for entry in t.getmembers():
            parts = Path(entry.name).parts
            if (not parts or parts[0] != 'content' or '..' in parts
                    or Path(entry.name).is_absolute()):
                raise Refused('invalid_archive_path')
            if not (entry.isfile() or entry.isdir()):
                raise Refused('unsupported_archive_entry')
            if entry.isfile():
                hashes[entry.name] = hashlib.sha256(t.extractfile(entry).read()).hexdigest()
        t.extractall(destination, filter='data')
    actual = {str(p.relative_to(destination)): hashlib.sha256(p.read_bytes()).hexdigest()
              for p in (destination / 'content').rglob('*') if p.is_file()}
    if actual != hashes:
        raise Refused('content_hash_mismatch')
    return hashes


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('snapshot', type=Path)
    parser.add_argument('--postgres-image', default='postgres:16',
                        help='Local PostgreSQL image/digest; use the retained kit image ID offline')
    parser.add_argument('--keep-instance', action='store_true',
                        help='Keep the isolated internal network/database for a content rehearsal')
    parser.add_argument('--work-root', type=Path, default=Path('/tmp'),
                        help='Existing absolute parent for the private restored content directory')
    args = parser.parse_args()
    if not args.work_root.is_absolute() or not args.work_root.is_dir():
        raise Refused('absolute_existing_work_root_required')
    source = args.snapshot.resolve()
    receipt = json.loads((source / 'receipt.json').read_text())
    if receipt.get('status') != 'captured' or not receipt.get('writers_quiesced'):
        raise Refused('snapshot_not_paired')
    for name, key in [('database.dump', 'database_sha256'), ('content.tar.gz', 'content_sha256')]:
        if hashlib.sha256((source / name).read_bytes()).hexdigest() != receipt[key]:
            raise Refused('artifact_checksum_mismatch')
    started = time.monotonic()
    name = 'csoc-recovery-' + time.strftime('%Y%m%d%H%M%S', time.gmtime())
    network = name + '-net'
    work = Path(tempfile.mkdtemp(prefix=name + '-', dir=args.work_root))
    image = run(['docker', 'image', 'inspect', args.postgres_image,
                 '--format', '{{.Id}}']).decode().strip()
    if not re.fullmatch(r'sha256:[a-f0-9]{64}', image):
        raise Refused('unresolved_postgres_image')
    created = False
    network_created = False
    verified = False
    try:
        run(['docker', 'network', 'create', '--internal', network])
        network_created = True
        run(['docker', 'run', '-d', '--name', name, '--network', network,
             '--network-alias', 'recovery-db', '-e', 'POSTGRES_HOST_AUTH_METHOD=trust',
             '-e', 'POSTGRES_DB=coulomb_social_restore', image])
        created = True
        for _ in range(45):
            # The image's temporary initialization server accepts Unix-socket
            # connections before creating the target DB and restarting. Require
            # a query against the final TCP server before starting pg_restore.
            p = subprocess.run(['docker', 'exec', name, 'psql', '-h', '127.0.0.1',
                                '-U', 'postgres', '-d', 'coulomb_social_restore',
                                '-tAc', 'SELECT 1'],
                               capture_output=True)
            if p.returncode == 0:
                break
            time.sleep(1)
        else:
            raise Refused('restore_db_not_ready')
        archive = (source / 'database.dump').read_bytes()
        original_sql = run(['docker', 'exec', '-i', name, 'pg_restore', '--data-only', '-f', '-'],
                           payload=archive)
        baseline = data_fingerprint(original_sql)
        run(['docker', 'exec', '-i', name, 'pg_restore', '-U', 'postgres',
             '-d', 'coulomb_social_restore', '--no-owner', '--no-acl', '--exit-on-error'],
            payload=archive)
        restored_sql = run(['docker', 'exec', name, 'pg_dump', '-U', 'postgres',
                            '-d', 'coulomb_social_restore', '--data-only',
                            '--no-owner', '--no-acl'])
        if data_fingerprint(restored_sql) != baseline:
            raise Refused('database_fingerprint_mismatch')
        files = content_fingerprint(source / 'content.tar.gz', work)
        result = {'status': 'restored_verified', 'database_tables': len(baseline['tables']),
                  'database_rows': sum(t['rows'] for t in baseline['tables'].values()),
                  'database_rows_and_sequences_match': True, 'content_files': len(files),
                  'content_hashes_match': True, 'network_internal': True, 'published_ports': 0,
                  'postgres_image_id': image, 'elapsed_seconds': round(time.monotonic()-started, 2),
                  'content_was_empty': not files,
                  'database_container': name, 'network': network, 'restored_directory': str(work)}
        path = work / 'restore-receipt.json'
        fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
        with os.fdopen(fd, 'w') as f:
            json.dump(result, f, indent=2)
        verified = True
        print(json.dumps(result))
    finally:
        if not (args.keep_instance and verified):
            if created:
                run(['docker', 'rm', '-f', '-v', name])
            if network_created:
                run(['docker', 'network', 'rm', network])
    return 0


if __name__ == '__main__':
    try:
        sys.exit(main())
    except Exception as exc:
        print(json.dumps({'status': 'failed', 'failure_type': type(exc).__name__}))
        sys.exit(1)