memed/memed.py
2018-02-24 22:31:57 -03:00

240 lines
6.5 KiB
Python

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