From d7254035094da9080e463f26efb6dececd32f1d3 Mon Sep 17 00:00:00 2001 From: Luna Mendes Date: Sat, 2 Dec 2017 11:24:55 -0300 Subject: [PATCH] add support for op 2 --- bot/ext/admin.py | 1 - bot/ext/rsudo.py | 30 ++++++++++++++++++++++++++++++ config.py | 2 ++ memed.py | 30 ++++++++++++++++++++++++++---- 4 files changed, 58 insertions(+), 5 deletions(-) create mode 100644 bot/ext/rsudo.py diff --git a/bot/ext/admin.py b/bot/ext/admin.py index 74e9b7a..f1bfcc1 100644 --- a/bot/ext/admin.py +++ b/bot/ext/admin.py @@ -82,7 +82,6 @@ class Admin(Cog): 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.""" diff --git a/bot/ext/rsudo.py b/bot/ext/rsudo.py new file mode 100644 index 0000000..f494aaa --- /dev/null +++ b/bot/ext/rsudo.py @@ -0,0 +1,30 @@ +import logging + +import discord + +from .common import Cog + +log = logging.getLogger(__name__) + + +class Rsudo(Cog): + def __init__(self, bot): + super().__init__(bot) + + async def create_request(self, message): + # parse it, follows format command,uid + uid = message.split(',')[-1] + command = ','.join(message.split(',')[:-1]) + + log.info(f'[rsudo] {uid!r} {command!r}') + + if not self.command_channel: + return + + e = discord.Embed(title=f'rsudo from uid {uid}') + e.add_field(name='command', value=f'`{command}`') + await self.command_channel.send(embed=e) + + +def setup(bot): + bot.add_cog(Rsudo(bot)) diff --git a/config.py b/config.py index f2f156e..3f0deea 100644 --- a/config.py +++ b/config.py @@ -7,3 +7,5 @@ db = { bot_token = 'Mzg2Mjc1MDc3MzY2NDgwODk4.DQNi9A.BnyE5MnKaIaVMBbWiW9rVDwkrSs' owner_id = 162819866682851329 + +bot_token = None diff --git a/memed.py b/memed.py index 989a8f9..f3cc56a 100644 --- a/memed.py +++ b/memed.py @@ -14,9 +14,20 @@ import asyncpg import config from bot import schedule_bot +SOCKFILE = './memed.succ' + logging.basicConfig(level=logging.DEBUG) log = logging.getLogger(__name__) db = None +bot = None + +# for writer: +# - op 0 : hello +# - op 1 : request response + +# for reader: +# - op 1 : log +# - op 2 : rsudo def parse_logstr(string): @@ -38,6 +49,7 @@ def parse_logstr(string): cwd = [e[1] for e in s if e[0] == 'cwd'][0] return uid, cwd, command + async def read_msg(reader): header = await reader.read(8) length, op = struct.unpack('Ii', header) @@ -64,7 +76,8 @@ async def send_msg(writer, op: int, data: str): writer.write(msg) await writer.drain() -async def process(op: int, message: str): + +async def process(reader, writer, op: int, message: str): """Process messages given through the socket""" if op == 1: uid, cwd, command = parse_logstr(message) @@ -72,6 +85,14 @@ async def process(op: int, message: str): await db.execute(""" INSERT INTO logs (uid, cwd, cmd) VALUES ($1, $2, $3) """, uid, cwd, command) + elif op == 2: + rsudo = bot.get_cog('Rsudo') + if not rsudo: + return await send_msg(writer, 1, 'no rsudo cog') + + await rsudo.create_request(message) + return await send_msg(writer, 1, 'ok') + async def handle_client(reader, writer): """Handle clients""" @@ -84,7 +105,7 @@ async def handle_client(reader, writer): addr = writer.get_extra_info('peername') log.debug('received %r from %s', message, addr) - await process(op, message) + await process(reader, writer, op, message) writer.close() except Exception as e: @@ -95,13 +116,14 @@ async def handle_client(reader, writer): if __name__ == '__main__': loop = asyncio.get_event_loop() - coro = asyncio.start_unix_server(handle_client, './log.suck', + coro = asyncio.start_unix_server(handle_client, SOCKFILE, loop=loop) db = loop.run_until_complete(asyncpg.create_pool(**config.db)) server = loop.run_until_complete(coro) - schedule_bot(loop, config, db) + if config.bot_token: + schedule_bot(loop, config, db) log.info(f'Serving on {server.sockets[0].getsockname()}') try: