diff --git a/requirements.txt b/requirements.txt index 9d7693d..ac2382b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,6 @@ opencv-python numpy pillow -pycord +py-cord python-dotenv markovchain diff --git a/src/bot.py b/src/bot.py index d0194f5..2530833 100644 --- a/src/bot.py +++ b/src/bot.py @@ -36,7 +36,7 @@ async def on_thread_join(thread): await thread.join() @client.event -async def on_message(message): +async def on_message(message: discord.Message): try: # Ignore the bot's own messages if message.author == client.user: @@ -63,6 +63,8 @@ async def on_message(message): await message.channel.send('Hello!', delete_after=300) await asyncio.sleep(300) await message.delete() + elif message.content.lower().startswith('good bot'): + await message.channel.send('Thanks. Good human.') except BaseException as error: logging.exception(f'{datetime.now(timezone.utc)} main on_message failed') await message.channel.send("Something went wrong :(") diff --git a/src/faces.py b/src/faces.py index ed61cbc..a8c1bf2 100644 --- a/src/faces.py +++ b/src/faces.py @@ -7,32 +7,43 @@ from asyncio import run BYTES_IN_MEGABYTE = 1048576 +MAX_SIZE = BYTES_IN_MEGABYTE * 8 +class TooBigException(Exception): + pass allowable_attachment_types = ['image/png', 'image/jpeg'] -async def on_message(message, client, conf): - if message.content.startswith("$face"): - recursions = 0 - original_message = message - while recursions < 2: - if len(message.attachments) > 0: - for attachment in message.attachments: - # check if attachment is a picture and is smaller than 8mb - if attachment.content_type in allowable_attachment_types and attachment.size < (8 * BYTES_IN_MEGABYTE): - if attachment.proxy_url != '404': - await do_face(original_message, 'attachment_picture.png', lambda : attachment.save('attachment_picture.png', use_cached=True)) +async def on_message(message: discord.Message, client, conf): + try: + if message.content.startswith("$face"): + recursions = 0 + original_message = message + while recursions < 2: + if len(message.attachments) > 0: + for attachment in message.attachments: + # check if attachment is a picture and is smaller than 8mb + if attachment.content_type in allowable_attachment_types: + if attachment.size > MAX_SIZE: + raise TooBigException + if attachment.proxy_url != '404': + await do_face(original_message, 'attachment_picture.png', lambda : attachment.save('attachment_picture.png', use_cached=True)) + return + if len(message.embeds) > 0: + for embed in message.embeds: + if embed.type == 'image': + await do_face(original_message, 'attachment_picture.png', lambda : download_file(embed.url, 'attachment_picture.png')) return - if len(message.embeds) > 0: - for embed in message.embeds: - if embed.type == 'image': - await do_face(original_message, 'attachment_picture.png', lambda : download_file(embed.url, 'attachment_picture.png')) - return - if message.reference != None: - message = await message.channel.fetch_message(message.reference.message_id) - recursions += 1 + if message.reference != None: + message = await message.channel.fetch_message(message.reference.message_id) + recursions += 1 + else: + return + except TooBigException: + await message.channel.send('Your image is too big :(') -async def do_face(message, name, do_download): +async def do_face(message: discord.Message, name, do_download): try: + logging.info(f'{datetime.now(timezone.utc)} Starting face processing for user {message.author.display_name} {message.author.id}') await do_download() found, path = faces_util.get_face_replace(name) if found: @@ -43,12 +54,12 @@ async def do_face(message, name, do_download): os.remove('face_detected.png') if os.path.exists(name): os.remove(name) - except discord.HTTPException as error: - logging.exception(f'{datetime.now(timezone.utc)} Face on_message') - await message.channel.send('Could not download attachment :(') - except discord.NotFound as error: + except discord.NotFound: logging.exception(f'{datetime.now(timezone.utc)} Face on_message') await message.channel.send('Attachment not found :(') + except discord.HTTPException: + logging.exception(f'{datetime.now(timezone.utc)} Face on_message') + await message.channel.send('Could not download attachment :(') async def download_file(url, out_filename): async with aiohttp.ClientSession() as session: @@ -56,3 +67,6 @@ async def download_file(url, out_filename): img = await resp.read() with open(out_filename, 'wb') as f: f.write(img) + if f.tell() > MAX_SIZE: + os.remove(out_filename) + raise TooBigException