From d40b8dfcaa1d0fb0c0680c75efc764ab3b733ac4 Mon Sep 17 00:00:00 2001 From: Luna Mendes Date: Fri, 1 Dec 2017 22:04:35 -0300 Subject: [PATCH] add the bot --- Makefile | 4 ++ __init__.py | 0 bot/__init__.py | 1 + bot/bot.py | 56 +++++++++++++++++ bot/ext/admin.py | 126 ++++++++++++++++++++++++++++++++++++++ bot/ext/basic.py | 53 ++++++++++++++++ bot/ext/common.py | 9 +++ bot/ext/utils/__init__.py | 2 + bot/ext/utils/mousey.py | 97 +++++++++++++++++++++++++++++ config.py | 3 + memed.py | 9 ++- requirements.txt | 2 + 12 files changed, 360 insertions(+), 2 deletions(-) create mode 100644 Makefile create mode 100644 __init__.py create mode 100644 bot/__init__.py create mode 100644 bot/bot.py create mode 100644 bot/ext/admin.py create mode 100644 bot/ext/basic.py create mode 100644 bot/ext/common.py create mode 100644 bot/ext/utils/__init__.py create mode 100644 bot/ext/utils/mousey.py diff --git a/Makefile b/Makefile new file mode 100644 index 0000000..1130004 --- /dev/null +++ b/Makefile @@ -0,0 +1,4 @@ +OBFUSCATE="./env/bin/pyminifier --obfuscate --gzip" + +all: + ${OBFUSCATE} -o build/memed.py memed.py diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/bot/__init__.py b/bot/__init__.py new file mode 100644 index 0000000..79d300f --- /dev/null +++ b/bot/__init__.py @@ -0,0 +1 @@ +from .bot import * diff --git a/bot/bot.py b/bot/bot.py new file mode 100644 index 0000000..7de0e10 --- /dev/null +++ b/bot/bot.py @@ -0,0 +1,56 @@ +import logging +import time + +import discord +from discord.ext import commands + +log = logging.getLogger(__name__) + +exts = [ + 'basic', + 'admin' +] + + +class SexDungeonMistressBot(commands.Bot): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.init_time = time.monotonic() + + async def on_command(self, ctx): + # thanks dogbot ur a good + content = ctx.message.content + + author = ctx.message.author + guild = ctx.guild + checks = [c.__qualname__.split('.')[0] for c in ctx.command.checks] + location = '[DM]' if isinstance(ctx.channel, discord.DMChannel) else \ + f'[Guild {guild.name} {guild.id}]' + + log.info('%s [cmd] %s(%d) "%s" checks=%s', location, author, + author.id, content, ','.join(checks) or '(none)') + + +def schedule_bot(loop, config, db): + mute = ['discord', 'websockets'] + + for l in mute: + d = logging.getLogger(l) + d.setLevel(logging.INFO) + + bot = SexDungeonMistressBot( + command_prefix='sex ', + description='sexhouse management bot', + owner_id=getattr(config, 'owner_id', None), + ) + + bot.db = db + + try: + for ext in exts: + bot.load_extension(f'bot.ext.{ext}') + log.info('loaded %s', ext) + + loop.create_task(bot.start(config.bot_token)) + except: + log.exception('failed to load %s', ext) diff --git a/bot/ext/admin.py b/bot/ext/admin.py new file mode 100644 index 0000000..74e9b7a --- /dev/null +++ b/bot/ext/admin.py @@ -0,0 +1,126 @@ +import logging +import traceback + +import asyncpg + +from discord.ext import commands + +from .common import Cog +from .utils import Table, Timer + +log = logging.getLogger(__name__) + + +def no_codeblock(text: str) -> str: + """ + Removes codeblocks (grave accents), python and sql syntax highlight + indicators from a text if present. + .. note:: only the start of a string is checked, the text is allowed + to have grave accents in the middle + """ + if text.startswith('```'): + text = text[3:-3] + + if text.startswith(('py', 'sql')): + # cut off the first line as this removes the + # highlight indicator regardless of length + text = '\n'.join(text.split('\n')[1:]) + + if text.startswith('`'): + text = text[1:-1] + + return text + + +class Admin(Cog): + async def __local_check(self, ctx): + return ctx.author.id == self.bot.owner_id + + @commands.command() + @commands.is_owner() + async def shutdown(self, ctx): + log.info('Logging out! %s', ctx.author) + await ctx.send("dude rip") + await self.bot.logout() + + @commands.command() + @commands.is_owner() + async def load(self, ctx, *exts: str): + """Loads an extension.""" + for ext in exts: + try: + self.bot.load_extension("bot.ext." + ext) + except Exception as e: + await ctx.send(f'Oops. ```py\n{traceback.format_exc()}\n```') + return + log.info(f'Loaded {ext}') + m = ctx.send(f':ok_hand: `{ext}` loaded.') + self.bot.loop.create_task(m) + + @commands.command(hidden=True) + @commands.is_owner() + async def unload(self, ctx, *exts: str): + """Unloads an extension.""" + for ext in exts: + self.bot.unload_extension('bot.ext.' + ext) + log.info(f'Unloaded {ext}') + m = ctx.send(f':ok_hand: `{ext}` unloaded.') + self.bot.loop.create_task(m) + + @commands.command(hidden=True) + @commands.is_owner() + async def reload(self, ctx, *extensions: str): + """Reloads an extension""" + for ext in extensions: + try: + self.bot.unload_extension('bot.ext.' + ext) + self.bot.load_extension('bot.ext.' + ext) + except Exception as err: + await ctx.send(f'```{traceback.format_exc()}```') + return + log.info(f'Reloaded {ext}') + m = ctx.send(f':ok_hand: Reloaded `{ext}`') + self.bot.loop.create_task(m) + + + @commands.command(typing=True) + async def sql(self, ctx, *, statement: no_codeblock): + """Execute SQL.""" + # this is probably not the ideal solution + if 'select' in statement.lower(): + coro = self.db.fetch + else: + coro = self.db.execute + + try: + with Timer() as timer: + result = await coro(statement) + except asyncpg.PostgresError as e: + return await ctx.send(f':boom: Failed to execute! {type(e).__name__}: {e}') + + # execute returns the status as a string + if isinstance(result, str): + return await ctx.send(f'```py\n{result}```took {timer}') + + if not result: + return await ctx.send(f'no results, took {timer}') + + # render output of statement + columns = list(result[0].keys()) + table = Table(*columns) + + for row in result: + values = [str(x) for x in row] + table.add_row(*values) + + rendered = await table.render(self.loop) + + # properly emulate the psql console + rows = len(result) + rows = f'({rows} row{"s" if rows > 1 else ""})' + + await ctx.send(f'```py\n{rendered}\n{rows}```took {timer}') + + +def setup(bot): + bot.add_cog(Admin(bot)) diff --git a/bot/ext/basic.py b/bot/ext/basic.py new file mode 100644 index 0000000..d394f68 --- /dev/null +++ b/bot/ext/basic.py @@ -0,0 +1,53 @@ +import time +import discord + +from discord.ext import commands + +from .common import Cog + + +class Basic(Cog): + @commands.command(aliases=['p']) + async def ping(self, ctx): + """Ping. + """ + + t1 = time.monotonic() + m = await ctx.send('pinging...') + t2 = time.monotonic() + + rtt = (t2 - t1) * 1000 + gw = self.bot.latency * 1000 + + await m.edit(content=f'rtt: `{rtt:.1f}ms`, gw: `{gw:.1f}ms`') + + @commands.command() + async def uptime(self, ctx): + """Show uptime""" + sec = round(time.monotonic() - self.bot.init_time) + + m, s = divmod(sec, 60) + h, m = divmod(m, 60) + d, h = divmod(h, 24) + + await ctx.send(f'Uptime: **`{d} days, {h} hours, ' + f'{m} minutes, {s} seconds`**') + + @commands.command() + async def about(self, ctx): + """Show stuff.""" + + em = discord.Embed(title='Sex Dungeon Mistress') + em.add_field(name='About', value='SDM is a bot made to manage' + 'memework vps\' (sexhouse) operations to users') + + appinfo = await self.bot.application_info() + owner = appinfo.owner + em.add_field(name='Owner', + value=f'{owner.mention}, {owner}, (`{owner.id}`)') + + await ctx.send(embed=em) + + +def setup(bot): + bot.add_cog(Basic(bot)) diff --git a/bot/ext/common.py b/bot/ext/common.py new file mode 100644 index 0000000..c6fb2a5 --- /dev/null +++ b/bot/ext/common.py @@ -0,0 +1,9 @@ + +class Cog: + def __init__(self, bot): + self.bot = bot + self.loop = bot.loop + + @property + def db(self): + return self.bot.db diff --git a/bot/ext/utils/__init__.py b/bot/ext/utils/__init__.py new file mode 100644 index 0000000..2b39b32 --- /dev/null +++ b/bot/ext/utils/__init__.py @@ -0,0 +1,2 @@ + +from .mousey import * diff --git a/bot/ext/utils/mousey.py b/bot/ext/utils/mousey.py new file mode 100644 index 0000000..4d4d42a --- /dev/null +++ b/bot/ext/utils/mousey.py @@ -0,0 +1,97 @@ +import time +import asyncio +import functools + +from typing import List + + +class Timer: + """Context manager to measure how long the indented block takes to run.""" + + def __init__(self): + self.start = None + self.end = None + + def __enter__(self): + self.start = time.perf_counter() + return self + + async def __aenter__(self): + return self.__enter__() + + def __exit__(self, exc_type, exc_val, exc_tb): + self.end = time.perf_counter() + + async def __aexit__(self, exc_type, exc_val, exc_tb): + return self.__exit__(exc_type, exc_val, exc_tb) + + def __str__(self): + return f'{self.duration:.3f}ms' + + @property + def duration(self): + """Duration in ms.""" + return (self.end - self.start) * 1000 + + +class Table: + def __init__(self, *column_titles: str): + self._rows = [column_titles] + self._widths = [] + + for index, entry in enumerate(column_titles): + self._widths.append(len(entry)) + + def _update_widths(self, row: tuple): + for index, entry in enumerate(row): + width = len(entry) + if width > self._widths[index]: + self._widths[index] = width + + def add_row(self, *row: str): + """ + Add a row to the table. + .. note :: There's no check for the number of items entered, this may cause issues rendering if not correct. + """ + self._rows.append(row) + self._update_widths(row) + + def add_rows(self, *rows: List[str]): + for row in rows: + self.add_row(*row) + + def _render(self): + def draw_row(row_): + columns = [] + + for index, field in enumerate(row_): + # digits get aligned to the right + if field.isdigit(): + columns.append(f" {field:>{self._widths[index]}} ") + continue + + # regular text gets aligned to the left + columns.append(f" {field:<{self._widths[index]}} ") + + return "|".join(columns) + + # column title is centered in the middle of each field + title_row = "|".join(f" {field:^{self._widths[index]}} " for index, field in enumerate(self._rows[0])) + separator_row = "+".join("-" * (width + 2) for width in self._widths) + + drawn = [title_row, separator_row] + # remove the title row from the rows + self._rows = self._rows[1:] + + for row in self._rows: + row = draw_row(row) + drawn.append(row) + + return "\n".join(drawn) + + async def render(self, loop: asyncio.AbstractEventLoop=None): + """Returns a rendered version of the table.""" + loop = loop or asyncio.get_event_loop() + + func = functools.partial(self._render) + return await loop.run_in_executor(None, func) diff --git a/config.py b/config.py index 6c3c085..f2f156e 100644 --- a/config.py +++ b/config.py @@ -4,3 +4,6 @@ db = { 'database': 'memed', 'host': 'localhost' } + +bot_token = 'Mzg2Mjc1MDc3MzY2NDgwODk4.DQNi9A.BnyE5MnKaIaVMBbWiW9rVDwkrSs' +owner_id = 162819866682851329 diff --git a/memed.py b/memed.py index dd7b608..989a8f9 100644 --- a/memed.py +++ b/memed.py @@ -12,6 +12,7 @@ import logging import asyncpg import config +from bot import schedule_bot logging.basicConfig(level=logging.DEBUG) log = logging.getLogger(__name__) @@ -64,6 +65,7 @@ async def send_msg(writer, op: int, data: str): await writer.drain() async def process(op: int, message: str): + """Process messages given through the socket""" if op == 1: uid, cwd, command = parse_logstr(message) @@ -71,7 +73,8 @@ async def process(op: int, message: str): INSERT INTO logs (uid, cwd, cmd) VALUES ($1, $2, $3) """, uid, cwd, command) -async def handle_echo(reader, writer): +async def handle_client(reader, writer): + """Handle clients""" try: await send_msg(writer, 0, 'hello') @@ -92,12 +95,14 @@ async def handle_echo(reader, writer): if __name__ == '__main__': loop = asyncio.get_event_loop() - coro = asyncio.start_unix_server(handle_echo, './log.suck', + coro = asyncio.start_unix_server(handle_client, './log.suck', loop=loop) db = loop.run_until_complete(asyncpg.create_pool(**config.db)) server = loop.run_until_complete(coro) + schedule_bot(loop, config, db) + log.info(f'Serving on {server.sockets[0].getsockname()}') try: loop.run_forever() diff --git a/requirements.txt b/requirements.txt index ad26ec1..3d5b879 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1 +1,3 @@ +git+https://github.com/Rapptz/discord.py@rewrite asyncpg==0.13.0 +