141 lines
3.4 KiB
Python
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 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
|
|
|
|
|
|
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, sys.argv[1],
|
|
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()
|