memed/memed.py

"""
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()