Fix infinite loop and check file size for URLs

Fix bad dependency in requirements

requirements.txt

@@ -1,6 +1,6 @@
 opencv-python
 numpy
 pillow
-pycord
+py-cord
 python-dotenv
 markovchain

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

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,16 +54,19 @@ 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:
     async with session.get(url) as resp:
       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