add the bot

Makefile

@@ -0,0 +1,4 @@
+OBFUSCATE="./env/bin/pyminifier --obfuscate --gzip"
+
+all:
+	${OBFUSCATE} -o build/memed.py memed.py

bot/__init__.py

@@ -0,0 +1 @@
+from .bot import *

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)

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))

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))

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

bot/ext/utils/__init__.py

@@ -0,0 +1,2 @@
+
+from .mousey import *

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)

config.py

@@ -4,3 +4,6 @@ db = {
     'database': 'memed',
     'host': 'localhost'
 }
+
+bot_token = 'Mzg2Mjc1MDc3MzY2NDgwODk4.DQNi9A.BnyE5MnKaIaVMBbWiW9rVDwkrSs'
+owner_id = 162819866682851329

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()

requirements.txt

@@ -1 +1,3 @@
+git+https://github.com/Rapptz/discord.py@rewrite
 asyncpg==0.13.0
+