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__)
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 e:
        log.warning('connection err: %r', e)
    except:
        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
    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 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:
        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(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.

        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)

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

            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')

            rsudo = bot.get_cog('Rsudo')
            if not rsudo:
                return await self.send_msg(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')

            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 client_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"""
    client = MemeClient(reader, writer)

    await client.send_msg(0, 'hello')
    await client.client_loop()


if __name__ == '__main__':
    loop = asyncio.get_event_loop()
    coro = asyncio.start_unix_server(handle_client, sys.argv[1],
                                     loop=loop)

    db = loop.create_task(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()