From eb1dc48a11334848f801bb86f4c06ebf7276b227 Mon Sep 17 00:00:00 2001 From: trueold89 Date: Wed, 7 Aug 2024 18:48:04 +0300 Subject: [PATCH] Add Controller object - init controller.py - Add Controller class - Add AUTH_PASSWD env var - Add new method 'is_user_exist' to CacheDB --- tubot/db/abc.py | 9 +++ tubot/db/cache.py | 16 +++++ tubot/static/controller.py | 119 +++++++++++++++++++++++++++++++++++++ tubot/static/env.py | 9 +++ tubot/static/init.py | 13 +++- 5 files changed, 164 insertions(+), 2 deletions(-) create mode 100644 tubot/static/controller.py diff --git a/tubot/db/abc.py b/tubot/db/abc.py index d6bdc7a..03b0435 100644 --- a/tubot/db/abc.py +++ b/tubot/db/abc.py @@ -43,6 +43,15 @@ class CacheDB(IValidatable, ABC): """ raise NotImplementedError + @abstractmethod + async def chech_user_existing(self, tg_id: int): + """ + Checks if user exist in db + + :param tg_id: User telegram id + """ + raise NotImplementedError + # Dirs @abstractmethod diff --git a/tubot/db/cache.py b/tubot/db/cache.py index d34c84c..9976663 100644 --- a/tubot/db/cache.py +++ b/tubot/db/cache.py @@ -68,6 +68,13 @@ class PythonCache(CacheDB): user_data = self.users[tg_id] return User.from_dict(user_data) + async def chech_user_existing(self, tg_id: int) -> bool: + try: + await self.read_user(tg_id) + return True + except KeyError: + return False + # Dirs async def cache_dirs(self, dirs: dict, expire: int) -> None: @@ -112,10 +119,19 @@ class RedisCache(CacheDB): f"redis://{self.host}", encoding="utf-8", decode_responses=True ) as redis: json = await redis.get(str(tg_id)) + if json is None: + raise KeyError user_data = json_loads(json) user_data["state"] = UserStates(user_data["state"]) return User.from_dict(user_data) + async def chech_user_existing(self, tg_id: int) -> bool: + try: + await self.read_user(tg_id) + return True + except KeyError: + return False + # Dirs async def cache_dirs(self, dirs: dict, expire: int) -> None: diff --git a/tubot/static/controller.py b/tubot/static/controller.py new file mode 100644 index 0000000..c4fd20f --- /dev/null +++ b/tubot/static/controller.py @@ -0,0 +1,119 @@ +# -*- coding: utf-8 -*- + +############################ +# Controller static module # +############################ + +# Imports +from tubot.torrent.abc import TorrentAPI, TorrentObj +from tubot.dirgetter.abc import DirGetter +from tubot.db.abc import CacheDB +from tubot.db.types import User, UserStates +from tubot.static.functions import validate +from tubot.static.env import AUTH_PASSWD + + +class Controller(object): + """ + Controller object + """ + + torrent: TorrentAPI + getter: DirGetter + cache: CacheDB + + def __init__(self, torrent_api: TorrentAPI, dg: DirGetter, cache: CacheDB) -> None: + """ + :param torrent_api: TorrentAPI module + :param dg: DirectoryGetter module + :param cache: CacheDB module + """ + self.torrent = torrent_api + self.getter = dg + self.cache = cache + + # DG + + async def get_dirs(self) -> dict: + """ + Returns dict of dirs from cache / dg + """ + dirs = await self.cache.get_dirs + if len(dirs) > 0: + return dirs + dirs = await self.getter.folders + if len(dirs) > 0: + return dirs + raise ValueError("No dirs found") + + # Torrent + + async def upload_torrent(self, torrent: TorrentObj) -> None: + """ + Add torrent to query + + :param torrent: Torrent object + """ + await validate(torrent) + await self.torrent.upload(torrent) + + async def get_torrent_list(self) -> str: + """ + Returns message with current torrents list + """ + return await self.torrent.torrent_list + + # Users + + async def _create_user(self, tg_id: int, name: str) -> User: + """ + Creates user in database + + :param tg_id: Telegram id + :param name: Telegram profile name + """ + user_obj = User(tg_id=tg_id, name=name) + await self.cache.write_user(tg_id, user_obj) + return user_obj + + async def _get_user_from_db(self, tg_id: int) -> User: + """ + Gets user from database + + :param tg_id: Telegram id + """ + return await self.cache.read_user(tg_id) + + async def get_user(self, tg_id: int, name: str) -> User: + """ + Returns user object + + :param tg_id: Telegram id + :param name: Telegram profile name + """ + if await self.cache.chech_user_existing(tg_id): + return await self._get_user_from_db(tg_id) + return await self._create_user(tg_id, name) + + async def auth_user(self, user: User, pwd: str) -> None: + """ + Auth user + + :param user: Current user object + """ + if user.auth: + raise ValueError("Already auth") + if pwd == AUTH_PASSWD()(): + user.auth = True + await self.cache.write_user(user.tg_id, user) + raise ValueError("Wrong password") + + async def set_user_state(self, user: User, state: UserStates) -> None: + """ + Change user status + + :param user: Current user object + :param state: New user status + """ + user.state = state + await self.cache.write_user(user.tg_id, user) diff --git a/tubot/static/env.py b/tubot/static/env.py index 501548b..d83540b 100644 --- a/tubot/static/env.py +++ b/tubot/static/env.py @@ -105,3 +105,12 @@ class TS_HOST(ENV): _name = "TS_HOST" DEFAULT = "http://localhost" + + +class AUTH_PASSWD(ENV): + """ + Password for users auth + """ + + _name = "AUTH_PASSWD" + DEFAULT = "changeme" diff --git a/tubot/static/init.py b/tubot/static/init.py index 587b0e7..25f1a90 100644 --- a/tubot/static/init.py +++ b/tubot/static/init.py @@ -77,7 +77,11 @@ class InitBuilder(object): await validate(self.CACHE) async def init_all_modules(self) -> None: - tasks = (create_task(self.init_ts()), create_task(self.init_dg()), create_task(self.init_cache())) + tasks = ( + create_task(self.init_ts()), + create_task(self.init_dg()), + create_task(self.init_cache()), + ) await gather(*tasks) @property @@ -89,6 +93,11 @@ async def init_modules() -> Iterable: ts = ServerTypes(env.TORRENT_SERVER()()) dg = GetterTypes(env.DIR_GETTER()()) cache = CacheDBTypes(env.CACHE_TYPE()()) - builder = InitBuilder().set_torrent_server(ts).set_directory_getter(dg).set_cache_type(cache) + builder = ( + InitBuilder() + .set_torrent_server(ts) + .set_directory_getter(dg) + .set_cache_type(cache) + ) await builder.init_all_modules() return builder.tuple