""" MemeD - the MEMEwork Daemon - This program manages the memework vps' command logging """ import asyncio import struct import json import logging import sys import asyncpg import config from bot import schedule_bot 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): # '2015-02-11T19:05:10+00:00 labrat-1 snoopy[896]: [uid:0 sid:11679 # tty:/dev/pts/2 cwd:/root filename:/usr/bin/cat]: cat /etc/fstab.BAK' # I really need to parse the uid, cwd and the command out of that. splitted = string.split(':') command = splitted[-1].strip() k = string.find('[') important = string[string.find('[', k + 1):] lst = important.replace('[', '').replace(']', '').split() # filder uid and cwd s = [s.split(':') for s in lst if 'uid' in s or 'cwd' in s] uid = [e[1] for e in s if e[0] == 'uid'][0] 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) data = await reader.read(length) data = data.decode() log.info('[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 send_msg(writer, op: int, data: str): header = struct.pack('Ii', len(data), op) msg = f'{header.decode()}{data}'.encode() log.info('[send] %d, %s -> %r', op, data, msg) writer.write(msg) asyncio.get_event_loop.create_task(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) await db.execute(""" INSERT INTO logs (uid, cwd, cmd) VALUES ($1, $2, $3) """, uid, cwd, command) elif op == 2: log.debug(bot.cogs) 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""" try: asyncio.get_event_loop().create_task(send_msg(writer, 0, 'hello')) 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 Exception as e: log.exception('error at handler coro') await send_msg(writer, -1, repr(e)) writer.close() if __name__ == '__main__': loop = asyncio.get_event_loop() coro = asyncio.start_unix_server(handle_client, sys.argv[1], loop=loop) db = loop.run_until_complete(asyncpg.create_pool(**config.db)) server = loop.run_until_complete(coro) if config.bot_token: bot = schedule_bot(loop, config, db) if bot: loop.create_task(bot.start(config.bot_token)) log.info(f'Serving on {server.sockets[0].getsockname()}') try: loop.run_forever() except KeyboardInterrupt: pass log.info('Closing server') server.close() loop.run_until_complete(server.wait_closed()) loop.close()