rbox/mywebdav/storage.py

import os
import aiofiles
from pathlib import Path
from typing import AsyncGenerator

from .settings import settings

class StorageManager:
    def __init__(self, base_path: str = settings.STORAGE_PATH):
        self.base_path = Path(base_path)
        self.base_path.mkdir(parents=True, exist_ok=True)

    async def _get_full_path(self, user_id: int, file_path: str) -> Path:
        # Ensure file_path is relative and safe
        relative_path = Path(file_path).relative_to('/') if str(file_path).startswith('/') else Path(file_path)
        full_path = self.base_path / str(user_id) / relative_path
        full_path.parent.mkdir(parents=True, exist_ok=True)
        return full_path

    async def save_file(self, user_id: int, file_path: str, file_content: bytes):
        full_path = await self._get_full_path(user_id, file_path)
        async with aiofiles.open(full_path, "wb") as f:
            await f.write(file_content)
        return str(full_path)

    async def get_file(self, user_id: int, file_path: str) -> AsyncGenerator:
        full_path = await self._get_full_path(user_id, file_path)
        if not full_path.exists():
            raise FileNotFoundError(f"File not found: {file_path}")
        
        async with aiofiles.open(full_path, "rb") as f:
            while chunk := await f.read(8192):
                yield chunk

    async def delete_file(self, user_id: int, file_path: str):
        full_path = await self._get_full_path(user_id, file_path)
        if full_path.exists():
            os.remove(full_path)

            parent_dir = full_path.parent
            while parent_dir != self.base_path and parent_dir.exists():
                try:
                    if not any(parent_dir.iterdir()):
                        parent_dir.rmdir()
                        parent_dir = parent_dir.parent
                    else:
                        break
                except OSError:
                    break

    async def file_exists(self, user_id: int, file_path: str) -> bool:
        full_path = await self._get_full_path(user_id, file_path)
        return full_path.exists()

storage_manager = StorageManager()