diff options
| author | Nakidai <nakidai@disroot.org> | 2024-10-28 15:55:04 +0300 |
|---|---|---|
| committer | Nakidai <nakidai@disroot.org> | 2024-10-28 15:55:04 +0300 |
| commit | baaf3995646630de931113d4c554fd4867977b5c (patch) | |
| tree | 1e54150e5cc7f1a6a51022ea08d9a5c969e269eb | |
| parent | 94f682c11acbda7af6ee23eafa7db515c8d71ae4 (diff) | |
| download | petthecord-baaf3995646630de931113d4c554fd4867977b5c.tar.gz petthecord-baaf3995646630de931113d4c554fd4867977b5c.zip | |
Refactor
- Move caching to another file - Rename some files - Make loggers look ok
| -rw-r--r-- | src/petthecord/bot.py | 67 | ||||
| -rw-r--r-- | src/petthecord/cache.py | 114 | ||||
| -rw-r--r-- | src/petthecord/main.py | 11 | ||||
| -rw-r--r-- | src/petthecord/runner.py | 57 | ||||
| -rw-r--r-- | src/petthecord/server.py | 103 |
5 files changed, 187 insertions, 165 deletions
diff --git a/src/petthecord/bot.py b/src/petthecord/bot.py index a7e08e3..59827f2 100644 --- a/src/petthecord/bot.py +++ b/src/petthecord/bot.py @@ -1,26 +1,21 @@ from logging import getLogger -from aiohttp.web import AppRunner, TCPSite -from discord import app_commands, Interaction, Intents, User +from discord import app_commands, Interaction, User from discord.ext import commands -from .server import Server - - -class PatTheCordCog(commands.Cog): - def __init__(self, client: commands.Bot, origin: str = "https://ptc.pwn3t.ru") -> None: - self.client = client - self.origin = origin +class PetTheCordCog(commands.Cog): + def __init__(self, origin: str = "https://ptc.pwn3t.ru") -> None: + self._origin = origin super().__init__() - self._logger = getLogger("petthecord.bot") + self._logger = getLogger(__name__) @app_commands.allowed_installs(users=True) @app_commands.allowed_contexts(guilds=True, dms=True, private_channels=True) @app_commands.command( name="petpet", - description="Pat some user" + description="Pet some user" ) @app_commands.describe( user="User to pet", @@ -34,52 +29,4 @@ class PatTheCordCog(commands.Cog): r: str = "" ) -> None: self._logger.info(f"Petting {user.id} for {interaction.user.id}") - await interaction.response.send_message(f"{self.origin}/{user.id}.{r}.gif") - - -class Bot(commands.Bot): - def __init__( - self, - - host: str = "127.0.0.1", - port: int = 8080, - origin: str = "https://ptc.pwn3t.ru", - caching: bool = True, - cache_path: str = "/var/cache/petthecord", - cache_lifetime: int = 86400, - cache_gc_delay: int = 14400, - ) -> None: - super().__init__( - command_prefix="!", - intents=Intents.default() - ) - self._host = host - self._port = port - self._origin = origin - self._caching = caching - self._cache_path = cache_path - self._cache_lifetime = cache_lifetime - self._cache_gc_delay = cache_gc_delay - - self._logger = getLogger("petthecord") - - async def on_ready(self) -> None: - await self.add_cog(PatTheCordCog(self, self._origin)) - await self.tree.sync() - - server = Server( - self, - self._caching, - self._cache_path, - self._cache_lifetime, - self._cache_gc_delay - ) - runner = AppRunner(server) - await runner.setup() - site = TCPSite(runner, self._host, self._port) - await site.start() - - getLogger("petthecord.server").info(f"Started serving on {self._host}:{self._port}") - - if self._caching: - await server.clean_cache() + await interaction.response.send_message(f"{self._origin}/{user.id}.{r}.gif") diff --git a/src/petthecord/cache.py b/src/petthecord/cache.py new file mode 100644 index 0000000..1b4cd3a --- /dev/null +++ b/src/petthecord/cache.py @@ -0,0 +1,114 @@ +from asyncio import sleep +from json import load, dump +from io import BytesIO +from logging import getLogger +from os import PathLike, makedirs, remove, listdir +from os.path import getmtime +from pathlib import Path +from time import time +from typing import NoReturn + +import discord +from discord import Client +from petpetgif import petpet + + +class NotFound(Exception): + pass + + +class HTTPException(Exception): + pass + + +class CachedPet: + def __init__( + self, + + client: Client, + caching: bool = True, + path: str | PathLike = "/var/cache/petthecord", + lifetime: int = 86400, + gc_delay: int = 14400, + ) -> None: + self._client = client + self._caching = caching + self._path = Path(path).resolve() + self._lifetime = lifetime + self._gc_delay = gc_delay + + self._logger = getLogger(__name__) + + index_path = self._path / "index.json" + try: + if not index_path.exists(): + self._logger.warning("`index.json` doesn't exist, trying to create...") + if not self._path.exists(): + self._logger.warning("Cache folder doesnt exist, trying to create...") + makedirs(self._path, mode=0o755, exist_ok=True) + with open(index_path, "w") as f: + f.write("{}") + except OSError: + self._logger.error("Cannot create environment") + raise + + with open(index_path, "r") as f: + self._cache = load(f) + + async def petpet(self, uid: int) -> bytes: + try: + user = await self._client.fetch_user(uid) + except discord.NotFound: + raise NotFound + except discord.HTTPException: + raise HTTPException + + if user.avatar is None: + raise NotFound + + if self._caching: + avatar_path = self._path / f"{user.id}_{user.avatar.key}.gif" + if (path := self._cache.get(user.id)) != str(avatar_path): + self._logger.debug("Generating new gif for {user.id}") + if path: + remove(path) + self._cache[user.id] = str(avatar_path) + with open(self._path / "index.json", "w") as f: + dump(self._cache, f) + + if not avatar_path.exists(): + with open(avatar_path, "wb") as f: + image = await user.avatar.read() + petpet.make(BytesIO(image), f) + + avatar_path.touch() + + with open(avatar_path, "rb") as f: + return f.read() + else: + with BytesIO() as f: + image = await user.avatar.read() + petpet.make(BytesIO(image), f) + + f.seek(0) + + return f.read() + + async def gc_loop(self) -> NoReturn: + while True: + if self._caching: + self._logger.info("Starting new cache's gc iteration") + + for filename in listdir(self._path): + path = (self._path / filename) + if path.is_file() and filename != "index.json": + if (time() - getmtime(path) > self._lifetime): + self._logger.debug(f"Removing {filename}") + del self._cache[filename.split('_')[0]] + remove(path) + with open(self._path / "index.json", "w") as f: + dump(self._cache, f) + + self._logger.debug("Finished collecting old cache") + + await sleep(self._gc_delay) diff --git a/src/petthecord/main.py b/src/petthecord/main.py index f85d473..0b1d144 100644 --- a/src/petthecord/main.py +++ b/src/petthecord/main.py @@ -1,9 +1,8 @@ from argparse import ArgumentParser -from logging import Formatter, StreamHandler, Logger, getLogger from os import getenv from sys import argv, stderr -from .bot import Bot +from .runner import Bot def main() -> None: @@ -58,14 +57,6 @@ def main() -> None: ) args = parser.parse_args() - # logger = getLogger() - # handler = StreamHandler() - # handler.setFormatter( - # Formatter("[{asctime}] [{levelname:<8}] {name}: {message}", '%Y-%m-%d %H:%M:%S', style="{") - # ) - # logger.addHandler(handler) - # logger.log(1, "Hello!") - bot = Bot( args.host, args.port, diff --git a/src/petthecord/runner.py b/src/petthecord/runner.py new file mode 100644 index 0000000..f43f148 --- /dev/null +++ b/src/petthecord/runner.py @@ -0,0 +1,57 @@ +from logging import getLogger + +from aiohttp.web import AppRunner, TCPSite +from discord import Intents +from discord.ext import commands + +from .bot import PetTheCordCog +from .cache import CachedPet +from .server import Server + + +class Bot(commands.Bot): + def __init__( + self, + + host: str = "127.0.0.1", + port: int = 8080, + origin: str = "https://ptc.pwn3t.ru", + caching: bool = True, + cache_path: str = "/var/cache/petthecord", + cache_lifetime: int = 86400, + cache_gc_delay: int = 14400, + ) -> None: + super().__init__( + command_prefix="!", + intents=Intents.default() + ) + self._host = host + self._port = port + self._origin = origin + self._caching = caching + self._cache_path = cache_path + self._cache_lifetime = cache_lifetime + self._cache_gc_delay = cache_gc_delay + + self._logger = getLogger(__name__) + + async def on_ready(self) -> None: + await self.add_cog(PetTheCordCog(self._origin)) + await self.tree.sync() + + petter = CachedPet( + self, + self._caching, + self._cache_path, + self._cache_lifetime, + self._cache_gc_delay + ) + runner = AppRunner(Server(petter)) + await runner.setup() + site = TCPSite(runner, self._host, self._port) + await site.start() + + self._logger.info(f"Started serving on {self._host}:{self._port}") + + if self._caching: + await petter.gc_loop() diff --git a/src/petthecord/server.py b/src/petthecord/server.py index 28b49ee..08d7389 100644 --- a/src/petthecord/server.py +++ b/src/petthecord/server.py @@ -1,55 +1,17 @@ -from asyncio import sleep -from json import load, dumps -from io import BytesIO from logging import getLogger -from os import listdir, makedirs, remove -from os.path import getmtime -from pathlib import Path -from time import time from typing import NoReturn from aiohttp.web import Application, StreamResponse, get, HTTPFound, Request, Response -from discord import NotFound, HTTPException -from discord.ext import commands -from petpetgif import petpet +from .cache import CachedPet, HTTPException, NotFound -class Server(Application): - def __init__( - self, - client: commands.Bot, - caching: bool = True, - cache_path: str = "/var/cache/petthecord", - cache_lifetime: int = 86400, - cache_gc_delay: int = 14400, - ) -> None: - self.client = client - self.caching = caching - self.cache_path = Path(cache_path).resolve() - self.cache_lifetime = cache_lifetime - self.cache_gc_delay = cache_gc_delay +class Server(Application): + def __init__(self, petter: CachedPet) -> None: + self._petter = petter super().__init__() - self._logger = getLogger("petthecord.server") - self._cache_logger = getLogger("petthecord.cache") - - if self.caching: - index_path = self.cache_path / "index.json" - try: - if not index_path.exists(): - self._cache_logger.warning("index.json doesn't exist, creating...") - if not self.cache_path.exists(): - self._cache_logger.warning(f"{self.cache_path} doesn't exist, creating..") - makedirs(self.cache_path, mode=0o755, exist_ok=True) - with open(index_path, "w") as f: - f.write("{}") - except OSError: - self._cache_logger.error("Cannot create environment, turning caching off") - self.caching = False - - with open(index_path, "r") as f: - self.cache = load(f) + self._logger = getLogger(__name__) self.add_routes( [ @@ -62,65 +24,16 @@ class Server(Application): raise HTTPFound("https://github.com/nakidai/petthecord") async def petpet(self, request: Request) -> StreamResponse: - self._logger.info(f"Incoming '{request.rel_url}' request from {request.remote}") try: uid = int(request.match_info["uid"][:request.match_info["uid"].find('.')]) except ValueError: return Response(status=400) + self._logger.info(f"Petting {uid} for {request.remote}") + try: - user = await self.client.fetch_user(uid) + return Response(body=await self._petter.petpet(uid), content_type="image/gif") except NotFound: return Response(status=404) except HTTPException: return Response(status=403) - - if user.avatar is None: - return Response(status=404) - - avatar_path = str(self.cache_path / f"{user.id}_{user.avatar.key}.gif") - if self.caching: - if (path := self.cache.get(user.id)) != avatar_path: - self._cache_logger.debug("Regenerating cached avatar {user.id}") - if path: - remove(path) - self.cache[user.id] = avatar_path - with open(self.cache_path / "index.json", "w") as f: - f.write(dumps(self.cache)) - - if not Path(avatar_path).exists(): - with open(avatar_path, "wb") as f: - image = await user.avatar.read() - petpet.make(BytesIO(image), f) - - Path(avatar_path).touch() - - with open(avatar_path, "rb") as f: - return Response(body=f.read(), content_type="image/gif") - else: - with BytesIO() as f: - image = await user.avatar.read() - petpet.make(BytesIO(image), f) - - f.seek(0) - - return Response(body=f.read(), content_type="image/gif") - - async def clean_cache(self) -> None: - while True: - if self.caching: - self._cache_logger.info("Starting new cache's gc iteration") - - for filename in listdir(self.cache_path): - path = (self.cache_path / filename) - if path.is_file() and filename != "index.json": - if (time() - getmtime(path) > self.cache_lifetime): - self._cache_logger.debug(f"Removing {filename}") - del self.cache[filename.split('_')[0]] - remove(path) - with open(self.cache_path / "index.json", "w") as f: - f.write(dumps(self.cache)) - - self._cache_logger.debug("Finished collecting old cache") - - await sleep(self.cache_gc_delay) |