add the bot
This commit is contained in:
parent
3779ed4279
commit
d40b8dfcaa
12 changed files with 360 additions and 2 deletions
4
Makefile
Normal file
4
Makefile
Normal file
|
@ -0,0 +1,4 @@
|
|||
OBFUSCATE="./env/bin/pyminifier --obfuscate --gzip"
|
||||
|
||||
all:
|
||||
${OBFUSCATE} -o build/memed.py memed.py
|
0
__init__.py
Normal file
0
__init__.py
Normal file
1
bot/__init__.py
Normal file
1
bot/__init__.py
Normal file
|
@ -0,0 +1 @@
|
|||
from .bot import *
|
56
bot/bot.py
Normal file
56
bot/bot.py
Normal file
|
@ -0,0 +1,56 @@
|
|||
import logging
|
||||
import time
|
||||
|
||||
import discord
|
||||
from discord.ext import commands
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
exts = [
|
||||
'basic',
|
||||
'admin'
|
||||
]
|
||||
|
||||
|
||||
class SexDungeonMistressBot(commands.Bot):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.init_time = time.monotonic()
|
||||
|
||||
async def on_command(self, ctx):
|
||||
# thanks dogbot ur a good
|
||||
content = ctx.message.content
|
||||
|
||||
author = ctx.message.author
|
||||
guild = ctx.guild
|
||||
checks = [c.__qualname__.split('.')[0] for c in ctx.command.checks]
|
||||
location = '[DM]' if isinstance(ctx.channel, discord.DMChannel) else \
|
||||
f'[Guild {guild.name} {guild.id}]'
|
||||
|
||||
log.info('%s [cmd] %s(%d) "%s" checks=%s', location, author,
|
||||
author.id, content, ','.join(checks) or '(none)')
|
||||
|
||||
|
||||
def schedule_bot(loop, config, db):
|
||||
mute = ['discord', 'websockets']
|
||||
|
||||
for l in mute:
|
||||
d = logging.getLogger(l)
|
||||
d.setLevel(logging.INFO)
|
||||
|
||||
bot = SexDungeonMistressBot(
|
||||
command_prefix='sex ',
|
||||
description='sexhouse management bot',
|
||||
owner_id=getattr(config, 'owner_id', None),
|
||||
)
|
||||
|
||||
bot.db = db
|
||||
|
||||
try:
|
||||
for ext in exts:
|
||||
bot.load_extension(f'bot.ext.{ext}')
|
||||
log.info('loaded %s', ext)
|
||||
|
||||
loop.create_task(bot.start(config.bot_token))
|
||||
except:
|
||||
log.exception('failed to load %s', ext)
|
126
bot/ext/admin.py
Normal file
126
bot/ext/admin.py
Normal file
|
@ -0,0 +1,126 @@
|
|||
import logging
|
||||
import traceback
|
||||
|
||||
import asyncpg
|
||||
|
||||
from discord.ext import commands
|
||||
|
||||
from .common import Cog
|
||||
from .utils import Table, Timer
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def no_codeblock(text: str) -> str:
|
||||
"""
|
||||
Removes codeblocks (grave accents), python and sql syntax highlight
|
||||
indicators from a text if present.
|
||||
.. note:: only the start of a string is checked, the text is allowed
|
||||
to have grave accents in the middle
|
||||
"""
|
||||
if text.startswith('```'):
|
||||
text = text[3:-3]
|
||||
|
||||
if text.startswith(('py', 'sql')):
|
||||
# cut off the first line as this removes the
|
||||
# highlight indicator regardless of length
|
||||
text = '\n'.join(text.split('\n')[1:])
|
||||
|
||||
if text.startswith('`'):
|
||||
text = text[1:-1]
|
||||
|
||||
return text
|
||||
|
||||
|
||||
class Admin(Cog):
|
||||
async def __local_check(self, ctx):
|
||||
return ctx.author.id == self.bot.owner_id
|
||||
|
||||
@commands.command()
|
||||
@commands.is_owner()
|
||||
async def shutdown(self, ctx):
|
||||
log.info('Logging out! %s', ctx.author)
|
||||
await ctx.send("dude rip")
|
||||
await self.bot.logout()
|
||||
|
||||
@commands.command()
|
||||
@commands.is_owner()
|
||||
async def load(self, ctx, *exts: str):
|
||||
"""Loads an extension."""
|
||||
for ext in exts:
|
||||
try:
|
||||
self.bot.load_extension("bot.ext." + ext)
|
||||
except Exception as e:
|
||||
await ctx.send(f'Oops. ```py\n{traceback.format_exc()}\n```')
|
||||
return
|
||||
log.info(f'Loaded {ext}')
|
||||
m = ctx.send(f':ok_hand: `{ext}` loaded.')
|
||||
self.bot.loop.create_task(m)
|
||||
|
||||
@commands.command(hidden=True)
|
||||
@commands.is_owner()
|
||||
async def unload(self, ctx, *exts: str):
|
||||
"""Unloads an extension."""
|
||||
for ext in exts:
|
||||
self.bot.unload_extension('bot.ext.' + ext)
|
||||
log.info(f'Unloaded {ext}')
|
||||
m = ctx.send(f':ok_hand: `{ext}` unloaded.')
|
||||
self.bot.loop.create_task(m)
|
||||
|
||||
@commands.command(hidden=True)
|
||||
@commands.is_owner()
|
||||
async def reload(self, ctx, *extensions: str):
|
||||
"""Reloads an extension"""
|
||||
for ext in extensions:
|
||||
try:
|
||||
self.bot.unload_extension('bot.ext.' + ext)
|
||||
self.bot.load_extension('bot.ext.' + ext)
|
||||
except Exception as err:
|
||||
await ctx.send(f'```{traceback.format_exc()}```')
|
||||
return
|
||||
log.info(f'Reloaded {ext}')
|
||||
m = ctx.send(f':ok_hand: Reloaded `{ext}`')
|
||||
self.bot.loop.create_task(m)
|
||||
|
||||
|
||||
@commands.command(typing=True)
|
||||
async def sql(self, ctx, *, statement: no_codeblock):
|
||||
"""Execute SQL."""
|
||||
# this is probably not the ideal solution
|
||||
if 'select' in statement.lower():
|
||||
coro = self.db.fetch
|
||||
else:
|
||||
coro = self.db.execute
|
||||
|
||||
try:
|
||||
with Timer() as timer:
|
||||
result = await coro(statement)
|
||||
except asyncpg.PostgresError as e:
|
||||
return await ctx.send(f':boom: Failed to execute! {type(e).__name__}: {e}')
|
||||
|
||||
# execute returns the status as a string
|
||||
if isinstance(result, str):
|
||||
return await ctx.send(f'```py\n{result}```took {timer}')
|
||||
|
||||
if not result:
|
||||
return await ctx.send(f'no results, took {timer}')
|
||||
|
||||
# render output of statement
|
||||
columns = list(result[0].keys())
|
||||
table = Table(*columns)
|
||||
|
||||
for row in result:
|
||||
values = [str(x) for x in row]
|
||||
table.add_row(*values)
|
||||
|
||||
rendered = await table.render(self.loop)
|
||||
|
||||
# properly emulate the psql console
|
||||
rows = len(result)
|
||||
rows = f'({rows} row{"s" if rows > 1 else ""})'
|
||||
|
||||
await ctx.send(f'```py\n{rendered}\n{rows}```took {timer}')
|
||||
|
||||
|
||||
def setup(bot):
|
||||
bot.add_cog(Admin(bot))
|
53
bot/ext/basic.py
Normal file
53
bot/ext/basic.py
Normal file
|
@ -0,0 +1,53 @@
|
|||
import time
|
||||
import discord
|
||||
|
||||
from discord.ext import commands
|
||||
|
||||
from .common import Cog
|
||||
|
||||
|
||||
class Basic(Cog):
|
||||
@commands.command(aliases=['p'])
|
||||
async def ping(self, ctx):
|
||||
"""Ping.
|
||||
"""
|
||||
|
||||
t1 = time.monotonic()
|
||||
m = await ctx.send('pinging...')
|
||||
t2 = time.monotonic()
|
||||
|
||||
rtt = (t2 - t1) * 1000
|
||||
gw = self.bot.latency * 1000
|
||||
|
||||
await m.edit(content=f'rtt: `{rtt:.1f}ms`, gw: `{gw:.1f}ms`')
|
||||
|
||||
@commands.command()
|
||||
async def uptime(self, ctx):
|
||||
"""Show uptime"""
|
||||
sec = round(time.monotonic() - self.bot.init_time)
|
||||
|
||||
m, s = divmod(sec, 60)
|
||||
h, m = divmod(m, 60)
|
||||
d, h = divmod(h, 24)
|
||||
|
||||
await ctx.send(f'Uptime: **`{d} days, {h} hours, '
|
||||
f'{m} minutes, {s} seconds`**')
|
||||
|
||||
@commands.command()
|
||||
async def about(self, ctx):
|
||||
"""Show stuff."""
|
||||
|
||||
em = discord.Embed(title='Sex Dungeon Mistress')
|
||||
em.add_field(name='About', value='SDM is a bot made to manage'
|
||||
'memework vps\' (sexhouse) operations to users')
|
||||
|
||||
appinfo = await self.bot.application_info()
|
||||
owner = appinfo.owner
|
||||
em.add_field(name='Owner',
|
||||
value=f'{owner.mention}, {owner}, (`{owner.id}`)')
|
||||
|
||||
await ctx.send(embed=em)
|
||||
|
||||
|
||||
def setup(bot):
|
||||
bot.add_cog(Basic(bot))
|
9
bot/ext/common.py
Normal file
9
bot/ext/common.py
Normal file
|
@ -0,0 +1,9 @@
|
|||
|
||||
class Cog:
|
||||
def __init__(self, bot):
|
||||
self.bot = bot
|
||||
self.loop = bot.loop
|
||||
|
||||
@property
|
||||
def db(self):
|
||||
return self.bot.db
|
2
bot/ext/utils/__init__.py
Normal file
2
bot/ext/utils/__init__.py
Normal file
|
@ -0,0 +1,2 @@
|
|||
|
||||
from .mousey import *
|
97
bot/ext/utils/mousey.py
Normal file
97
bot/ext/utils/mousey.py
Normal file
|
@ -0,0 +1,97 @@
|
|||
import time
|
||||
import asyncio
|
||||
import functools
|
||||
|
||||
from typing import List
|
||||
|
||||
|
||||
class Timer:
|
||||
"""Context manager to measure how long the indented block takes to run."""
|
||||
|
||||
def __init__(self):
|
||||
self.start = None
|
||||
self.end = None
|
||||
|
||||
def __enter__(self):
|
||||
self.start = time.perf_counter()
|
||||
return self
|
||||
|
||||
async def __aenter__(self):
|
||||
return self.__enter__()
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
self.end = time.perf_counter()
|
||||
|
||||
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
||||
return self.__exit__(exc_type, exc_val, exc_tb)
|
||||
|
||||
def __str__(self):
|
||||
return f'{self.duration:.3f}ms'
|
||||
|
||||
@property
|
||||
def duration(self):
|
||||
"""Duration in ms."""
|
||||
return (self.end - self.start) * 1000
|
||||
|
||||
|
||||
class Table:
|
||||
def __init__(self, *column_titles: str):
|
||||
self._rows = [column_titles]
|
||||
self._widths = []
|
||||
|
||||
for index, entry in enumerate(column_titles):
|
||||
self._widths.append(len(entry))
|
||||
|
||||
def _update_widths(self, row: tuple):
|
||||
for index, entry in enumerate(row):
|
||||
width = len(entry)
|
||||
if width > self._widths[index]:
|
||||
self._widths[index] = width
|
||||
|
||||
def add_row(self, *row: str):
|
||||
"""
|
||||
Add a row to the table.
|
||||
.. note :: There's no check for the number of items entered, this may cause issues rendering if not correct.
|
||||
"""
|
||||
self._rows.append(row)
|
||||
self._update_widths(row)
|
||||
|
||||
def add_rows(self, *rows: List[str]):
|
||||
for row in rows:
|
||||
self.add_row(*row)
|
||||
|
||||
def _render(self):
|
||||
def draw_row(row_):
|
||||
columns = []
|
||||
|
||||
for index, field in enumerate(row_):
|
||||
# digits get aligned to the right
|
||||
if field.isdigit():
|
||||
columns.append(f" {field:>{self._widths[index]}} ")
|
||||
continue
|
||||
|
||||
# regular text gets aligned to the left
|
||||
columns.append(f" {field:<{self._widths[index]}} ")
|
||||
|
||||
return "|".join(columns)
|
||||
|
||||
# column title is centered in the middle of each field
|
||||
title_row = "|".join(f" {field:^{self._widths[index]}} " for index, field in enumerate(self._rows[0]))
|
||||
separator_row = "+".join("-" * (width + 2) for width in self._widths)
|
||||
|
||||
drawn = [title_row, separator_row]
|
||||
# remove the title row from the rows
|
||||
self._rows = self._rows[1:]
|
||||
|
||||
for row in self._rows:
|
||||
row = draw_row(row)
|
||||
drawn.append(row)
|
||||
|
||||
return "\n".join(drawn)
|
||||
|
||||
async def render(self, loop: asyncio.AbstractEventLoop=None):
|
||||
"""Returns a rendered version of the table."""
|
||||
loop = loop or asyncio.get_event_loop()
|
||||
|
||||
func = functools.partial(self._render)
|
||||
return await loop.run_in_executor(None, func)
|
|
@ -4,3 +4,6 @@ db = {
|
|||
'database': 'memed',
|
||||
'host': 'localhost'
|
||||
}
|
||||
|
||||
bot_token = 'Mzg2Mjc1MDc3MzY2NDgwODk4.DQNi9A.BnyE5MnKaIaVMBbWiW9rVDwkrSs'
|
||||
owner_id = 162819866682851329
|
||||
|
|
9
memed.py
9
memed.py
|
@ -12,6 +12,7 @@ import logging
|
|||
import asyncpg
|
||||
|
||||
import config
|
||||
from bot import schedule_bot
|
||||
|
||||
logging.basicConfig(level=logging.DEBUG)
|
||||
log = logging.getLogger(__name__)
|
||||
|
@ -64,6 +65,7 @@ async def send_msg(writer, op: int, data: str):
|
|||
await writer.drain()
|
||||
|
||||
async def process(op: int, message: str):
|
||||
"""Process messages given through the socket"""
|
||||
if op == 1:
|
||||
uid, cwd, command = parse_logstr(message)
|
||||
|
||||
|
@ -71,7 +73,8 @@ async def process(op: int, message: str):
|
|||
INSERT INTO logs (uid, cwd, cmd) VALUES ($1, $2, $3)
|
||||
""", uid, cwd, command)
|
||||
|
||||
async def handle_echo(reader, writer):
|
||||
async def handle_client(reader, writer):
|
||||
"""Handle clients"""
|
||||
try:
|
||||
await send_msg(writer, 0, 'hello')
|
||||
|
||||
|
@ -92,12 +95,14 @@ async def handle_echo(reader, writer):
|
|||
|
||||
if __name__ == '__main__':
|
||||
loop = asyncio.get_event_loop()
|
||||
coro = asyncio.start_unix_server(handle_echo, './log.suck',
|
||||
coro = asyncio.start_unix_server(handle_client, './log.suck',
|
||||
loop=loop)
|
||||
|
||||
db = loop.run_until_complete(asyncpg.create_pool(**config.db))
|
||||
server = loop.run_until_complete(coro)
|
||||
|
||||
schedule_bot(loop, config, db)
|
||||
|
||||
log.info(f'Serving on {server.sockets[0].getsockname()}')
|
||||
try:
|
||||
loop.run_forever()
|
||||
|
|
|
@ -1 +1,3 @@
|
|||
git+https://github.com/Rapptz/discord.py@rewrite
|
||||
asyncpg==0.13.0
|
||||
|
||||
|
|
Loading…
Reference in a new issue