about summary refs log tree commit diff
diff options
context:
space:
mode:
authorNakidai <nakidai@disroot.org>2024-10-28 15:55:04 +0300
committerNakidai <nakidai@disroot.org>2024-10-28 15:55:04 +0300
commitbaaf3995646630de931113d4c554fd4867977b5c (patch)
tree1e54150e5cc7f1a6a51022ea08d9a5c969e269eb
parent94f682c11acbda7af6ee23eafa7db515c8d71ae4 (diff)
downloadpetthecord-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.py67
-rw-r--r--src/petthecord/cache.py114
-rw-r--r--src/petthecord/main.py11
-rw-r--r--src/petthecord/runner.py57
-rw-r--r--src/petthecord/server.py103
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)