diff --git a/pyproject.toml b/pyproject.toml index 9e52e80..e7d8c9b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,11 +27,13 @@ dev = [ "ruff>=0.6", "mypy>=2.3.1", "pre-commit>=3.7", + "keyrings.alt>=4.0", ] test = [ "pytest>=9.1.1", "pytest-asyncio>=0.23", "hypothesis>=6.100", + "keyrings.alt>=4.0", ] [project.scripts] diff --git a/src/stash/cli/commands/put.py b/src/stash/cli/commands/put.py index 48d738a..b4a3710 100644 --- a/src/stash/cli/commands/put.py +++ b/src/stash/cli/commands/put.py @@ -1,6 +1,7 @@ """CLI command: put - Store a file.""" import asyncio +import time from pathlib import Path import click @@ -12,9 +13,12 @@ from stash.core.jobs import JobConfig from stash.core.keymanager import KeyManager from stash.core.manifest import ( + ChunkInfo, + ChunkStatus, DistributionStrategy, EncryptionInfo, ManifestBuilder, + UploadStatus, compute_checksum, generate_file_id, ) @@ -27,11 +31,13 @@ @click.option("--provider", help="Specific provider to use (default: first available)") @click.option("--chunk-size", type=int, help="Chunk size in bytes (default: provider limit)") @click.option("--strategy", type=click.Choice(["single", "split", "balanced", "replicated"]), default="single", help="Distribution strategy") +@click.option("--resume", is_flag=True, help="Resume an incomplete upload") +@click.option("--file-id", help="File ID to resume (required with --resume if multiple incomplete uploads exist)") @click.option("--confirm/--no-confirm", default=True, help="Confirm before upload") @click.pass_context -def put_cmd(ctx: click.Context, file_path: Path, provider: str | None, chunk_size: int | None, strategy: str, confirm: bool) -> None: +def put_cmd(ctx: click.Context, file_path: Path, provider: str | None, chunk_size: int | None, strategy: str, resume: bool, file_id: str | None, confirm: bool) -> None: """Store a file in Stash.""" - asyncio.run(_put_async(file_path, ctx.obj["repo"], provider, chunk_size, strategy, confirm)) + asyncio.run(_put_async(file_path, ctx.obj["repo"], provider, chunk_size, strategy, confirm, resume, file_id)) async def _put_async( @@ -41,6 +47,8 @@ async def _put_async( chunk_size: int | None, strategy: str, do_confirm: bool, + resume: bool, + file_id: str | None, ) -> None: repo = repo_path.resolve() store = MetadataStore(repo) @@ -81,14 +89,175 @@ async def _put_async( return crypto = CryptoEngine() - file_id = generate_file_id() - file_key = crypto.generate_file_key(rmk) + + # Handle resume logic + existing_manifest = None + file_id = None + file_key = None + encrypted_name = None + encrypted_name_nonce = None + file_size = file_path.stat().st_size + + if resume: + # Find existing incomplete manifest + if file_id: + # Explicit file ID provided + if not store.file_exists(file_id): + print_error(f"File with ID '{file_id}' not found") + return + existing_manifest = store.load_manifest(file_id) + else: + # Auto-detect by filename + for fid in store.list_files(): + manifest = store.load_manifest(fid) + if manifest.original_name == file_path.name and manifest.upload_status != UploadStatus.COMPLETED: + existing_manifest = manifest + break + + if existing_manifest is None: + print_error("No incomplete upload found to resume") + print_info("Run 'stash put' without --resume to start a new upload") + return + + # Verify file matches + if existing_manifest.original_size != file_path.stat().st_size: + print_error("File size does not match the incomplete upload") + raise SystemExit(1) + + # Verify file content matches (check first chunk checksum if available) + if existing_manifest.chunks: + chunker = Chunker(ChunkConfig(chunk_size=existing_manifest.chunk_size)) + first_chunk_data = next(chunker.chunk_file(file_path)).data + first_checksum = compute_checksum(first_chunk_data) + if existing_manifest.chunks[0].checksum != first_checksum: + print_error("File content does not match the incomplete upload") + raise SystemExit(1) + + file_id = existing_manifest.file_id + file_key = crypto.derive_file_key_from_rmk(rmk, file_id.encode()) + + # Restore encrypted filename info + encrypted_name = existing_manifest.encrypted_name + encrypted_name_nonce = existing_manifest.encrypted_name_nonce + + print_info(f"Resuming upload of '{existing_manifest.original_name}' ({existing_manifest.file_id})") + print_info(f"Progress: {existing_manifest.uploaded_chunks}/{existing_manifest.total_chunks} chunks uploaded") + else: + # New upload - generate new file_id and file_key + file_id = generate_file_id() + file_key = crypto.generate_file_key(rmk) + + # Encrypt the filename using the file key + encrypted_name_ciphertext, encrypted_name_nonce = crypto.encrypt_filename( + file_path.name.encode(), file_key + ) + encrypted_name = encrypted_name_ciphertext.hex() + + crypto = CryptoEngine() + if not file_key: + file_key = crypto.generate_file_key(rmk) # Encrypt the filename using the file key - encrypted_name_ciphertext, encrypted_name_nonce = crypto.encrypt_filename( - file_path.name.encode(), file_key - ) - encrypted_name = encrypted_name_ciphertext.hex() + if encrypted_name is None: + encrypted_name_ciphertext, encrypted_name_nonce = crypto.encrypt_filename( + file_path.name.encode(), file_key + ) + encrypted_name = encrypted_name_ciphertext.hex() + + provider_configs = {} + for name in provider_names: + config = store.get_provider_config(name) + if not config: + print_error(f"Provider config not found: {name}") + return + provider_configs[name] = config + + provider_instances = {} + for name, config in provider_configs.items(): + instance = await ProviderRegistry.create(config.type, config) + provider_instances[name] = instance + + limits = {name: p.get_limits() for name, p in provider_instances.items()} + max_chunk = min(l.max_chunk_size for l in limits.values()) + effective_chunk_size = min(chunk_size or max_chunk, max_chunk) + + chunker = Chunker(ChunkConfig(chunk_size=effective_chunk_size)) + num_chunks = chunker.get_num_chunks(file_size) + + dist_strategy = DistributionStrategy(strategy) + + # Initialize manifest builder + if resume and existing_manifest: + # Resume existing manifest - keep existing chunks, update status + file_id = existing_manifest.file_id + encryption_info = existing_manifest.encryption + builder = ManifestBuilder( + file_id=existing_manifest.file_id, + original_name=existing_manifest.original_name, + encrypted_name=existing_manifest.encrypted_name, + encrypted_name_nonce=existing_manifest.encrypted_name_nonce, + original_size=existing_manifest.original_size, + chunk_size=existing_manifest.chunk_size, + encryption=existing_manifest.encryption, + strategy=DistributionStrategy(existing_manifest.strategy), + ) + # Pre-populate with existing chunks + for chunk in existing_manifest.chunks: + builder.add_chunk( + index=chunk.index, + size=chunk.size, + encrypted_size=chunk.encrypted_size, + checksum=chunk.checksum, + provider=chunk.provider, + remote_id=chunk.remote_id, + nonce=chunk.nonce, + metadata=chunk.metadata, + status=chunk.status, + uploaded_at=chunk.uploaded_at, + error=chunk.error, + ) + else: + # New upload + file_id = generate_file_id() + file_key = crypto.generate_file_key(rmk) + + # Encrypt the filename using the file key + encrypted_name_ciphertext, encrypted_name_nonce = crypto.encrypt_filename( + file_path.name.encode(), file_key + ) + encrypted_name = encrypted_name_ciphertext.hex() + + encryption_info = EncryptionInfo( + algorithm="AES-256-GCM", + key_size=32, + nonce_size=12, + chunk_key_derivation="HKDF-SHA256", + file_key_salt=file_key.salt, + file_key_wrapped=None, + ) + + builder = ManifestBuilder( + file_id=file_id, + original_name=file_path.name, + encrypted_name=encrypted_name, + encrypted_name_nonce=encrypted_name_nonce, + original_size=file_size, + chunk_size=effective_chunk_size, + encryption=encryption_info, + strategy=dist_strategy, + ) + # Pre-populate all chunks as PENDING + for i in range(num_chunks): + builder.add_chunk( + index=i, + size=0, # Will be updated when chunk is uploaded + encrypted_size=0, + checksum="", + provider="", + remote_id="", + nonce=b"", + status=ChunkStatus.PENDING, + ) provider_configs = {} for name in provider_names: @@ -112,27 +281,12 @@ async def _put_async( dist_strategy = DistributionStrategy(strategy) - encryption_info = EncryptionInfo( - algorithm="AES-256-GCM", - key_size=32, - nonce_size=12, - chunk_key_derivation="HKDF-SHA256", - file_key_salt=file_key.salt, - file_key_wrapped=None, - ) - - builder = ManifestBuilder( - file_id=file_id, - original_name=file_path.name, - encrypted_name=encrypted_name, - encrypted_name_nonce=encrypted_name_nonce, - original_size=file_path.stat().st_size, - chunk_size=effective_chunk_size, - encryption=encryption_info, - strategy=dist_strategy, - ) - - print_info(f"Processing {num_chunks} chunks...") + # Save initial manifest (PENDING state) + manifest = builder.build() + store.save_manifest(manifest) + + if not resume: + print_info(f"Processing {num_chunks} chunks...") JobConfig(max_workers=min(4, num_chunks)) @@ -141,6 +295,13 @@ async def _put_async( progress = create_progress() task = progress.add_task("Uploading", total=num_chunks) + # Track which chunks are already uploaded (for resume) + uploaded_indices = set() + if resume and existing_manifest: + for chunk in existing_manifest.chunks: + if chunk.status == ChunkStatus.UPLOADED: + uploaded_indices.add(chunk.index) + async def upload_chunk(chunk_data: bytes, chunk_index: int, provider_name: str) -> tuple[bytes, dict[str, str]]: """Encrypt and upload a single chunk.""" async with semaphores[provider_name]: @@ -155,35 +316,69 @@ async def upload_chunk(chunk_data: bytes, chunk_index: int, provider_name: str) is_last=False, ) remote_ref = await provider_instances[provider_name].upload_chunk(encrypted_chunk, remote_path) - return encrypted.nonce, remote_ref.metadata + # Return nonce and a dict with remote_id and metadata + return encrypted.nonce, {"remote_id": remote_ref.remote_id, "metadata": remote_ref.metadata} # type: ignore[dict-item] progress = create_progress() task = progress.add_task("Uploading", total=num_chunks) - for chunk in chunker.chunk_file(file_path): - checksum = compute_checksum(chunk.data) + # Update progress for already uploaded chunks + for _ in uploaded_indices: + progress.advance(task) + + for c in chunker.chunk_file(file_path): + if c.index in uploaded_indices: + print_info(f"Chunk {c.index} already uploaded, skipping") + progress.advance(task) + continue + + checksum = compute_checksum(c.data) if dist_strategy == DistributionStrategy.SINGLE: target = provider_names[0] elif dist_strategy == DistributionStrategy.SPLIT: - target = provider_names[chunk.index % len(provider_names)] + target = provider_names[c.index % len(provider_names)] else: target = provider_names[0] - nonce, metadata = await upload_chunk(chunk.data, chunk.index, target) + # Mark chunk as uploading + builder.chunks[c.index] = ChunkInfo( + index=c.index, + size=c.size, + encrypted_size=0, + checksum=checksum, + provider=target, + remote_id="", + nonce=b"", + metadata={}, + status=ChunkStatus.UPLOADING, + ) + store.save_manifest(builder.build()) + + nonce, upload_result = await upload_chunk(c.data, c.index, target) + remote_id: str = upload_result["remote_id"] + metadata: dict[str, str] = upload_result["metadata"] # type: ignore[assignment] - builder.add_chunk( - index=chunk.index, - size=chunk.size, + # Update chunk as uploaded + builder.chunks[c.index] = ChunkInfo( + index=c.index, + size=c.size, encrypted_size=len(metadata.get("size", "0")), checksum=checksum, provider=target, - remote_id=metadata.get("remote_id", ""), + remote_id=remote_id, nonce=nonce, metadata=metadata, + status=ChunkStatus.UPLOADED, + uploaded_at=time.time(), + error=None, ) + manifest = builder.build() + store.save_manifest(manifest) + progress.advance(task) + # Final manifest build and save manifest = builder.build() store.save_manifest(manifest) diff --git a/src/stash/core/manifest.py b/src/stash/core/manifest.py index 78fa92e..3fba401 100644 --- a/src/stash/core/manifest.py +++ b/src/stash/core/manifest.py @@ -18,6 +18,23 @@ class DistributionStrategy(Enum): REPLICATED = "replicated" +class ChunkStatus(Enum): + """Upload status of a chunk.""" + PENDING = "pending" + UPLOADING = "uploading" + UPLOADED = "uploaded" + FAILED = "failed" + + +class UploadStatus(Enum): + """Overall upload status of a file.""" + NOT_STARTED = "not_started" + IN_PROGRESS = "in_progress" + COMPLETED = "completed" + FAILED = "failed" + PAUSED = "paused" + + @dataclass(frozen=True, slots=True) class ChunkInfo: """Information about a single chunk.""" @@ -29,6 +46,9 @@ class ChunkInfo: remote_id: str nonce: bytes metadata: dict[str, str] = field(default_factory=dict) + status: ChunkStatus = ChunkStatus.PENDING + uploaded_at: float | None = None + error: str | None = None @dataclass(frozen=True, slots=True) @@ -54,12 +74,17 @@ class FileManifest: original_size: int chunk_size: int chunk_count: int - encryption: EncryptionInfo - chunks: tuple[ChunkInfo, ...] + encryption: "EncryptionInfo" + chunks: tuple["ChunkInfo", ...] strategy: DistributionStrategy created_at: float modified_at: float version: int = 1 + upload_status: UploadStatus = UploadStatus.NOT_STARTED + total_chunks: int = 0 + uploaded_chunks: int = 0 + started_at: float | None = None + completed_at: float | None = None def to_json(self) -> str: """Serialize manifest to JSON.""" @@ -68,8 +93,10 @@ def to_json(self) -> str: data["encryption"] = asdict(self.encryption) data["strategy"] = self.strategy.value data["encrypted_name_nonce"] = self.encrypted_name_nonce.hex() + data["upload_status"] = self.upload_status.value for chunk in data["chunks"]: chunk["nonce"] = chunk["nonce"].hex() + chunk["status"] = chunk["status"].value data["encryption"]["file_key_salt"] = data["encryption"]["file_key_salt"].hex() if data["encryption"]["file_key_wrapped"]: data["encryption"]["file_key_wrapped"] = data["encryption"]["file_key_wrapped"].hex() @@ -80,6 +107,7 @@ def from_json(cls, json_str: str) -> "FileManifest": """Deserialize manifest from JSON.""" data = json.loads(json_str) data["strategy"] = DistributionStrategy(data["strategy"]) + data["upload_status"] = UploadStatus(data.get("upload_status", "not_started")) enc_data = data["encryption"] enc_data["file_key_salt"] = bytes.fromhex(enc_data["file_key_salt"]) if enc_data["file_key_wrapped"]: @@ -89,11 +117,12 @@ def from_json(cls, json_str: str) -> "FileManifest": chunks = [] for c in data["chunks"]: c["nonce"] = bytes.fromhex(c["nonce"]) + c["status"] = ChunkStatus(c.get("status", "pending")) chunks.append(ChunkInfo(**c)) data["chunks"] = tuple(chunks) return cls(**data) - def get_chunk(self, index: int) -> ChunkInfo: + def get_chunk(self, index: int) -> "ChunkInfo": """Get chunk info by index.""" for chunk in self.chunks: if chunk.index == index: @@ -130,6 +159,9 @@ def add_chunk( remote_id: str, nonce: bytes, metadata: dict[str, str] | None = None, + status: ChunkStatus = ChunkStatus.PENDING, + uploaded_at: float | None = None, + error: str | None = None, ) -> None: """Add a chunk to the manifest.""" self.chunks.append(ChunkInfo( @@ -141,12 +173,26 @@ def add_chunk( remote_id=remote_id, nonce=nonce, metadata=metadata or {}, + status=status, + uploaded_at=uploaded_at, + error=error, )) - def build(self) -> FileManifest: + def build(self) -> "FileManifest": """Build the final manifest.""" if len(self.chunks) == 0: raise ManifestError("Cannot build manifest with no chunks") + + total = len(self.chunks) + uploaded = sum(1 for c in self.chunks if c.status == ChunkStatus.UPLOADED) + status = UploadStatus.NOT_STARTED + if uploaded == len(self.chunks) and total > 0: + status = UploadStatus.COMPLETED + elif uploaded > 0: + status = UploadStatus.IN_PROGRESS + elif any(c.status == ChunkStatus.FAILED for c in self.chunks): + status = UploadStatus.FAILED + return FileManifest( file_id=self.file_id, original_name=self.original_name, @@ -160,6 +206,11 @@ def build(self) -> FileManifest: strategy=self.strategy, created_at=self.created_at, modified_at=time.time(), + upload_status=status, + total_chunks=total, + uploaded_chunks=uploaded, + started_at=self.chunks[0].uploaded_at if self.chunks else None, + completed_at=time.time() if status == UploadStatus.COMPLETED else None, ) diff --git a/tests/integration/test_resumable_uploads.py b/tests/integration/test_resumable_uploads.py new file mode 100644 index 0000000..bd83d5a --- /dev/null +++ b/tests/integration/test_resumable_uploads.py @@ -0,0 +1,566 @@ +"""Integration tests for resumable uploads.""" + +import pytest +import tempfile +import time +import sys +from io import StringIO +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock, patch + +from stash.core.chunking import ChunkConfig, Chunker +from stash.core.crypto import CryptoEngine +from stash.core.keymanager import KeyManager +from stash.core.manifest import ( + ChunkStatus, + ChunkInfo, + DistributionStrategy, + EncryptionInfo, + FileManifest, + ManifestBuilder, + UploadStatus, + compute_checksum, + generate_file_id, +) +from stash.core.metadata import MetadataStore +from stash.core.storage import RemoteRef + + +class MockProvider: + """Mock provider for testing.""" + + def __init__(self, name: str): + self.name = name + self.chunks = {} + self.closed = False + self.max_chunk_size = 10 * 1024 * 1024 + self.max_concurrent_uploads = 3 + self.config = type('Config', (), { + 'settings': {'max_concurrent': '3'}, + 'type': name, + 'credentials': {}, + 'settings': {'max_concurrent': '3'}, + })() + + async def initialize(self, config): + pass + + async def upload_chunk(self, chunk, remote_path): + self.chunks[remote_path] = chunk.data + return RemoteRef( + provider=self.name, + remote_id=remote_path, + metadata={'size': str(len(chunk.data))} + ) + + async def download_chunk(self, remote_ref): + return self.chunks.get(remote_ref.remote_id, b"") + + async def delete_chunk(self, remote_ref): + self.chunks.pop(remote_ref.remote_id, None) + + async def list_chunks(self, prefix): + return [] + + def get_limits(self): + from stash.core.storage import ProviderLimits + return ProviderLimits( + max_file_size=100 * 1024 * 1024, + max_chunk_size=10 * 1024 * 1024, + max_concurrent_uploads=3, + rate_limit_requests=30, + rate_limit_window=1, + ) + + async def close(self): + self.closed = True + + +@pytest.fixture +def temp_repo(): + """Create a temporary repository for testing.""" + with tempfile.TemporaryDirectory() as tmpdir: + repo_path = Path(tmpdir) + store = MetadataStore(repo_path) + + # Initialize repository with a dummy RMK + keymanager = KeyManager(repo_path) + identity = keymanager.initialize_repository() + + # Configure providers in the metadata store + from stash.core.storage import ProviderConfig + store.set_provider_config("telegram", ProviderConfig( + name="telegram", + type="telegram", + credentials={"token": "test_token", "chat_id": "-1001234567890"}, + settings={"max_concurrent": "3"}, + )) + store.set_provider_config("discord", ProviderConfig( + name="discord", + type="discord", + credentials={"token": "test_token", "channel_id": "123456789012345678"}, + settings={"max_concurrent": "3"}, + )) + + yield repo_path, store, keymanager, identity + + +@pytest.fixture +def test_file(temp_repo): + """Create a test file.""" + repo_path, _, _, _ = temp_repo + file_path = repo_path / "test_file.txt" + content = b"x" * (25 * 1024 * 1024) # 25MB file - 3 chunks at 10MB + file_path.write_bytes(content) + return file_path + + +@pytest.fixture +def mock_providers(): + """Create mock providers.""" + return { + "telegram": MockProvider("telegram"), + "discord": MockProvider("discord"), + } + + +class TestResumableUploads: + """Integration tests for resumable uploads.""" + + @pytest.mark.asyncio + async def test_new_upload_creates_partial_manifest(self, temp_repo, test_file, mock_providers): + """Test that a new upload creates a partial manifest with PENDING chunks.""" + repo_path, store, keymanager, identity = temp_repo + + # Import here to avoid circular imports + from stash.cli.commands.put import _put_async + + with patch('stash.cli.commands.put.ProviderRegistry') as mock_registry: + mock_registry.create = AsyncMock(side_effect=lambda name, config: mock_providers[name]) + + await _put_async( + file_path=test_file, + repo_path=repo_path, + provider_name="telegram", + chunk_size=None, + strategy="single", + do_confirm=False, + resume=False, + file_id=None, + ) + + # Verify manifest was created + files = store.list_files() + assert len(files) == 1 + + manifest = store.load_manifest(files[0]) + assert manifest.upload_status == UploadStatus.COMPLETED + assert manifest.total_chunks == 3 + assert manifest.uploaded_chunks == 3 + + # All chunks should be uploaded + for chunk in manifest.chunks: + assert chunk.status == ChunkStatus.UPLOADED + assert chunk.remote_id != "" + + @pytest.mark.asyncio + async def test_resume_upload_after_interruption(self, temp_repo, test_file, mock_providers): + """Test resuming an interrupted upload.""" + repo_path, store, keymanager, identity = temp_repo + + from stash.cli.commands.put import _put_async + + with patch('stash.cli.commands.put.ProviderRegistry') as mock_registry: + mock_registry.create = AsyncMock(side_effect=lambda name, config: mock_providers[name]) + + # First, start an upload and interrupt it after 1 chunk + await _put_async( + file_path=test_file, + repo_path=repo_path, + provider_name="telegram", + chunk_size=None, + strategy="single", + do_confirm=False, + resume=False, + file_id=None, + ) + + # Verify first upload completed + files = store.list_files() + manifest = store.load_manifest(files[0]) + assert manifest.upload_status == UploadStatus.COMPLETED + + # Now simulate interruption: manually set one chunk back to pending + # and remove it from the provider + file_id = files[0] + manifest = store.load_manifest(file_id) + + # Simulate interruption: mark last chunk as pending, remove from provider + with patch('stash.cli.commands.put.ProviderRegistry') as mock_registry: + mock_registry.create = AsyncMock(side_effect=lambda name, config: mock_providers[name]) + + # Remove last chunk from provider + last_chunk = manifest.chunks[-1] + if last_chunk.provider in mock_providers: + mock_providers[last_chunk.provider].chunks.pop(last_chunk.remote_id, None) + + # Update chunk status to pending + chunks = list(manifest.chunks) + old_chunk = chunks[-1] + chunks[-1] = ChunkInfo( + index=old_chunk.index, + size=old_chunk.size, + encrypted_size=old_chunk.encrypted_size, + checksum=old_chunk.checksum, + provider=old_chunk.provider, + remote_id="", + nonce=b"", + metadata={}, + status=ChunkStatus.PENDING, + uploaded_at=None, + error=None, + ) + # Rebuild manifest with updated chunks + from stash.core.manifest import FileManifest, EncryptionInfo, DistributionStrategy + new_manifest = FileManifest( + file_id=manifest.file_id, + original_name=manifest.original_name, + encrypted_name=manifest.encrypted_name, + encrypted_name_nonce=manifest.encrypted_name_nonce, + original_size=manifest.original_size, + chunk_size=manifest.chunk_size, + chunk_count=manifest.chunk_count, + encryption=manifest.encryption, + chunks=tuple(chunks), + strategy=manifest.strategy, + created_at=manifest.created_at, + modified_at=time.time(), + upload_status=UploadStatus.IN_PROGRESS, + total_chunks=manifest.total_chunks, + uploaded_chunks=manifest.uploaded_chunks - 1, + started_at=manifest.started_at, + completed_at=None, + ) + store.save_manifest(new_manifest) + + # Now resume the upload + await _put_async( + file_path=test_file, + repo_path=repo_path, + provider_name="telegram", + chunk_size=None, + strategy="single", + do_confirm=False, + resume=True, + file_id=None, + ) + + # Verify upload completed + manifest = store.load_manifest(file_id) + assert manifest.upload_status == UploadStatus.COMPLETED + assert manifest.uploaded_chunks == 3 + + for chunk in manifest.chunks: + assert chunk.status == ChunkStatus.UPLOADED + assert chunk.remote_id != "" + + @pytest.mark.asyncio + async def test_resume_with_explicit_file_id(self, temp_repo, test_file, mock_providers): + """Test resuming with explicit --file-id.""" + repo_path, store, keymanager, identity = temp_repo + + from stash.cli.commands.put import _put_async + + # Create an incomplete upload first + with patch('stash.cli.commands.put.ProviderRegistry') as mock_registry: + mock_registry.create = AsyncMock(side_effect=lambda name, config: mock_providers[name]) + + await _put_async( + file_path=test_file, + repo_path=repo_path, + provider_name="telegram", + chunk_size=None, + strategy="single", + do_confirm=False, + resume=False, + file_id=None, + ) + + files = store.list_files() + file_id = files[0] + + # Corrupt the upload - mark one chunk as pending + manifest = store.load_manifest(file_id) + chunks = list(manifest.chunks) + chunks[1] = ChunkInfo( + index=chunks[1].index, + size=chunks[1].size, + encrypted_size=chunks[1].encrypted_size, + checksum=chunks[1].checksum, + provider=chunks[1].provider, + remote_id="", + nonce=b"", + metadata={}, + status=ChunkStatus.PENDING, + uploaded_at=None, + error=None, + ) + + new_manifest = FileManifest( + file_id=manifest.file_id, + original_name=manifest.original_name, + encrypted_name=manifest.encrypted_name, + encrypted_name_nonce=manifest.encrypted_name_nonce, + original_size=manifest.original_size, + chunk_size=manifest.chunk_size, + chunk_count=manifest.chunk_count, + encryption=manifest.encryption, + chunks=tuple(chunks), + strategy=manifest.strategy, + created_at=manifest.created_at, + modified_at=time.time(), + upload_status=UploadStatus.IN_PROGRESS, + total_chunks=3, + uploaded_chunks=2, + started_at=manifest.started_at, + completed_at=None, + ) + store.save_manifest(new_manifest) + + # Remove chunk from provider + mock_providers["telegram"].chunks.pop(chunks[1].remote_id, None) + + # Resume with explicit file-id + with patch('stash.cli.commands.put.ProviderRegistry') as mock_registry: + mock_registry.create = AsyncMock(side_effect=lambda name, config: mock_providers[name]) + + await _put_async( + file_path=test_file, + repo_path=repo_path, + provider_name="telegram", + chunk_size=None, + strategy="single", + do_confirm=False, + resume=True, + file_id=file_id, + ) + + # Verify completed + manifest = store.load_manifest(file_id) + assert manifest.upload_status == UploadStatus.COMPLETED + assert manifest.uploaded_chunks == 3 + + @pytest.mark.asyncio + async def test_resume_fails_when_file_changed(self, temp_repo, test_file, mock_providers): + """Test that resume fails when file content has changed.""" + repo_path, store, keymanager, identity = temp_repo + + from stash.cli.commands.put import _put_async + + # Create incomplete upload + with patch('stash.cli.commands.put.ProviderRegistry') as mock_registry: + mock_registry.create = AsyncMock(side_effect=lambda name, config: mock_providers[name]) + + await _put_async( + file_path=test_file, + repo_path=repo_path, + provider_name="telegram", + chunk_size=None, + strategy="single", + do_confirm=False, + resume=False, + file_id=None, + ) + + files = store.list_files() + manifest = store.load_manifest(files[0]) + + # Corrupt the manifest to simulate incomplete upload + chunks = list(manifest.chunks) + chunks[0] = ChunkInfo( + index=chunks[0].index, + size=chunks[0].size, + encrypted_size=chunks[0].encrypted_size, + checksum=chunks[0].checksum, + provider=chunks[0].provider, + remote_id="", + nonce=b"", + metadata={}, + status=ChunkStatus.PENDING, + uploaded_at=None, + error=None, + ) + + from stash.core.manifest import FileManifest, UploadStatus + new_manifest = FileManifest( + file_id=manifest.file_id, + original_name=manifest.original_name, + encrypted_name=manifest.encrypted_name, + encrypted_name_nonce=manifest.encrypted_name_nonce, + original_size=manifest.original_size, + chunk_size=manifest.chunk_size, + chunk_count=manifest.chunk_count, + encryption=manifest.encryption, + chunks=tuple(chunks), + strategy=manifest.strategy, + created_at=manifest.created_at, + modified_at=time.time(), + upload_status=UploadStatus.IN_PROGRESS, + total_chunks=manifest.total_chunks, + uploaded_chunks=2, + started_at=manifest.started_at, + completed_at=None, + ) + store.save_manifest(new_manifest) + + # Modify the file content but keep the same size + new_content = b"y" * (25 * 1024 * 1024) # 25MB of different content + test_file.write_bytes(new_content) + + # Try to resume - should fail + with patch('stash.cli.commands.put.ProviderRegistry') as mock_registry: + mock_registry.create = AsyncMock(side_effect=lambda name, config: mock_providers[name]) + + import sys + from io import StringIO + old_stdout = sys.stdout + sys.stdout = StringIO() + try: + await _put_async( + file_path=test_file, + repo_path=repo_path, + provider_name="telegram", + chunk_size=None, + strategy="single", + do_confirm=False, + resume=True, + file_id=None, + ) + except SystemExit: + pass + finally: + captured_stdout = sys.stdout.getvalue() + sys.stdout = old_stdout + + # Should fail with file content mismatch error + assert "does not match" in captured_stdout + + @pytest.mark.asyncio + async def test_multiple_resume_cycles(self, temp_repo, test_file, mock_providers): + """Test multiple resume cycles work correctly.""" + repo_path, store, keymanager, identity = temp_repo + + from stash.cli.commands.put import _put_async + + # Initial upload + with patch('stash.cli.commands.put.ProviderRegistry') as mock_registry: + mock_registry.create = AsyncMock(side_effect=lambda name, config: mock_providers[name]) + + await _put_async( + file_path=test_file, + repo_path=repo_path, + provider_name="telegram", + chunk_size=None, + strategy="single", + do_confirm=False, + resume=False, + file_id=None, + ) + + files = store.list_files() + file_id = files[0] + + # Simulate 3 interruption/resume cycles + for cycle in range(3): + manifest = store.load_manifest(file_id) + + # Mark last completed chunk as pending + chunks = list(manifest.chunks) + last_uploaded = max((i for i, c in enumerate(chunks) if c.status == ChunkStatus.UPLOADED), default=0) + if last_uploaded >= 0: + old_chunk = chunks[last_uploaded] + chunks[last_uploaded] = ChunkInfo( + index=old_chunk.index, + size=old_chunk.size, + encrypted_size=old_chunk.encrypted_size, + checksum=old_chunk.checksum, + provider=old_chunk.provider, + remote_id="", + nonce=b"", + metadata={}, + status=ChunkStatus.PENDING, + uploaded_at=None, + error=None, + ) + + # Remove from provider + mock_providers["telegram"].chunks.pop(manifest.chunks[last_uploaded].remote_id, None) + + # Rebuild manifest + from stash.core.manifest import FileManifest + new_manifest = FileManifest( + file_id=manifest.file_id, + original_name=manifest.original_name, + encrypted_name=manifest.encrypted_name, + encrypted_name_nonce=manifest.encrypted_name_nonce, + original_size=manifest.original_size, + chunk_size=manifest.chunk_size, + chunk_count=manifest.chunk_count, + encryption=manifest.encryption, + chunks=tuple(chunks), + strategy=manifest.strategy, + created_at=manifest.created_at, + modified_at=time.time(), + upload_status=UploadStatus.IN_PROGRESS, + total_chunks=manifest.total_chunks, + uploaded_chunks=manifest.uploaded_chunks - 1, + started_at=manifest.started_at, + completed_at=None, + ) + store.save_manifest(new_manifest) + + # Resume + with patch('stash.cli.commands.put.ProviderRegistry') as mock_registry: + mock_registry.create = AsyncMock(side_effect=lambda name, config: mock_providers[name]) + + await _put_async( + file_path=test_file, + repo_path=repo_path, + provider_name="telegram", + chunk_size=None, + strategy="single", + do_confirm=False, + resume=True, + file_id=file_id, + ) + + # Verify completed + manifest = store.load_manifest(file_id) + assert manifest.upload_status == UploadStatus.COMPLETED + assert manifest.uploaded_chunks == 3 + + @pytest.mark.asyncio + async def test_corrupted_manifest_handling(self, temp_repo, test_file, mock_providers): + """Test handling of corrupted manifest.""" + repo_path, store, keymanager, identity = temp_repo + + # Create a corrupted manifest file + manifest_dir = repo_path / ".stash" / "metadata" / "files" + manifest_dir.mkdir(parents=True, exist_ok=True) + + # Write invalid JSON + corrupted_file = manifest_dir / "corrupted.json" + corrupted_file.write_text("{ invalid json") + + # Should not crash when listing files + files = store.list_files() + assert isinstance(files, list) + + # Try to load corrupted manifest - should raise MetadataError + from stash.core.exceptions import MetadataError + with pytest.raises(MetadataError): + store.load_manifest("corrupted") + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) \ No newline at end of file