rbox/mywebdav/concurrency/webdav_locks.py

import asyncio
import time
import hashlib
import uuid
from typing import Dict, Optional
from dataclasses import dataclass, field, asdict
import json
import logging

logger = logging.getLogger(__name__)


@dataclass
class WebDAVLockInfo:
    token: str
    path: str
    path_hash: str
    owner: str
    user_id: int
    scope: str = "exclusive"
    depth: str = "0"
    timeout: int = 3600
    created_at: float = field(default_factory=time.time)

    @property
    def is_expired(self) -> bool:
        return time.time() - self.created_at > self.timeout

    @property
    def remaining_seconds(self) -> int:
        remaining = self.timeout - (time.time() - self.created_at)
        return max(0, int(remaining))

    def to_dict(self) -> dict:
        return asdict(self)

    @classmethod
    def from_dict(cls, data: dict) -> "WebDAVLockInfo":
        return cls(**data)


class PersistentWebDAVLocks:
    def __init__(self, db_manager=None):
        self.db_manager = db_manager
        self.locks: Dict[str, WebDAVLockInfo] = {}
        self._lock = asyncio.Lock()
        self._cleanup_task: Optional[asyncio.Task] = None
        self._running = False
        self._persistence_enabled = False

    async def start(self, db_manager=None):
        if db_manager:
            self.db_manager = db_manager
            self._persistence_enabled = True
            await self._load_locks_from_db()
        self._running = True
        self._cleanup_task = asyncio.create_task(self._background_cleanup())
        logger.info("PersistentWebDAVLocks started")

    async def stop(self):
        self._running = False
        if self._cleanup_task:
            self._cleanup_task.cancel()
            try:
                await self._cleanup_task
            except asyncio.CancelledError:
                pass
        if self._persistence_enabled:
            await self._save_all_locks()
        logger.info("PersistentWebDAVLocks stopped")

    def _hash_path(self, path: str) -> str:
        return hashlib.sha256(path.encode()).hexdigest()[:32]

    async def acquire_lock(
        self,
        path: str,
        owner: str,
        user_id: int,
        timeout: int = 3600,
        scope: str = "exclusive",
        depth: str = "0"
    ) -> Optional[str]:
        path = path.strip("/")
        path_hash = self._hash_path(path)

        async with self._lock:
            if path_hash in self.locks:
                existing = self.locks[path_hash]
                if not existing.is_expired:
                    if existing.user_id == user_id:
                        existing.created_at = time.time()
                        existing.timeout = timeout
                        await self._persist_lock(existing)
                        return existing.token
                    return None
                else:
                    del self.locks[path_hash]
                    await self._delete_lock_from_db(path_hash)

            token = f"opaquelocktoken:{uuid.uuid4()}"
            lock_info = WebDAVLockInfo(
                token=token,
                path=path,
                path_hash=path_hash,
                owner=owner,
                user_id=user_id,
                scope=scope,
                depth=depth,
                timeout=timeout
            )
            self.locks[path_hash] = lock_info
            await self._persist_lock(lock_info)
            logger.debug(f"WebDAV lock acquired: {path} by {owner}")
            return token

    async def refresh_lock(self, path: str, token: str, timeout: int = 3600) -> bool:
        path = path.strip("/")
        path_hash = self._hash_path(path)

        async with self._lock:
            if path_hash not in self.locks:
                return False
            lock_info = self.locks[path_hash]
            if lock_info.token != token:
                return False
            lock_info.created_at = time.time()
            lock_info.timeout = timeout
            await self._persist_lock(lock_info)
            return True

    async def release_lock(self, path: str, token: str) -> bool:
        path = path.strip("/")
        path_hash = self._hash_path(path)

        async with self._lock:
            if path_hash not in self.locks:
                return False
            lock_info = self.locks[path_hash]
            if lock_info.token != token:
                return False
            del self.locks[path_hash]
            await self._delete_lock_from_db(path_hash)
            logger.debug(f"WebDAV lock released: {path}")
            return True

    async def check_lock(self, path: str) -> Optional[WebDAVLockInfo]:
        path = path.strip("/")
        path_hash = self._hash_path(path)

        async with self._lock:
            if path_hash in self.locks:
                lock_info = self.locks[path_hash]
                if lock_info.is_expired:
                    del self.locks[path_hash]
                    await self._delete_lock_from_db(path_hash)
                    return None
                return lock_info
            return None

    async def get_lock_by_token(self, token: str) -> Optional[WebDAVLockInfo]:
        async with self._lock:
            for lock_info in self.locks.values():
                if lock_info.token == token and not lock_info.is_expired:
                    return lock_info
            return None

    async def is_locked(self, path: str, user_id: Optional[int] = None) -> bool:
        lock_info = await self.check_lock(path)
        if not lock_info:
            return False
        if user_id is not None and lock_info.user_id == user_id:
            return False
        return True

    async def force_unlock(self, path: str, user_id: int) -> bool:
        path = path.strip("/")
        path_hash = self._hash_path(path)

        async with self._lock:
            if path_hash not in self.locks:
                return False
            lock_info = self.locks[path_hash]
            if lock_info.user_id != user_id:
                return False
            del self.locks[path_hash]
            await self._delete_lock_from_db(path_hash)
            return True

    async def get_user_locks(self, user_id: int) -> list:
        async with self._lock:
            return [
                lock_info for lock_info in self.locks.values()
                if lock_info.user_id == user_id and not lock_info.is_expired
            ]

    async def _persist_lock(self, lock_info: WebDAVLockInfo):
        if not self._persistence_enabled or not self.db_manager:
            return
        try:
            async with self.db_manager.get_master_connection() as conn:
                await conn.execute("""
                    INSERT OR REPLACE INTO webdav_locks
                    (path_hash, path, token, owner, user_id, scope, timeout, created_at)
                    VALUES (?, ?, ?, ?, ?, ?, ?, ?)
                """, (
                    lock_info.path_hash,
                    lock_info.path,
                    lock_info.token,
                    lock_info.owner,
                    lock_info.user_id,
                    lock_info.scope,
                    lock_info.timeout,
                    lock_info.created_at
                ))
                await conn.commit()
        except Exception as e:
            logger.error(f"Failed to persist WebDAV lock: {e}")

    async def _delete_lock_from_db(self, path_hash: str):
        if not self._persistence_enabled or not self.db_manager:
            return
        try:
            async with self.db_manager.get_master_connection() as conn:
                await conn.execute(
                    "DELETE FROM webdav_locks WHERE path_hash = ?",
                    (path_hash,)
                )
                await conn.commit()
        except Exception as e:
            logger.error(f"Failed to delete WebDAV lock from db: {e}")

    async def _load_locks_from_db(self):
        if not self.db_manager:
            return
        try:
            async with self.db_manager.get_master_connection() as conn:
                cursor = await conn.execute("SELECT * FROM webdav_locks")
                rows = await cursor.fetchall()
                for row in rows:
                    lock_info = WebDAVLockInfo(
                        token=row[3],
                        path=row[2],
                        path_hash=row[1],
                        owner=row[4],
                        user_id=row[5],
                        scope=row[6],
                        timeout=row[7],
                        created_at=row[8]
                    )
                    if not lock_info.is_expired:
                        self.locks[lock_info.path_hash] = lock_info
                    else:
                        await self._delete_lock_from_db(lock_info.path_hash)
                logger.info(f"Loaded {len(self.locks)} WebDAV locks from database")
        except Exception as e:
            logger.error(f"Failed to load WebDAV locks: {e}")

    async def _save_all_locks(self):
        if not self._persistence_enabled:
            return
        for lock_info in list(self.locks.values()):
            if not lock_info.is_expired:
                await self._persist_lock(lock_info)

    async def _background_cleanup(self):
        while self._running:
            try:
                await asyncio.sleep(60)
                await self._cleanup_expired()
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"Error in WebDAV lock cleanup: {e}")

    async def _cleanup_expired(self):
        async with self._lock:
            expired = [
                path_hash for path_hash, lock_info in self.locks.items()
                if lock_info.is_expired
            ]
            for path_hash in expired:
                del self.locks[path_hash]
                await self._delete_lock_from_db(path_hash)
            if expired:
                logger.debug(f"Cleaned up {len(expired)} expired WebDAV locks")

    async def get_stats(self) -> dict:
        async with self._lock:
            total = len(self.locks)
            active = sum(1 for l in self.locks.values() if not l.is_expired)
            return {
                "total_locks": total,
                "active_locks": active,
                "expired_locks": total - active,
            }


_webdav_locks: Optional[PersistentWebDAVLocks] = None


async def init_webdav_locks(db_manager=None) -> PersistentWebDAVLocks:
    global _webdav_locks
    _webdav_locks = PersistentWebDAVLocks()
    await _webdav_locks.start(db_manager)
    return _webdav_locks


async def shutdown_webdav_locks():
    global _webdav_locks
    if _webdav_locks:
        await _webdav_locks.stop()
        _webdav_locks = None


def get_webdav_locks() -> PersistentWebDAVLocks:
    if not _webdav_locks:
        raise RuntimeError("WebDAV locks not initialized")
    return _webdav_locks