commit acbccb4f2886f59fdb9e332ba1dd3819ce3754c4 Author: timofey Date: Thu Aug 24 19:22:12 2023 +0000 initial commit diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..39d1cb0 --- /dev/null +++ b/.gitignore @@ -0,0 +1,164 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +#poetry.lock + +# pdm +# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. +#pdm.lock +# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it +# in version control. +# https://pdm.fming.dev/#use-with-ide +.pdm.toml + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# PyCharm +# JetBrains specific template is maintained in a separate JetBrains.gitignore that can +# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore +# and can be added to the global gitignore or merged into this file. For a more nuclear +# option (not recommended) you can uncomment the following to ignore the entire idea folder. +#.idea/ + +# Environment variables files +*.env +.secrets diff --git a/.vscode/settings.json b/.vscode/settings.json new file mode 100644 index 0000000..941b5b7 --- /dev/null +++ b/.vscode/settings.json @@ -0,0 +1,13 @@ +{ + "python.analysis.typeCheckingMode": "basic", + "python.analysis.extraPaths": [ + "starbot" + ], + "python.testing.pytestEnabled": false, + "python.testing.unittestEnabled": true, + "search.exclude": { + "**/venv": true + }, + "isort.args": ["--profile", "black"], + "black-formatter.args": ["--line-length", "120"] +} diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000..1ee2793 --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,26 @@ +volumes: + stardata: + external: false + +services: + starbot: + build: starbot + volumes: + - "./starbot/starbot:/app/starbot:ro" + - "stardata:/app/data/:rw" + env_file: + - .secrets/starbot.env + environment: + - DBF_MODULE=starbot.db_ptvp35 + deploy: + resources: + limits: + cpus: '2' + memory: 200M + restart_policy: + condition: on-failure + delay: 5s + max_attempts: 3 + window: 120s + tty: true + stop_signal: SIGINT diff --git a/starbot/Dockerfile b/starbot/Dockerfile new file mode 100644 index 0000000..4588ce2 --- /dev/null +++ b/starbot/Dockerfile @@ -0,0 +1,6 @@ +FROM python:3.11 +WORKDIR /app/ +COPY requirements.txt requirements.txt +RUN pip install -r requirements.txt +COPY starbot starbot +CMD ["python3", "-m", "starbot"] diff --git a/starbot/pyproject.toml b/starbot/pyproject.toml new file mode 100644 index 0000000..e69de29 diff --git a/starbot/requirements.txt b/starbot/requirements.txt new file mode 100644 index 0000000..9f07e2d --- /dev/null +++ b/starbot/requirements.txt @@ -0,0 +1,4 @@ +aiohttp>=3.7.4,<4 +aiosqlite~=0.19 +discord.py~=2.3.2 +ptvp35 @ git+https://gitea.parrrate.ru/PTV/ptvp35.git@f8ee5d20f4e159df2e20c40dbf3b81e925c2db36 diff --git a/starbot/starbot/__main__.py b/starbot/starbot/__main__.py new file mode 100644 index 0000000..7e05169 --- /dev/null +++ b/starbot/starbot/__main__.py @@ -0,0 +1,57 @@ +from __future__ import annotations + +import asyncio +import importlib +import os + +import discord +from discord.ext import commands + +from .bot import StarBot, state_manager +from .db import AbstractDbFactory + + +async def _main(): + token = os.getenv("DISCORD_TOKEN") + assert token is not None + dbfm = os.getenv("DBF_MODULE") + assert dbfm is not None + DBF: AbstractDbFactory = importlib.import_module(dbfm).DBF + async with state_manager(DBF) as state: + bot = StarBot(state) + await bot.login(token) + yield bot + await bot.load_extension("starbot.stars") + await bot.connect() + yield bot + + +async def aclose(client: discord.Client): + if not client.is_closed(): + await client.change_presence(status=discord.Status.offline) + await client.close() + + +def close(client: discord.Client, loop: asyncio.AbstractEventLoop): + loop.run_until_complete(aclose(client)) + + +def main(): + loop = asyncio.new_event_loop() + gen = _main() + bot: commands.Bot = loop.run_until_complete(anext(gen)) + discord.utils.setup_logging() + + async def complete(): + return await anext(gen) + + task = loop.create_task(complete()) + try: + loop.run_until_complete(task) + except (KeyboardInterrupt, InterruptedError, RuntimeError): + close(bot, loop) + loop.run_until_complete(task) + + +if __name__ == "__main__": + main() diff --git a/starbot/starbot/bot.py b/starbot/starbot/bot.py new file mode 100644 index 0000000..b90e2ca --- /dev/null +++ b/starbot/starbot/bot.py @@ -0,0 +1,29 @@ +from __future__ import annotations + +from contextlib import AsyncExitStack, asynccontextmanager +from pathlib import Path + +import discord +from discord.ext import commands + +from .db import AbstractConnection, AbstractDbFactory + + +class StarState: + def __init__(self, connection: AbstractConnection) -> None: + self.connection = connection + + +@asynccontextmanager +async def state_manager(dbf: AbstractDbFactory): + async with dbf.from_path(Path("/app/data/stars.db")) as connection: + yield StarState(connection) + + +class StarBot(commands.Bot): + def __init__(self, state: StarState) -> None: + super().__init__( + command_prefix="⭐", + intents=discord.Intents(message_content=True, guild_messages=True, guild_reactions=True, guilds=True), + ) + self.starstate: StarState = state diff --git a/starbot/starbot/db.py b/starbot/starbot/db.py new file mode 100644 index 0000000..c531e5c --- /dev/null +++ b/starbot/starbot/db.py @@ -0,0 +1,26 @@ +from __future__ import annotations + +from collections.abc import Hashable +from pathlib import Path +from typing import Any, Protocol + + +class AbstractConnection(Protocol): + def get(self, key: Hashable, default: Any, /) -> Any: + raise NotImplementedError + + async def set(self, key: Hashable, value: Any, /) -> None: + raise NotImplementedError + + +class AbstractDbManager(Protocol): + async def __aenter__(self) -> AbstractConnection: + raise NotImplementedError + + async def __aexit__(self, et, ev, tb, /): + raise NotImplementedError + + +class AbstractDbFactory(Protocol): + def from_path(self, path: Path) -> AbstractDbManager: + raise NotImplementedError diff --git a/starbot/starbot/db_aiosqlite.py b/starbot/starbot/db_aiosqlite.py new file mode 100644 index 0000000..52b07b1 --- /dev/null +++ b/starbot/starbot/db_aiosqlite.py @@ -0,0 +1,56 @@ +import json +from collections.abc import Hashable +from contextlib import asynccontextmanager +from pathlib import Path +from typing import Any + +import aiosqlite + +from starbot.db import AbstractDbManager + +from .db import AbstractConnection, AbstractDbFactory, AbstractDbManager + + +def _load_key(key: Any, /) -> Hashable: + """note: unstable signature.""" + match key: + case Hashable(): + return key + case list(): + return tuple(map(_load_key, key)) + case dict(): + return tuple((_load_key(k), _load_key(v)) for k, v in key.items()) + case _: + raise TypeError("unknown json key type, cannot convert to hashable") + + +class Adapter(AbstractConnection): + def __init__(self, connection: aiosqlite.Connection, db: dict[Hashable, Any]) -> None: + self.connection = connection + self.db = db + + def get(self, key: Hashable, default: Any) -> Any: + return self.db.get(key, default) + + async def set(self, key: Hashable, value: Any) -> None: + self.db[key] = value + await self.connection.execute("REPLACE INTO kv(key, value) VALUES (?, ?)", (json.dumps(key), json.dumps(value))) + await self.connection.commit() + + +@asynccontextmanager +async def manager(path: Path): + async with aiosqlite.connect(path) as connection: + db: dict[Hashable, Any] = {} + await connection.execute("CREATE TABLE IF NOT EXISTS kv(key, value)") + for key, value in await connection.execute_fetchall("SELECT key, value FROM kv"): + db[_load_key(json.loads(key))] = json.loads(value) + yield Adapter(connection, db) + + +class Factory(AbstractDbFactory): + def from_path(self, path: Path) -> AbstractDbManager: + return manager(path.with_suffix(path.suffix + "sqlite")) + + +DBF: AbstractDbFactory = Factory() diff --git a/starbot/starbot/db_ptvp35.py b/starbot/starbot/db_ptvp35.py new file mode 100644 index 0000000..9450bb4 --- /dev/null +++ b/starbot/starbot/db_ptvp35.py @@ -0,0 +1,13 @@ +from pathlib import Path + +from ptvp35 import DbManager, KVJson + +from .db import AbstractDbFactory, AbstractDbManager + + +class Ptvp35: + def from_path(self, path: Path) -> AbstractDbManager: + return DbManager(path, kvfactory=KVJson()) + + +DBF: AbstractDbFactory = Ptvp35() diff --git a/starbot/starbot/stars.py b/starbot/starbot/stars.py new file mode 100644 index 0000000..84f5392 --- /dev/null +++ b/starbot/starbot/stars.py @@ -0,0 +1,167 @@ +from __future__ import annotations + +import asyncio +from collections.abc import Hashable + +import discord +from discord.ext import commands + +from .bot import StarBot, StarState +from .db import AbstractConnection + + +class Locks: + def __init__(self) -> None: + self.locks: dict[Hashable, asyncio.Lock] = {} + + def lock_for(self, key: Hashable) -> asyncio.Lock: + if key in self.locks: + return self.locks[key] + else: + return self.locks.setdefault(key, asyncio.Lock()) + + +locks = Locks() +lock_for = locks.lock_for + + +class StarCtx: + def __init__(self, ctx: commands.Context) -> None: + self.ctx = ctx + assert ctx.guild + self.guild: discord.Guild = ctx.guild + self.bot: StarBot = ctx.bot + self.state: StarState = self.bot.starstate + self.connection: AbstractConnection = self.state.connection + + async def assign(self, count: int) -> None: + await self.connection.set(("assign", self.guild.id), {"channel": self.ctx.channel.id, "count": count}) + + async def unassign(self) -> None: + await self.connection.set(("assign", self.guild.id), None) + + +class ReactionCtx: + def __init__(self, bot: StarBot, event: discord.RawReactionActionEvent): + self.bot = bot + self.guild_id = event.guild_id + self.channel_id = event.channel_id + self.message_id = event.message_id + self.name = event.emoji.name + self.state: StarState = self.bot.starstate + self.connection: AbstractConnection = self.state.connection + + async def get_channel(self, id_: int) -> discord.abc.Messageable: + channel = self.bot.get_channel(id_) + if channel is None: + channel = await self.bot.fetch_channel(id_) + if isinstance(channel, discord.CategoryChannel): + raise TypeError + if isinstance(channel, discord.ForumChannel): + raise TypeError + if isinstance(channel, discord.abc.PrivateChannel): + raise TypeError + return channel + + async def on(self) -> None: + if self.name != "⭐": + return + assignment: dict[str, int] | None = self.connection.get(("assign", self.guild_id), None) + if assignment is None: + return + assigned_to, count = assignment["channel"], assignment["count"] + if self.channel_id == assigned_to: + return + async with lock_for(self.message_id): + assigned_channel, event_channel = await asyncio.gather( + self.get_channel(assigned_to), self.get_channel(self.channel_id) + ) + message = await event_channel.fetch_message(self.message_id) + reaction = next((reaction for reaction in message.reactions if reaction.emoji == "⭐"), None) + if reaction is None: + return + if reaction.me: + return + if reaction.count >= count: + guild = message.guild + if guild is None: + return + member = guild.get_member(message.author.id) + if member is None: + member = await guild.fetch_member(message.author.id) + embed = discord.Embed(description=message.content or None) + avatar = member.avatar + embed.set_author(name=member.display_name, url=message.jump_url, icon_url=avatar and avatar.url) + image = next( + ( + attachment.url + for attachment in message.attachments + if (attachment.content_type or "").startswith("image/") + ), + None, + ) + if image is not None: + embed.set_image(url=image) + await asyncio.gather( + message.add_reaction("⭐"), + assigned_channel.send(embed=embed), + ) + + +class Stars(commands.Cog): + def __init__(self, bot: StarBot) -> None: + super().__init__() + self.__bot = bot + + @commands.hybrid_command() + async def ping(self, ctx: commands.Context): + print("ping pong") + await ctx.reply("pong", mention_author=False) + + @commands.hybrid_command() + @commands.is_owner() + async def reload(self, ctx: commands.Context): + print("reload") + bot: commands.Bot = ctx.bot + try: + await bot.reload_extension("starbot.stars") + except commands.ExtensionNotLoaded: + await ctx.reply("not loaded") + print("reloaded") + await ctx.reply("reloaded") + + @commands.hybrid_command() + @commands.is_owner() + async def sync(self, ctx: commands.Context): + await ctx.bot.tree.sync() + print("synced") + await ctx.reply("synced") + + @commands.hybrid_command() + @commands.has_permissions(administrator=True) + async def assign(self, ctx: commands.Context, count: int): + await StarCtx(ctx).assign(count) + await ctx.reply("assigned") + + @commands.hybrid_command() + @commands.has_permissions(administrator=True) + async def unassign(self, ctx: commands.Context): + await StarCtx(ctx).unassign() + await ctx.reply("unassigned") + + @commands.Cog.listener() + async def on_raw_reaction_add(self, event: discord.RawReactionActionEvent): + await ReactionCtx(self.__bot, event).on() + + +async def setup(bot: StarBot): + global cog + cog = Stars(bot) + await bot.add_cog(cog) + + +async def teardown(bot: StarBot): + global cog + await bot.remove_cog(cog.qualified_name) + del cog + print("torn down")