memed/memed.py
2017-12-02 12:53:47 -03:00

141 lines
3.4 KiB
Python

"""
MemeD - the MEMEwork Daemon
- This program manages the memework vps'
command logging
"""
import asyncio
import struct
import json
import logging
import asyncpg
import config
from bot import schedule_bot
SOCKFILE = './memed.succ'
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)
await 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:
await 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, SOCKFILE,
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()