from __future__ import annotations import asyncio import json import logging import os import uuid from datetime import UTC, datetime from pathlib import Path logger = logging.getLogger(__name__) _META_SUFFIX = ".meta.json" def _now_iso() -> str: return datetime.now(UTC).isoformat() def _date_dir_name() -> str: return datetime.now(UTC).strftime("%Y-%m-%d") class MediaStore: def __init__( self, base_dir: Path, max_total_bytes: int, max_file_bytes: int, ttl_seconds: int, ): self.base_dir = base_dir self.max_total_bytes = max_total_bytes self.max_file_bytes = max_file_bytes self.ttl_seconds = ttl_seconds self._size_initialized = False self._running_total: int = 0 self._size_lock = asyncio.Lock() def _ensure_dir(self, dir_path: Path) -> None: dir_path.mkdir(parents=True, exist_ok=True) def _make_file_id(self, date_dir: str, uid: str, ext: str) -> str: return f"{date_dir}/{uid}{ext}" def _get_file_path(self, file_id: str) -> Path: return self.base_dir / file_id def _get_meta_path(self, file_id: str) -> Path: return self.base_dir / f"{file_id}{_META_SUFFIX}" def _dir_size(self, dir_path: Path) -> int: total = 0 try: for entry in dir_path.rglob("*"): if entry.is_file(): try: total += entry.stat().st_size except OSError: logger.debug("Failed to stat file during size calculation: %s", entry) except OSError: logger.debug("Failed to calculate directory size: %s", dir_path) return total async def _check_quota(self, incoming_bytes: int) -> bool: if incoming_bytes > self.max_file_bytes: logger.warning( "File size %d exceeds per-file limit %d", incoming_bytes, self.max_file_bytes, ) return False current_total = await self._get_running_total() if current_total + incoming_bytes > self.max_total_bytes: logger.warning( "Storage quota exceeded: current=%d, incoming=%d, max=%d", current_total, incoming_bytes, self.max_total_bytes, ) return False return True async def _get_running_total(self) -> int: async with self._size_lock: if not self._size_initialized: self._running_total = self._dir_size(self.base_dir) self._size_initialized = True return self._running_total def _infer_ext(self, content_type: str | None, filename: str | None) -> str: if filename: suffix = Path(filename).suffix if suffix: return suffix.lower() if content_type: ct = content_type.split(";")[0].strip().lower() mapping = { "image/jpeg": ".jpg", "image/png": ".png", "image/gif": ".gif", "image/webp": ".webp", "image/bmp": ".bmp", "image/svg+xml": ".svg", "audio/mpeg": ".mp3", "audio/wav": ".wav", "audio/ogg": ".ogg", "video/mp4": ".mp4", "application/pdf": ".pdf", } if ct in mapping: return mapping[ct] return ".bin" async def store( self, data: bytes, content_type: str | None = None, filename: str | None = None, source_url: str | None = None, ) -> str | None: if not await self._check_quota(len(data)): return None date_dir = _date_dir_name() uid = uuid.uuid4().hex ext = self._infer_ext(content_type, filename) file_id = self._make_file_id(date_dir, uid, ext) file_path = self._get_file_path(file_id) meta_path = self._get_meta_path(file_id) self._ensure_dir(file_path.parent) try: file_path.write_bytes(data) except OSError: logger.exception("Failed to write media file: %s", file_path) return None meta = { "original_filename": filename, "content_type": content_type, "source_url": source_url, "stored_at": _now_iso(), "ttl_seconds": self.ttl_seconds, "size": len(data), } try: meta_path.write_text(json.dumps(meta, ensure_ascii=False), encoding="utf-8") except OSError: logger.exception("Failed to write media metadata: %s", meta_path) try: file_path.unlink(missing_ok=True) except OSError: logger.debug("Failed to clean up orphaned media file: %s", file_path) return None logger.info("Stored media: file_id=%s, size=%d", file_id, len(data)) async with self._size_lock: self._running_total += len(data) return file_id def resolve_path(self, file_id: str) -> Path | None: file_path = self._get_file_path(file_id) if file_path.exists(): return file_path return None def read_meta(self, file_id: str) -> dict | None: meta_path = self._get_meta_path(file_id) try: return json.loads(meta_path.read_text(encoding="utf-8")) except (OSError, json.JSONDecodeError): return None async def delete(self, file_id: str) -> bool: file_path = self._get_file_path(file_id) meta_path = self._get_meta_path(file_id) deleted = False for p in (file_path, meta_path): if not p.exists(): continue try: p.unlink() deleted = True except OSError: logger.debug("Failed to delete file during media cleanup: %s", p) if deleted: logger.info("Deleted media: file_id=%s", file_id) return deleted def get_total_size(self) -> int: return self._dir_size(self.base_dir) async def cleanup_expired(self) -> int: deleted_count = 0 now = datetime.now(UTC) try: for entry in sorted(self.base_dir.rglob(f"*{_META_SUFFIX}")): try: meta = json.loads(entry.read_text(encoding="utf-8")) except (OSError, json.JSONDecodeError): continue stored_at_str = meta.get("stored_at") ttl_s = meta.get("ttl_seconds", self.ttl_seconds) if not stored_at_str: continue try: stored_at = datetime.fromisoformat(stored_at_str) except ValueError: continue if (now - stored_at).total_seconds() <= ttl_s: continue file_id = str(entry.relative_to(self.base_dir)).removesuffix(_META_SUFFIX) if await self.delete(file_id): deleted_count += 1 self._remove_empty_dirs() except OSError: logger.exception("Error during media cleanup scan") if deleted_count: logger.info("Media cleanup: deleted %d expired files", deleted_count) return deleted_count def _remove_empty_dirs(self) -> None: try: for dirpath, dirnames, filenames in os.walk(self.base_dir, topdown=False): if dirpath == str(self.base_dir): continue if not dirnames and not filenames: try: os.rmdir(dirpath) except OSError: logger.debug("Failed to remove empty directory: %s", dirpath) except OSError: logger.debug("Failed to walk directory for cleanup: %s", self.base_dir)