import os import shutil import hashlib from datetime import datetime, timezone from pathlib import Path from typing import List, Tuple, Dict, Any, Optional import aiofiles from app.core.config import settings from app.storage.base import BaseStorageProvider class LocalStorageProvider(BaseStorageProvider): """Local filesystem storage provider designed for Proxmox VE dedicated mount volumes.""" def __init__(self, root_dir: Optional[str] = None, temp_dir: Optional[str] = None): self.root_dir = Path(root_dir or settings.STORAGE_ROOT).resolve() self.temp_dir = Path(temp_dir or settings.STORAGE_TEMP_ROOT).resolve() self.root_dir.mkdir(parents=True, exist_ok=True) self.temp_dir.mkdir(parents=True, exist_ok=True) def _get_session_temp_dir(self, session_code: str) -> Path: return self.temp_dir / session_code async def init_session_storage(self, session_code: str) -> None: session_path = self._get_session_temp_dir(session_code) session_path.mkdir(parents=True, exist_ok=True) async def save_chunk( self, session_code: str, chunk_index: int, chunk_data: bytes, expected_sha256: Optional[str] = None ) -> bool: # Validate individual chunk hash if provided if expected_sha256: actual_chunk_hash = hashlib.sha256(chunk_data).hexdigest() if actual_chunk_hash.lower() != expected_sha256.lower(): raise ValueError( f"Chunk {chunk_index} checksum mismatch: expected {expected_sha256}, got {actual_chunk_hash}" ) session_path = self._get_session_temp_dir(session_code) session_path.mkdir(parents=True, exist_ok=True) chunk_file = session_path / f"{chunk_index:08d}.chunk" async with aiofiles.open(chunk_file, "wb") as f: await f.write(chunk_data) return True async def get_received_chunks(self, session_code: str) -> List[int]: session_path = self._get_session_temp_dir(session_code) if not session_path.exists(): return [] chunks = [] for file in session_path.glob("*.chunk"): try: index = int(file.stem) chunks.append(index) except ValueError: continue chunks.sort() return chunks async def assemble_file( self, session_code: str, client_code: str, job_code: str, filename: str, total_chunks: int, expected_sha256: str ) -> Tuple[str, str, int]: session_path = self._get_session_temp_dir(session_code) if not session_path.exists(): raise FileNotFoundError(f"Upload session temporary directory {session_code} does not exist") # Verify all chunks are present received_chunks = set(await self.get_received_chunks(session_code)) missing_chunks = [i for i in range(total_chunks) if i not in received_chunks] if missing_chunks: raise ValueError(f"Cannot assemble file. Missing {len(missing_chunks)} chunks: {missing_chunks[:10]}...") # Prepare client isolated destination directory dest_dir = self.root_dir / "clients" / client_code / (job_code or "DEFAULT") dest_dir.mkdir(parents=True, exist_ok=True) timestamp_str = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S") safe_filename = Path(filename).name target_filename = f"{timestamp_str}_{safe_filename}" target_path = dest_dir / target_filename relative_path = str(target_path.relative_to(self.root_dir)).replace("\\", "/") hasher = hashlib.sha256() total_bytes = 0 # Stream and concatenate all chunks in sequential order async with aiofiles.open(target_path, "wb") as out_file: for idx in range(total_chunks): chunk_file = session_path / f"{idx:08d}.chunk" if not chunk_file.exists(): # Clean up target on failure if target_path.exists(): target_path.unlink() raise FileNotFoundError(f"Missing chunk file {chunk_file}") async with aiofiles.open(chunk_file, "rb") as in_chunk: while True: buffer = await in_chunk.read(1024 * 1024) # 1MB buffer if not buffer: break hasher.update(buffer) total_bytes += len(buffer) await out_file.write(buffer) final_sha256 = hasher.hexdigest() # Strict integrity check against client's pre-calculated full SHA-256 if final_sha256.lower() != expected_sha256.lower(): if target_path.exists(): target_path.unlink() raise ValueError( f"Full file integrity check failed! Expected SHA-256: {expected_sha256}, Actual: {final_sha256}" ) # Cleanup temporary chunks upon confirmed assembly & integrity verification await self.delete_session_temp(session_code) return relative_path, final_sha256, total_bytes async def delete_session_temp(self, session_code: str) -> None: session_path = self._get_session_temp_dir(session_code) if session_path.exists(): shutil.rmtree(session_path, ignore_errors=True) async def delete_backup_file(self, relative_path: str) -> bool: full_path = (self.root_dir / relative_path).resolve() # Security guard: prevent path traversal outside root_dir if not str(full_path).startswith(str(self.root_dir)): raise PermissionError("Attempted path traversal outside storage root") if full_path.exists() and full_path.is_file(): full_path.unlink() return True return False async def get_file_path(self, relative_path: str) -> str: full_path = (self.root_dir / relative_path).resolve() if not str(full_path).startswith(str(self.root_dir)): raise PermissionError("Attempted path traversal outside storage root") if not full_path.exists(): raise FileNotFoundError(f"Backup file {relative_path} not found") return str(full_path) async def get_storage_stats(self) -> Dict[str, Any]: total, used, free = shutil.disk_usage(self.root_dir) usage_pct = round((used / total) * 100, 2) if total > 0 else 0 return { "total_bytes": total, "used_bytes": used, "free_bytes": free, "usage_percent": usage_pct, "storage_root": str(self.root_dir) } # Global singleton storage provider storage_provider = LocalStorageProvider()