from __future__ import annotations import asyncio import os from datetime import UTC, datetime from pathlib import Path import pytest from alembic import command from backup_tool.cli import ( build_alembic_config, import_recovery_payload, ) from backup_tool.cli import ( main as cli_main, ) from backup_tool.config import Settings from backup_tool.db.engine import create_engine from backup_tool.db.models import ( Backup, Execution, Job, Repository, RepositoryDataKeyEpoch, Restore, Source, ) from backup_tool.ids import new_uuid7 from backup_tool.repository import initialize from backup_tool.security.recovery_bundle import RecoveryBundleError, decrypt_bundle from backup_tool.snapshot import finalize_publication, publish_full_snapshot from backup_tool.worker import Worker from sqlalchemy import select from sqlalchemy.ext.asyncio import async_sessionmaker from tests.conftest import make_settings PASSPHRASE = b"correct horse battery staple" def _passphrase_fd(value: bytes) -> int: read_fd, write_fd = os.pipe() os.write(write_fd, value + b"\n") os.close(write_fd) return read_fd def _recovered_settings(tmp_path: Path, original: Settings) -> Settings: data_dir = tmp_path / "recovered-data" source_root = tmp_path / "recovered-sources" restore_root = tmp_path / "recovered-restores" for path in (data_dir, source_root, restore_root): path.mkdir() master_key = tmp_path / "recovered-master.key" master_key.write_bytes(b"recovered-host-master-key-material-32-bytes") master_key.chmod(0o600) return Settings( data_dir=data_dir, database_url=f"sqlite+aiosqlite:///{data_dir / 'metadata.db'}", repository_roots=original.repository_roots, local_source_roots=(source_root,), restore_roots=(restore_root,), master_key_file=master_key, min_free_bytes=1, ) def test_recovery_import_restores_encrypted_snapshot_after_host_loss( tmp_path: Path, ) -> None: original_path = tmp_path / "original" original_path.mkdir() original = make_settings(original_path) command.upgrade(build_alembic_config(original), "head") source_root = original.local_source_roots[0] / "project" source_root.mkdir() plaintext = b"recovery host-loss content\n" (source_root / "document.txt").write_bytes(plaintext) initialized = initialize(original, "encrypted", "none", "aes-256-gcm") assert initialized.data_key_id is not None expected_created_at = datetime(2024, 1, 2, 3, 4, 5, tzinfo=UTC) expected_tombstoned_at = datetime(2024, 2, 3, 4, 5, 6, tzinfo=UTC) async def create_snapshot() -> tuple[str, str]: engine = create_engine(original) sessions = async_sessionmaker(engine, expire_on_commit=False) try: async with sessions() as db: repository = Repository( name="encrypted", root=str(initialized.root), format_version=initialized.format_version, compression=initialized.compression, encryption=initialized.encryption, signing_key_id=initialized.signing_key_id, signing_public_key=initialized.signing_public_key, active_data_key_id=initialized.data_key_id, ) source = Source( name="project", kind="local", public_config={"root": str(source_root)}, secret_refs=[], ) db.add_all([repository, source]) await db.flush() db.add( RepositoryDataKeyEpoch( repository_id=repository.id, key_id=initialized.data_key_id, state="active", ) ) job = Job( name="encrypted-job", source_id=source.id, repository_id=repository.id, requested_mode="full", exclusions=[], retention={}, allow_empty=False, ) db.add(job) await db.flush() execution = Execution(job_id=job.id, trigger="manual", progress={}) db.add(execution) await db.flush() backup = await publish_full_snapshot( original, db, execution, job, source, repository ) backup.created_at = expected_created_at execution.state = "committed" tombstoned_execution = Execution(job_id=job.id, trigger="manual", progress={}) db.add(tombstoned_execution) await db.flush() tombstoned_backup = Backup( execution_id=tombstoned_execution.id, parent_backup_id=None, manifest_id=str(new_uuid7()), manifest_digest="0" * 64, logical_bytes=0, stored_bytes=0, integrity="verified", data_key_id=initialized.data_key_id, tombstoned_at=expected_tombstoned_at, created_at=expected_created_at, ) db.add(tombstoned_backup) await db.commit() finalize_publication(initialized.root, execution.id) return backup.id, tombstoned_backup.id finally: await engine.dispose() original_backup_id, tombstoned_backup_id = asyncio.run(create_snapshot()) bundle = tmp_path / "offline.btrec" export_fd = _passphrase_fd(PASSPHRASE) try: assert ( cli_main( [ "admin", "recovery", "export", "--output", str(bundle), "--passphrase-fd", str(export_fd), ], settings=original, ) == 0 ) finally: os.close(export_fd) retry_path = tmp_path / "retry-recovery" retry_path.mkdir() retry = _recovered_settings(retry_path, original) command.upgrade(build_alembic_config(retry), "head") retry_payload = decrypt_bundle(bundle.read_bytes(), PASSPHRASE) def interrupted_after_key_install() -> None: raise OSError("simulated crash before metadata commit") with pytest.raises(RecoveryBundleError, match="recovery import failed"): asyncio.run( import_recovery_payload( retry, retry_payload, after_key_install=interrupted_after_key_install, ) ) assert not list((retry.data_dir / "repository-keys").glob("*")) assert not list((retry.data_dir / "repository-data-keys").glob("*")) assert asyncio.run(import_recovery_payload(retry, retry_payload)) == 1 unsafe_path = tmp_path / "unsafe-recovery" unsafe_path.mkdir() unsafe = _recovered_settings(unsafe_path, original) command.upgrade(build_alembic_config(unsafe), "head") unsafe_payload = decrypt_bundle(bundle.read_bytes(), PASSPHRASE) unsafe_payload["catalog"]["repositories"][0]["root"] = str(tmp_path) with pytest.raises(RecoveryBundleError, match="recovery import failed"): asyncio.run(import_recovery_payload(unsafe, unsafe_payload)) assert not (unsafe.data_dir / "repository-keys").exists() assert not (unsafe.data_dir / "repository-data-keys").exists() recovered = _recovered_settings(tmp_path, original) command.upgrade(build_alembic_config(recovered), "head") import_fd = _passphrase_fd(PASSPHRASE) try: assert ( cli_main( [ "admin", "recovery", "import", "--input", str(bundle), "--passphrase-fd", str(import_fd), ], settings=recovered, ) == 0 ) finally: os.close(import_fd) async def restore_and_assert() -> None: engine = create_engine(recovered) sessions = async_sessionmaker(engine, expire_on_commit=False) try: async with sessions() as db: backup = await db.get(Backup, original_backup_id) tombstoned_backup = await db.get(Backup, tombstoned_backup_id) job = await db.scalar(select(Job)) recovered_source = await db.scalar(select(Source)) assert backup is not None assert tombstoned_backup is not None assert backup.created_at == expected_created_at assert tombstoned_backup.created_at == expected_created_at assert tombstoned_backup.tombstoned_at == expected_tombstoned_at assert job is not None assert recovered_source is not None assert recovered_source.state == "unavailable" assert job.state == "archived" assert not job.enabled restore = Restore( backup_id=backup.id, destination=str(recovered.restore_roots[0] / "restored"), selection=[], overwrite_policy="fail", ) db.add(restore) await db.commit() worker = Worker(recovered, owner="host-loss-restore") try: assert await worker.run_once() finally: await worker.engine.dispose() finally: await engine.dispose() asyncio.run(restore_and_assert()) assert (recovered.restore_roots[0] / "restored" / "document.txt").read_bytes() == plaintext payload = decrypt_bundle(bundle.read_bytes(), PASSPHRASE) with pytest.raises(RecoveryBundleError, match="destination is not empty"): asyncio.run(import_recovery_payload(recovered, payload))