mdrs-client-python/mdrsclient/cache.py

import dataclasses
import hashlib
import json
import os
from typing import Protocol, runtime_checkable

from pydantic import TypeAdapter, ValidationError
from pydantic.dataclasses import dataclass

from mdrsclient.exceptions import UnexpectedException
from mdrsclient.models import Laboratories, Token, User
from mdrsclient.settings import CONFIG_DIRNAME
from mdrsclient.utils import FileLock


@dataclass
class CacheData:
    user: User | None = None
    token: Token | None = None
    laboratories: Laboratories = dataclasses.field(default_factory=Laboratories)
    digest: str = ""

    def clear(self) -> None:
        self.user = None
        self.token = None
        self.laboratories.clear()
        self.digest = ""

    def update_digest(self) -> None:
        self.digest = self.__calc_digest()

    def verify_digest(self) -> bool:
        return self.digest == self.__calc_digest()

    def __calc_digest(self) -> str:
        return hashlib.sha256(
            json.dumps(
                [
                    None if self.user is None else dataclasses.asdict(self.user),
                    None if self.token is None else dataclasses.asdict(self.token),
                    dataclasses.asdict(self.laboratories),
                ]
            ).encode("utf-8")
        ).hexdigest()


@runtime_checkable
class CacheInterface(Protocol):
    @property
    def token(self) -> Token | None: ...
    @token.setter
    def token(self, token: Token) -> None: ...
    @token.deleter
    def token(self) -> None: ...

    @property
    def user(self) -> User | None: ...
    @user.setter
    def user(self, user: User) -> None: ...
    @user.deleter
    def user(self) -> None: ...

    @property
    def laboratories(self) -> Laboratories: ...
    @laboratories.setter
    def laboratories(self, laboratories: Laboratories) -> None: ...


class InMemoryCache(CacheInterface):
    def __init__(self) -> None:
        self.__data = CacheData()

    @property
    def token(self) -> Token | None:
        return self.__data.token

    @token.setter
    def token(self, token: Token) -> None:
        self.__data.token = token

    @token.deleter
    def token(self) -> None:
        if self.__data.token is not None:
            self.__data.token = None

    @property
    def user(self) -> User | None:
        return self.__data.user

    @user.setter
    def user(self, user: User) -> None:
        self.__data.user = user

    @user.deleter
    def user(self) -> None:
        if self.__data.user is not None:
            self.__data.user = None

    @property
    def laboratories(self) -> Laboratories:
        return self.__data.laboratories

    @laboratories.setter
    def laboratories(self, laboratories: Laboratories) -> None:
        self.__data.laboratories = laboratories


class CacheFile(CacheInterface):
    __serial: int
    __cache_dir: str
    __cache_file: str
    __data: CacheData

    def __init__(self, remote: str) -> None:
        self.__serial = -1
        self.__cache_dir = os.path.join(CONFIG_DIRNAME, "cache")
        self.__cache_file = os.path.join(self.__cache_dir, remote + ".json")
        self.__data = CacheData()

    @property
    def token(self) -> Token | None:
        self.__load()
        return self.__data.token

    @token.setter
    def token(self, token: Token) -> None:
        self.__load()
        self.__data.token = token
        self.__save()

    @token.deleter
    def token(self) -> None:
        if self.__data.token is not None:
            self.__clear()

    @property
    def user(self) -> User | None:
        return self.__data.user

    @user.setter
    def user(self, user: User) -> None:
        self.__load()
        self.__data.user = user
        self.__save()

    @user.deleter
    def user(self) -> None:
        if self.__data.user is not None:
            self.__clear()

    @property
    def laboratories(self) -> Laboratories:
        return self.__data.laboratories

    @laboratories.setter
    def laboratories(self, laboratories: Laboratories) -> None:
        self.__load()
        self.__data.laboratories = laboratories
        self.__save()

    def __clear(self) -> None:
        self.__data.clear()
        self.__save()

    def __load(self) -> None:
        if os.path.isfile(self.__cache_file):
            stat = os.stat(self.__cache_file)
            serial = hash((stat.st_uid, stat.st_gid, stat.st_mode, stat.st_size, stat.st_mtime))
            if self.__serial != serial:
                try:
                    with open(self.__cache_file) as f:
                        data = TypeAdapter(CacheData).validate_python(json.load(f))
                        if not data.verify_digest():
                            raise UnexpectedException("Cache data has been broken.")
                        self.__data = data
                except (ValidationError, UnexpectedException) as e:
                    self.__clear()
                    self.__save()
                    print(e)
                else:
                    self.__serial = serial
        else:
            self.__clear()
            self.__serial = -1

    def __save(self) -> None:
        self.__ensure_cache_dir()
        with open(self.__cache_file, "w") as f:
            FileLock.lock(f)
            self.__data.update_digest()
            f.write(json.dumps(dataclasses.asdict(self.__data)))
            FileLock.unlock(f)
        stat = os.stat(self.__cache_file)
        self.__serial = hash((stat.st_uid, stat.st_gid, stat.st_mode, stat.st_size, stat.st_mtime))
        # ensure file is secure.
        os.chmod(self.__cache_file, 0o600)

    def __ensure_cache_dir(self) -> None:
        if not os.path.exists(self.__cache_dir):
            os.makedirs(self.__cache_dir)
        # ensure directory is secure.
        os.chmod(self.__cache_dir, 0o700)