From 3412c5c0b90e80460c8747307eb407eb8ce04923 Mon Sep 17 00:00:00 2001 From: Luna Mendes Date: Thu, 7 Dec 2017 21:54:46 -0300 Subject: [PATCH] shit --- bot/ext/rsudo.py | 10 +++- config.py | 18 ++++-- memed.py | 140 ++++++++++++++++++++++++++++------------------- 3 files changed, 106 insertions(+), 62 deletions(-) diff --git a/bot/ext/rsudo.py b/bot/ext/rsudo.py index 73d81e9..b1f040c 100644 --- a/bot/ext/rsudo.py +++ b/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 diff --git a/config.py b/config.py index 608e313..dcd125d 100644 --- a/config.py +++ b/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 diff --git a/memed.py b/memed.py index 957eea7..ab0a703 100644 --- a/memed.py +++ b/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__':