""" 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__) # globals are bad, but who cares db = None bot = None # for writer: # - op 0 : hello # - op 1 : request response # for reader: # - op 1 : log # - op 2 : rsudo # - op 3 : rsudo with steroids async def wrap(coro): try: await coro() except ConnectionError as err: log.warning(f'connection err: {err!r}') except Exception: log.exception('error inside wrapped') 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. # '[uid:123 sid:440 tty:/dev/pts/4 cwd:/root filename:/bin/chmod]: AAAA BBBBBB CCCCCCCC' # THIS IS WHAT WE PARSE NOW. splitted = string.split(':') command = splitted[-1].strip() k = string.find('[') important = string[k:] lst = important.replace('[', '').replace(']', '').split() # filder uid and cwd spl = [s.split(':') for s in lst if 'uid' in s or 'cwd' in s] uid = [e[1] for e in spl if e[0] == 'uid'][0] cwd = [e[1] for e in spl if e[0] == 'cwd'][0] return int(uid), cwd, command class MemeClient: """MemeD client handler.""" def __init__(self, reader, writer): self.reader = reader self.writer = writer self.loop = asyncio.get_event_loop() async def read_msg(self) -> str: """Read one message from the socket.""" header = await self.reader.read(8) log.debug('[recv] %r', header) length, opcode = struct.unpack('Ii', header) data = await self.reader.read(length) data = data.decode() log.debug('[recv] %d %d %s', length, opcode, data) return opcode, data async def read_payload(self) -> dict: """Read a payload from the socket.""" opcode, message = await self.read_msg() # NOTE: this is kinda unused if opcode > 10: return opcode, json.loads(message) return opcode, message async def send_msg(self, op: int, data: str): """Send a message. This does not wait for the receiving end to properly flush their buffers. Arguments --------- op: int OP code to be sent. data: str Message to be sent with the op code. """ # create header, pack message, yadda yadda header = struct.pack('Ii', len(data), op).decode() msg = f'{header}{data}'.encode() log.debug('[send] %d, %s -> %r', op, data, msg) self.writer.write(msg) # clients 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(self, op: int, message: str) -> 'None': """Process a message given through the socket""" handler = getattr(self, f'handle_{op}', None) if not handler: # Ignore unknown OP codes. return await handler(message) async def handle_1(self, message: str): global db uid, cwd, command = parse_logstr(message) log.info('[process] Logging command ' f'uid={uid} cwd={cwd} cmd={command}') await db.execute(""" INSERT INTO logs (uid, cwd, cmd) VALUES ($1, $2, $3) """, uid, cwd, command) async def handle_2(self, message: str): """Handle an OP 2 packet. This is the RSudo handling, but without waiting for the approval of a mod (the 'nowait' behavior). """ if not bot: return await self.send_msg(1, 'no bot up') rsudo = bot.get_cog('Rsudo') if not rsudo: return await self.send_msg(1, 'no rsudo cog') log.info('[rsudo:nowait] %r', message) # this doesnt wait for the thing self.loop.create_task(rsudo.request(message)) return await self.send_msg(1, "true") async def handle_3(self, message: str): """Handle an OP 3 packet. This follows the same logic as OP 2, however the client gets a response back, a string ('true' or 'false'), depending on the request. NO COMMANDS WILL BE EXECUTED SERVER-SIDE FROM THIS OP. """ if not bot: return await self.send_msg(1, 'no bot started') rsudo = bot.get_cog('Rsudo') if not rsudo: return await self.send_msg(1, 'no rsudo cog') log.info('[rsudo:wait] %r', message) success = await rsudo.request(message, True) # the original idea was sending an int back # but we had no idea how to do that since # we were already working with strings all the time # so yeah, this is *another* string ok_str = 'true' if success else 'false' return await self.send_msg(1, ok_str) async def client_loop(self): """Enter a loop waiting for messages from the client.""" 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 coming in the socket, spawn a loop for them.""" client = MemeClient(reader, writer) await client.send_msg(0, 'hello') await client.client_loop() def main(): # YES GLOBALS ARE BAD I KNOW global bot global db loop = asyncio.get_event_loop() coro = asyncio.start_unix_server(handle_client, sys.argv[1], loop=loop) pool = loop.run_until_complete(asyncpg.create_pool(**config.db)) db = pool server = loop.run_until_complete(coro) if config.bot_token: bot = schedule_bot(loop, config, pool) if bot: loop.create_task(bot.start(config.bot_token)) log.info(f'memed serving at {server.sockets[0].getsockname()}') try: loop.run_forever() finally: log.info('Closing server') server.close() loop.run_until_complete(server.wait_closed()) loop.close() if __name__ == '__main__': main()