From f831e0e5d9f2790e1b33e6a0e030f96c18da0c97 Mon Sep 17 00:00:00 2001
From: Anon <yanderefedi@proton.me>
Date: Thu, 1 Sep 2022 01:17:02 -0700
Subject: [PATCH] Added max file size feature

---
 src/yandere_bot.py | 32 ++++++++++++++++++++++++++------
 1 file changed, 26 insertions(+), 6 deletions(-)

diff --git a/src/yandere_bot.py b/src/yandere_bot.py
index 2f219d2..dc40c0a 100644
--- a/src/yandere_bot.py
+++ b/src/yandere_bot.py
@@ -147,7 +147,10 @@ class YandereBot:
 			username = self.settings_credentials[backend_s]["username"]
 			password = self.settings_credentials[backend_s]["password"]
 			img = None
-			downloader = backend.downloader(username, password, tmp=self.settings_behavior["tmp_dir"])
+			downloader = backend.downloader(
+					username, password,
+					self.settings_behavior["tmp_dir"]
+					)
 
 			img = downloader.fetch_post(picked_profile)
 			
@@ -205,6 +208,13 @@ class YandereBot:
 		return mime_category in ("image", "video")
 
 
+	def valid_file_size(self, picked):
+		full_path = picked["full_path"]
+		max_size = self.settings_behavior["max_size"]
+		file_size = os.stat(full_path).st_size
+
+		return file_size <= max_size
+
 
 	def _post(self, picked):
 		# Validate picked
@@ -217,6 +227,9 @@ class YandereBot:
 		elif not self.valid_mimetype(picked):
 			raise InvalidMimeType("Invalid mime type")
 
+		elif not self.valid_file_size(picked):
+			raise FileTooLarge("File is too large to upload")
+
 		media_list = self.upload_media_list([picked["full_path"]])
 		content_type, message = self.get_post_text(picked, media_list)
 		if self.debug_mode:
@@ -238,7 +251,7 @@ class YandereBot:
 		if tag_select == "random":
 			pick_index = random.randint(0, len(profiles) - 1)
 		elif tag_select == "sequential":
-			pick_index = self.currentIndexCount % len(profiles)
+			pick_index = max(0, self.currentIndexCount % len(profiles))
 		return profiles[pick_index]
 
 
@@ -273,11 +286,14 @@ class YandereBot:
 		except InvalidPost as e:
 			print("Invalid post:", e)
 
-		# Failed post
-		except (InvalidMimeType, FileNotFoundError) as e:
-			# Decrement currentIndexCount to repost from the same profile
+		except (FileTooLarge, InvalidMimeType) as e:
+			os.remove(picked["full_path"])
 			self.currentIndexCount -= 1
-			print("Posting error:", e)
+			print("Unable to post:", e)
+
+		except FileNotFoundError as e:
+			self.currentIndexCount -= 1
+			print("File not found:", e)
 
 		# Check if the file limit has been reached
 		except MastodonAPIError as e:
@@ -405,6 +421,10 @@ class InvalidPost(Exception):
 	pass
 
 
+class FileTooLarge(Exception):
+	pass
+
+
 class InvalidMimeType(Exception):
 	pass