S.H.O.N.A.R._Desktop_Companion/backend/shonar/storage/__init__.py

"""Storage abstraction.

Backends:
- ``LocalStorage``: filesystem under ``SHONAR_STORAGE_PATH`` (dev + default).
- ``S3Storage``: any S3-compatible object store (extra: ``pip install
  shonar-backend[s3]``).

Storage keys are server-internal and validated against path traversal:
they are UUID-based by construction and never derived from user input.
At-rest encryption is NOT implemented; see docs/security.md for the
documented optional design.
"""

from __future__ import annotations

import asyncio
import shutil
from pathlib import Path, PurePosixPath
from typing import Protocol

from shonar.core.config import get_settings

_KEY_CHARSET = set("abcdefghijklmnopqrstuvwxyz0123456789-./")


class StorageError(Exception):
    pass


def validate_storage_key(key: str) -> str:
    """Reject anything that could escape the storage root."""
    if not key or len(key) > 500:
        raise StorageError("invalid storage key")
    pure = PurePosixPath(key)
    if pure.is_absolute() or ".." in pure.parts:
        raise StorageError("invalid storage key")
    if not set(key) <= _KEY_CHARSET:
        raise StorageError("invalid storage key")
    return key


class StorageBackend(Protocol):
    async def put(self, key: str, data: bytes) -> int: ...
    async def put_file(self, key: str, src_path: Path) -> int: ...
    async def get(self, key: str) -> bytes: ...
    async def open_path(self, key: str) -> Path | None:
        """Local file path if the backend can provide one, else None."""
        ...

    async def delete(self, key: str) -> None: ...
    async def exists(self, key: str) -> bool: ...


class LocalStorage:
    def __init__(self, root: str | Path):
        self.root = Path(root).resolve()
        self.root.mkdir(parents=True, exist_ok=True)

    def _path(self, key: str) -> Path:
        validate_storage_key(key)
        path = (self.root / key).resolve()
        # Defense in depth: resolved path must stay under root.
        if not path.is_relative_to(self.root):
            raise StorageError("storage key escapes storage root")
        return path

    async def put(self, key: str, data: bytes) -> int:
        path = self._path(key)
        path.parent.mkdir(parents=True, exist_ok=True)
        await asyncio.to_thread(path.write_bytes, data)
        return len(data)

    async def put_file(self, key: str, src_path: Path) -> int:
        path = self._path(key)
        path.parent.mkdir(parents=True, exist_ok=True)
        await asyncio.to_thread(shutil.copyfile, src_path, path)
        return path.stat().st_size

    async def get(self, key: str) -> bytes:
        path = self._path(key)
        if not path.exists():
            raise StorageError("object not found")
        return await asyncio.to_thread(path.read_bytes)

    async def open_path(self, key: str) -> Path | None:
        path = self._path(key)
        return path if path.exists() else None

    async def delete(self, key: str) -> None:
        path = self._path(key)
        if path.exists():
            await asyncio.to_thread(path.unlink)

    async def exists(self, key: str) -> bool:
        return self._path(key).exists()


class S3Storage:  # pragma: no cover - requires boto3 + a real/mini endpoint
    def __init__(
        self, endpoint_url: str, bucket: str, region: str, access_key: str, secret_key: str
    ):
        import boto3  # optional extra

        self.bucket = bucket
        self.s3 = boto3.client(
            "s3",
            endpoint_url=endpoint_url or None,
            region_name=region,
            aws_access_key_id=access_key or None,
            aws_secret_access_key=secret_key or None,
        )

    async def put(self, key: str, data: bytes) -> int:
        validate_storage_key(key)
        await asyncio.to_thread(self.s3.put_object, Bucket=self.bucket, Key=key, Body=data)
        return len(data)

    async def put_file(self, key: str, src_path: Path) -> int:
        validate_storage_key(key)
        await asyncio.to_thread(self.s3.upload_file, str(src_path), self.bucket, key)
        return src_path.stat().st_size

    async def get(self, key: str) -> bytes:
        validate_storage_key(key)

        def _get() -> bytes:
            obj = self.s3.get_object(Bucket=self.bucket, Key=key)
            return obj["Body"].read()

        return await asyncio.to_thread(_get)

    async def open_path(self, key: str) -> Path | None:
        return None  # callers must use get()/streaming

    async def delete(self, key: str) -> None:
        validate_storage_key(key)
        await asyncio.to_thread(self.s3.delete_object, Bucket=self.bucket, Key=key)

    async def exists(self, key: str) -> bool:
        validate_storage_key(key)

        def _head() -> bool:
            from botocore.exceptions import ClientError

            try:
                self.s3.head_object(Bucket=self.bucket, Key=key)
                return True
            except ClientError:
                return False

        return await asyncio.to_thread(_head)


_backend: StorageBackend | None = None


def get_storage() -> StorageBackend:
    global _backend
    if _backend is None:
        settings = get_settings()
        if settings.storage_backend == "local":
            _backend = LocalStorage(settings.storage_path)
        elif settings.storage_backend == "s3":
            _backend = S3Storage(
                settings.s3_endpoint_url,
                settings.s3_bucket,
                settings.s3_region,
                settings.s3_access_key_id,
                settings.s3_secret_access_key,
            )
        else:
            raise StorageError(f"unknown storage backend: {settings.storage_backend}")
    return _backend


def set_storage(backend: StorageBackend | None) -> None:
    """Test seam."""
    global _backend
    _backend = backend