shit
This commit is contained in:
Luna Mendes
parent
a0352ef1b1
commit
3412c5c0b9
3 changed files with 106 additions and 62 deletions
|
|
@ -76,7 +76,7 @@ class Rsudo(Cog):
|
|||
|
||||
return users[0]
|
||||
|
||||
async def create_request(self, message):
|
||||
async def request(self, message, wait=False):
|
||||
command_channel = self.bot.get_channel(
|
||||
self.bot.config.command_channel
|
||||
)
|
||||
|
|
@ -142,7 +142,10 @@ class Rsudo(Cog):
|
|||
await msg.edit(embed=embed)
|
||||
|
||||
# todo: does this not execute in the wrong directory?
|
||||
out = await shell(command)
|
||||
if wait:
|
||||
return True
|
||||
else:
|
||||
out = await shell(command)
|
||||
|
||||
# output may be very long
|
||||
paginator = commands.Paginator(prefix='```sh\n', suffix='```')
|
||||
|
|
@ -168,6 +171,9 @@ class Rsudo(Cog):
|
|||
|
||||
await msg.edit(embed=embed)
|
||||
|
||||
if wait:
|
||||
return False
|
||||
|
||||
try:
|
||||
# await user.send(f'Your request for the command {command} was denied.')
|
||||
pass
|
||||
|
|
|
|||
18
config.py
18
config.py
|
|
@ -1,3 +1,7 @@
|
|||
env = 'test'
|
||||
|
||||
owner_id = 162819866682851329
|
||||
|
||||
db = {
|
||||
'user': 'abc',
|
||||
'password': 'def',
|
||||
|
|
@ -5,10 +9,14 @@ db = {
|
|||
'host': 'localhost'
|
||||
}
|
||||
|
||||
bot_token = 'Mzg2Mjc1MDc3MzY2NDgwODk4.DQNi9A.BnyE5MnKaIaVMBbWiW9rVDwkrSs'
|
||||
owner_id = 162819866682851329
|
||||
if env == 'test':
|
||||
# test
|
||||
bot_token = 'MjExODg5MDEyNzQ5NzYyNTYw.DQttvA.dRkggolnj7csbzg3Y28d45KbAf4'
|
||||
else:
|
||||
# prod
|
||||
bot_token = 'Mzg2Mjc1MDc3MzY2NDgwODk4.DQNi9A.BnyE5MnKaIaVMBbWiW9rVDwkrSs'
|
||||
|
||||
command_channel = 386989864014446612
|
||||
sexr_chan = 386990164896907264
|
||||
command_channel = 212201893391237121
|
||||
sexr_chan = 212201893391237121
|
||||
|
||||
admin_role = 386566510803681280
|
||||
admin_role = 349563478530326529
|
||||
|
|
|
|||
140
memed.py
140
memed.py
|
|
@ -15,7 +15,7 @@ import asyncpg
|
|||
import config
|
||||
from bot import schedule_bot
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logging.basicConfig(level=logging.DEBUG)
|
||||
log = logging.getLogger(__name__)
|
||||
db = None
|
||||
bot = None
|
||||
|
|
@ -27,6 +27,7 @@ bot = None
|
|||
# for reader:
|
||||
# - op 1 : log
|
||||
# - op 2 : rsudo
|
||||
# - op 3 : rsudo with steroids
|
||||
|
||||
|
||||
async def wrap(coro):
|
||||
|
|
@ -62,77 +63,106 @@ def parse_logstr(string):
|
|||
return int(uid), cwd, command
|
||||
|
||||
|
||||
async def read_msg(reader):
|
||||
header = await reader.read(8)
|
||||
log.debug('[recv] %r', header)
|
||||
length, op = struct.unpack('Ii', header)
|
||||
data = await reader.read(length)
|
||||
data = data.decode()
|
||||
class MemeClient:
|
||||
"""MemeD client handler."""
|
||||
def __init__(self, reader, writer):
|
||||
self.reader = reader
|
||||
self.writer = writer
|
||||
|
||||
log.debug('[recv] %d %d %s', length, op, data)
|
||||
return op, data
|
||||
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(reader):
|
||||
op, message = await read_msg(reader)
|
||||
if op > 10:
|
||||
return op, json.loads(message)
|
||||
else:
|
||||
return op, message
|
||||
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.
|
||||
|
||||
async def send_msg(writer, op: int, data: str):
|
||||
header = struct.pack('Ii', len(data), op).decode()
|
||||
msg = f'{header}{data}'.encode()
|
||||
log.debug('[send] %d, %s -> %r', op, data, msg)
|
||||
This does not wait for the receiving end
|
||||
to properly finish their buffers.
|
||||
|
||||
writer.write(msg)
|
||||
asyncio.get_event_loop().create_task(wrap(writer.drain))
|
||||
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)
|
||||
|
||||
async def process(reader, writer, op: int, message: str):
|
||||
"""Process messages given through the socket"""
|
||||
if op == 1:
|
||||
uid, cwd, command = parse_logstr(message)
|
||||
# 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))
|
||||
|
||||
log.info('[process] Logging command uid=%d cwd=%r cmd=%r',
|
||||
uid, cwd, command)
|
||||
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:
|
||||
if not bot:
|
||||
return await send_msg(writer, 1, 'no bot up')
|
||||
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 send_msg(writer, 1, 'no rsudo cog')
|
||||
rsudo = bot.get_cog('Rsudo')
|
||||
if not rsudo:
|
||||
return await self.send_msg(1, 'no rsudo cog')
|
||||
|
||||
log.info('[process] got rsudo! %r', message)
|
||||
await rsudo.create_request(message)
|
||||
return await send_msg(writer, 1, 'ok')
|
||||
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 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"""
|
||||
try:
|
||||
asyncio.get_event_loop().create_task(send_msg(writer, 0, 'hello'))
|
||||
client = MemeClient(reader, writer)
|
||||
|
||||
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 ConnectionError as e:
|
||||
log.warning('connection err: %r', e)
|
||||
except Exception as e:
|
||||
log.exception('error at handler coro')
|
||||
writer.close()
|
||||
await client.send_msg(0, 'hello')
|
||||
await client.loop()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
|
|
|||
Loading…
Reference in a new issue