shit

bot/ext/rsudo.py

@@ -76,7 +76,7 @@ class Rsudo(Cog):
 
         return users[0]
 
-    async def create_request(self, message):
+    async def request(self, message, wait=False):
         command_channel = self.bot.get_channel(
             self.bot.config.command_channel
         )
@@ -142,7 +142,10 @@ class Rsudo(Cog):
                 await msg.edit(embed=embed)
 
                 # todo: does this not execute in the wrong directory?
-                out = await shell(command)
+                if wait:
+                    return True
+                else:
+                    out = await shell(command)
 
             # output may be very long
             paginator = commands.Paginator(prefix='```sh\n', suffix='```')
@@ -168,6 +171,9 @@ class Rsudo(Cog):
 
             await msg.edit(embed=embed)
 
+            if wait:
+                return False
+
             try:
                 # await user.send(f'Your request for the command {command} was denied.')
                 pass

config.py

@@ -1,3 +1,7 @@
+env = 'test'
+
+owner_id = 162819866682851329
+
 db = {
     'user': 'abc',
     'password': 'def',
@@ -5,10 +9,14 @@ db = {
     'host': 'localhost'
 }
 
-bot_token = 'Mzg2Mjc1MDc3MzY2NDgwODk4.DQNi9A.BnyE5MnKaIaVMBbWiW9rVDwkrSs'
-owner_id = 162819866682851329
+if env == 'test':
+    # test
+    bot_token = 'MjExODg5MDEyNzQ5NzYyNTYw.DQttvA.dRkggolnj7csbzg3Y28d45KbAf4'
+else:
+    # prod
+    bot_token = 'Mzg2Mjc1MDc3MzY2NDgwODk4.DQNi9A.BnyE5MnKaIaVMBbWiW9rVDwkrSs'
 
-command_channel = 386989864014446612
-sexr_chan = 386990164896907264
+    command_channel = 212201893391237121
+    sexr_chan = 212201893391237121
 
-admin_role = 386566510803681280
+    admin_role = 349563478530326529

memed.py

@@ -15,7 +15,7 @@ import asyncpg
 import config
 from bot import schedule_bot
 
-logging.basicConfig(level=logging.INFO)
+logging.basicConfig(level=logging.DEBUG)
 log = logging.getLogger(__name__)
 db = None
 bot = None
@@ -27,6 +27,7 @@ bot = None
 # for reader:
 # - op 1 : log
 # - op 2 : rsudo
+# - op 3 : rsudo with steroids
 
 
 async def wrap(coro):
@@ -62,77 +63,106 @@ def parse_logstr(string):
     return int(uid), cwd, command
 
 
-async def read_msg(reader):
-    header = await reader.read(8)
-    log.debug('[recv] %r', header)
-    length, op = struct.unpack('Ii', header)
-    data = await reader.read(length)
-    data = data.decode()
+class MemeClient:
+    """MemeD client handler."""
+    def __init__(self, reader, writer):
+        self.reader = reader
+        self.writer = writer
 
-    log.debug('[recv] %d %d %s', length, op, data)
-    return op, data
+    async def read_msg(self) -> str:
+        header = await self.reader.read(8)
+        log.debug('[recv] %r', header)
+        length, op = struct.unpack('Ii', header)
+        data = await self.reader.read(length)
+        data = data.decode()
 
+        log.debug('[recv] %d %d %s', length, op, data)
+        return op, data
 
-async def read_payload(reader):
-    op, message = await read_msg(reader)
-    if op > 10:
-        return op, json.loads(message)
-    else:
-        return op, message
+    async def read_payload(self) -> dict:
+        op, message = await self.read_msg()
+        if op > 10:
+            return op, json.loads(message)
+        else:
+            return op, message
 
+    async def send_msg(self, op: int, data: str) -> 'None':
+        """Send a message.
+        
+        This does not wait for the receiving end
+        to properly finish their buffers.
 
-async def send_msg(writer, op: int, data: str):
-    header = struct.pack('Ii', len(data), op).decode()
-    msg = f'{header}{data}'.encode()
-    log.debug('[send] %d, %s -> %r', op, data, msg)
+        Arguments
+        ---------
+        op: int
+            OP code to be sent.
+        data: str
+            Message to be sent with the op code.
+        """
+        header = struct.pack('Ii', len(data), op).decode()
+        msg = f'{header}{data}'.encode()
+        log.debug('[send] %d, %s -> %r', op, data, msg)
 
-    writer.write(msg)
-    asyncio.get_event_loop().create_task(wrap(writer.drain))
+        self.writer.write(msg)
 
+        # Utils can close this early
+        # and make writer.drain kill itself
+        # so we wrap on a task which is isolated
+        asyncio.get_event_loop().create_task(wrap(self.writer.drain))
 
-async def process(reader, writer, op: int, message: str):
-    """Process messages given through the socket"""
-    if op == 1:
-        uid, cwd, command = parse_logstr(message)
+    async def process(self, op: int, message: str) -> 'None':
+        """Process a message given through the socket"""
+        if op == 1:
+            uid, cwd, command = parse_logstr(message)
+            log.info('[process] Logging command uid=%d cwd=%r cmd=%r', 
+                    uid, cwd, command)
 
-        log.info('[process] Logging command uid=%d cwd=%r cmd=%r', 
-                uid, cwd, command)
+            await db.execute("""
+            INSERT INTO logs (uid, cwd, cmd) VALUES ($1, $2, $3)
+            """, uid, cwd, command)
+        elif op == 2:
+            # Handle rsudo without waiting
+            if not bot:
+                return await self.send_msg(1, 'no bot up')
 
-        await db.execute("""
-        INSERT INTO logs (uid, cwd, cmd) VALUES ($1, $2, $3)
-        """, uid, cwd, command)
-    elif op == 2:
-        if not bot:
-            return await send_msg(writer, 1, 'no bot up')
+            rsudo = bot.get_cog('Rsudo')
+            if not rsudo:
+                return await self.send_msg(1, 'no rsudo cog')
 
-        rsudo = bot.get_cog('Rsudo')
-        if not rsudo:
-            return await send_msg(writer, 1, 'no rsudo cog')
+            log.info('[process] got rsudo! %r', message)
+            self.loop.create_task(rsudo.request(message))
+            return await self.send_msg(1, int(True))
+        elif op == 3:
+            # handle rsudo, waitinG
+            if not bot:
+                return await self.send_msg(1, 'no bot')
 
-        log.info('[process] got rsudo! %r', message)
-        await rsudo.create_request(message)
-        return await send_msg(writer, 1, 'ok')
+            rsudo = bot.get_cog('Rsudo')
+            if not rsudo:
+                return await self.send_msg(1, 'no rsudo cog')
+
+            log.info('[process - wait] %r', message)
+            ok = await rsudo.request(message, True)
+            return await self.send_msg(1, int(ok))
+
+    async def loop(self):
+        try:
+            while True:
+                op, message = await self.read_msg()
+                await self.process(op, message)
+        except ConnectionError as e:
+            log.warning('conn err: %r', e)
+        except Exception:
+            log.exception('error at loop')
+            self.writer.close()
 
 
 async def handle_client(reader, writer):
     """Handle clients"""
-    try:
-        asyncio.get_event_loop().create_task(send_msg(writer, 0, 'hello'))
+    client = MemeClient(reader, writer)
 
-        while True:
-            op, message = await read_msg(reader)
-
-            addr = writer.get_extra_info('peername')
-            log.debug('received %r from %s', message, addr)
-
-            await process(reader, writer, op, message)
-
-        writer.close()
-    except ConnectionError as e:
-        log.warning('connection err: %r', e)
-    except Exception as e:
-        log.exception('error at handler coro')
-        writer.close()
+    await client.send_msg(0, 'hello')
+    await client.loop()
 
 
 if __name__ == '__main__':