diff --git a/requirements.txt b/requirements.txt index 63ed28b..80e5bb5 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,4 +7,5 @@ frozenlist>=1.4.1 idna>=3.7 multidict>=6.0.5 python-magic>=0.4.27 +redis>=5.0.8 yarl>=1.9.4 diff --git a/setup.py b/setup.py index 302ad49..508b17b 100644 --- a/setup.py +++ b/setup.py @@ -7,6 +7,11 @@ setup( author="ORUDO", author_email="root@orudo.ru", description="A simple Telegram bot that will allow you to upload torrent files / magnet links to a remote Torrent server (qBitTorrent, Transmission, etc.)", - install_requires=["aiohttp>=3.10.0", "aiofiles>=24.1.0", "aiofiles>=24.1.0"], + install_requires=[ + "aiohttp>=3.10.0", + "aiofiles>=24.1.0", + "aiofiles>=24.1.0", + "redis>=5.0.8", + ], packages=["tubot", "tubot.static", "tubot.torrent", "tubot.dirgetter", "db"], ) diff --git a/tubot/db/cache.py b/tubot/db/cache.py index 806c657..3553ec9 100644 --- a/tubot/db/cache.py +++ b/tubot/db/cache.py @@ -13,6 +13,9 @@ from aiofiles.ospath import isdir, isfile from aiofiles.os import mkdir from aiofiles import open from asyncio import sleep +from redis import asyncio as aioredis +from json import loads as json_loads +from json import dumps as json_dumps class PythonCache(CacheDB): @@ -100,3 +103,87 @@ class PythonCache(CacheDB): @property async def get_dirs(self) -> dict: return self.dirs + + +class RedisCache(CacheDB): + """ + Redis implementation of Cache DataBase + """ + + _ctype = CacheDBTypes.Redis + host: str + + def __init__(self, redis_host: str) -> None: + super().__init__() + self.host = redis_host + + async def __validate__(self) -> bool: + async with aioredis.from_url( + f"redis://{self.host}", encoding="utf-8", decode_responses=True + ) as redis: + return await redis.ping() + + # Users + + async def add_user(self, tg_id: int, name: str) -> None: + async with aioredis.from_url( + f"redis://{self.host}", encoding="utf-8", decode_responses=True + ) as redis: + user = User(tg_id, name) + json = json_dumps(user.dict) + await redis.set(str(tg_id), json) + + async def user_state(self, tg_id: int) -> UserStates: + async with aioredis.from_url( + f"redis://{self.host}", encoding="utf-8", decode_responses=True + ) as redis: + json = await redis.get(str(tg_id)) + user = json_loads(json) + return user["state"] + + async def change_user_state(self, tg_id: int, status: UserStates) -> None: + async with aioredis.from_url( + f"redis://{self.host}", encoding="utf-8", decode_responses=True + ) as redis: + json = await redis.get(str(tg_id)) + user = json_loads(json) + user["state"] = status + json = json_dumps(user) + await redis.set(str(tg_id), json) + + async def auth_user(self, tg_id: int) -> None: + async with aioredis.from_url( + f"redis://{self.host}", encoding="utf-8", decode_responses=True + ) as redis: + json = await redis.get(str(tg_id)) + user = json_loads(json) + user["auth"] = True + json = json_dumps(user) + await redis.set(str(tg_id), json) + + async def is_user_auth(self, tg_id: int) -> bool: + async with aioredis.from_url( + f"redis://{self.host}", encoding="utf-8", decode_responses=True + ) as redis: + json = await redis.get(str(tg_id)) + user = json_loads(json) + return bool(user["auth"]) + + # Dirs + + async def cache_dirs(self, dirs: dict, expire: int) -> None: + async with aioredis.from_url( + f"redis://{self.host}", encoding="utf-8", decode_responses=True + ) as redis: + json = json_dumps(dirs) + await redis.set("dirs", json, ex=expire) + + @property + async def get_dirs(self) -> dict: + async with aioredis.from_url( + f"redis://{self.host}", encoding="utf-8", decode_responses=True + ) as redis: + resp = await redis.get("dirs") + if resp is None: + return {} + return json_loads(resp) diff --git a/tubot/db/types.py b/tubot/db/types.py index d2cf3d5..c17f6bc 100644 --- a/tubot/db/types.py +++ b/tubot/db/types.py @@ -14,6 +14,7 @@ class CacheDBTypes(Enum): """ PythonPKL = "python" + Redis = "redis" class UserStates(Enum): diff --git a/tubot/static/abc.py b/tubot/static/abc.py index d2e3f58..2b93b4b 100644 --- a/tubot/static/abc.py +++ b/tubot/static/abc.py @@ -26,7 +26,6 @@ class IValidatable(ABC): class ENV(object): - _name: str | None = None DEFAULT: str diff --git a/tubot/static/env.py b/tubot/static/env.py index 9a4e66d..90a8596 100644 --- a/tubot/static/env.py +++ b/tubot/static/env.py @@ -15,3 +15,12 @@ class PN_CACHE(ENV): _name = "PN_CACHE" DEFAULT = "/etc/tubot" + + +class REDIS_HOST(ENV): + """ + Redis host adress + """ + + _name = "REDIS_HOST" + DEFAULT = "localhost:6379"