169 lines
6.7 KiB
Python
169 lines
6.7 KiB
Python
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()
|