src/danbooru_backend.py

@@ -9,11 +9,15 @@ import os
 
 
 def random_tag(*tags):
-	return len(tags) == 1 and tags[0].lower() == "random"
+	if len(tags) == 1 and tags[0].lower() == "random":
+		return True
+	return False
 
 
-def get_most_sever_rating(rating):
-	return rating in ("q", "e")
+def get_most_sever_rating(*ratings):
+	if "q" in ratings or "e" in ratings:
+		return True
+	return False
 
 
 class downloader:
@@ -38,8 +42,9 @@ class downloader:
 			print("Remote image request returned:", remote_image.status_code)
 			return None
 
-		with open(full_path, "wb") as f:
-			f.write(remote_image.content)
+		for d in full_path:
+			with open(d, "wb") as f:
+				f.write(remote_image.content)
 
 		return post
 
@@ -75,11 +80,11 @@ class downloader:
 			if tag_type in response:
 				tag_response.append(response[tag_type].strip())
 
-		nsfw = response["rating"]
+		nsfw = search_request.json()[0]["rating"]
 		nsfw = get_most_sever_rating(nsfw)
 
 		file_url = response["file_url"]
-
+		
 		basename = file_url.rsplit("/", 1)[1]
 		full_path = os.path.join(self.tmp, basename)
 
@@ -94,7 +99,7 @@ class downloader:
 			# Query results
 			"search_url":		search_url,
 			"file_url":			file_url,
-			"full_path":		full_path,
+			"full_path":		[full_path],
 			"tag_response":		" ".join(tag_response),
 			"nsfw":				nsfw
 			}

src/main.py

@@ -268,5 +268,7 @@ if __name__ == "__main__":
 		sys.exit(5)
 	except yandere_bot.BadCfgFile:
 		sys.exit(4)
+	except yandere_bot.BadPostSettings:
+		sys.exit(3)
 	except yandere_bot.FailedLogin:
 		sys.exit(2)

src/yandere_bot.py

@@ -24,7 +24,6 @@ import math
 import shutil
 import importlib
 import magic
-import random
 from threading import Event
 from mastodon import Mastodon, MastodonIllegalArgumentError, MastodonAPIError, MastodonVersionError
 
@@ -69,7 +68,6 @@ class YandereBot:
 		self.cfg = cfg
 		self.load_settings(self.cfg)
 		self.debug_mode = debug_mode or self.settings_behavior["debug"]
-		random.seed(os.urandom(16))
 		if prime_bot:
 			self.prime_bot()
 
@@ -144,9 +142,15 @@ class YandereBot:
 			self.currentSessionCount, self.failed_uploads)	)
 	
 
-	
+	def is_banned(self, tag_response):
+		for tag in self.settings_banned:
+			if tag in tag_response:
+				return True
+		return False
+
+
 	# Returns a list of media paths (without the hashes)
-	def download_media(self, picked):
+	def get_media_list(self, picked):
 		try:
 			backend_s = picked["backend"]
 			backend = importlib.import_module(backend_s)
@@ -155,12 +159,34 @@ class YandereBot:
 			img = None
 			downloader = backend.downloader(username, password, tmp=self.settings_behavior["tmp_dir"])
 
-			img = downloader.fetch_post(picked)
-			
-			if img is None:
-				raise InvalidPost("Img could not be downloaded")
-
-			return downloader.download_post(img)
+			while self.can_post():
+				img = downloader.fetch_post(picked)
+				if img is None:
+					self.eventSleep.wait(1)
+					break
+				if self.is_banned(img["tag_response"]):
+					print("Banned tag:", img["tag_response"])
+					self.eventSleep.wait(1)
+					continue
+				if downloader.download_post(img) is None:
+					print("Failed to download image")
+					self.eventSleep.wait(1)
+					continue
+				# Make sure the file is not malicious
+				full_path = img["full_path"]
+				mime = [magic.from_file(path, mime=True) for path in full_path]
+				if None in mime:
+					print("Unknown mime type.", mime)
+					self.eventSleep.wait(1)
+					continue
+				mime_categories = [c.split("/", 1)[0] for c in mime]
+				invalid_mime_categories = [c for c in mime_categories if c not in ("image", "video")]
+				if invalid_mime_categories:
+					print("mime type not allowed:", mime)
+					self.eventSleep.wait(1)
+					continue
+				break
+			return img
 		except ImportError:
 			print("Invalid Backend:", picked["backend"])
 			return None
@@ -198,43 +224,9 @@ class YandereBot:
 		string_post = content_newline.join(filter(None, (static_message, string_imglinks_joined)))
 
 		return content_type, string_post
-	
-
-	def is_banned(self, picked):
-		tag_response = picked["tag_response"]
-		for tag in self.settings_banned:
-			if tag in tag_response:
-				return True
-		return False
-
-
-	def valid_mimetype(self, picked):
-		full_path = picked["full_path"]
-		mime = magic.from_file(full_path, mime=True)
-
-		if mime is None:
-			return False
-
-		mime_category = mime.split("/", 1)[0]
-
-		return mime_category in ("image", "video")
-
 
 	def _post(self, picked):
-		# Validate picked
-		if picked is None:
-			raise InvalidPost("Picked post is None")
-		elif self.is_banned(picked):
-			raise BannedTag("Tag is banned")
-		elif not self.valid_mimetype(picked):
-			raise InvalidMimeType("Invalid mime type")
-
-		full_path = picked["full_path"]
-		if not os.path.isfile(full_path):
-			raise FileNotFoundError("File not found: {}".format(media_list))
-
-		media_list = self.upload_media_list([full_path])
-
+		media_list = self.upload_media_list(picked["full_path"])
 		content_type, message = self.get_post_text(picked, media_list)
 		if self.debug_mode:
 			return picked
@@ -247,7 +239,6 @@ class YandereBot:
 		)
 		return picked
 
-
 	def pick_profile(self):
 		profiles = self.settings_post
 		tag_select = self.settings_behavior["tag_select"].lower()
@@ -270,27 +261,37 @@ class YandereBot:
 	def post(self):
 		picked = None
 
+		# Flags that are set if an upload fails
+		timeout = False
+
 		# Attempt post
 		try:
 			# Post
-			picked_profile = self.pick_profile()
-			print("Posting...", picked_profile["name"])
+			while not picked and self.can_post():
+				picked_profile = self.pick_profile()
+				print("Posting...", picked_profile["name"])
+				picked = self.get_media_list(picked_profile)
+				self.currentIndexCount += 1
+
+			if not self.can_post():
+				return None
 			
-			picked = self.download_media(picked_profile)
 			self._post(picked)
 			
-			os.remove(picked["full_path"])
+			for path in picked["full_path"]:
+				os.remove(path)
 
 			# After a successful post
 			self.currentSessionCount += 1
-			self.currentIndexCount += 1
 
 			# The post was successful
 			return picked
 
-		# Failed post
-		except (BannedTag, UnknownMimeType, InvalidMimeType, FileNotFoundError) as e:
-			print("Posting error:", e)
+		# Attempted to post a file that doesn't exist (immediately repost ignoring retry_seconds)
+		except FileNotFoundError:
+			print("File not found:", picked["full_path"])
+			# Exception flags
+			timeout = False
 
 		# Check if the file limit has been reached
 		except MastodonAPIError as e:
@@ -300,6 +301,8 @@ class YandereBot:
 				file_limit_reached = (e.args[1] == 413)
 
 			print("API Error:", e)
+			# Exception flags
+			timeout = True
 
 		# Server Errors
 		# Assume all exceptions are on the server side
@@ -309,16 +312,17 @@ class YandereBot:
 		# 2. The server is down. Check to verify in a web browser (this is the default assumption since the
 		#    mastodon.py API will not specify why the connection timed out).
 		# The default assumption is #2
-		except Exception as e:
-			print("Unhandled Exception:", e)
+		# except Exception as e:
+		# 	print("Unhandled Exception:", e)
+		# 	# Exception flags
+		# 	timeout = True
 
 		# An exception occurred
 		self.failed_uploads += 1
 		self.currentIndexCount += 1
 		print("[Errors: {}]".format(self.failed_uploads))
-
-		# Sleep
-		self.eventSleep.wait(self.settings_behavior["retry_seconds"])
+		if timeout:
+			self.eventSleep.wait(self.settings_behavior["retry_seconds"])
 
 		# The post failed
 		return None
@@ -413,19 +417,7 @@ class Debug(Exception):
 	pass
 
 
-class InvalidPost(Exception):
-	pass
-
-
-class BannedTag(Exception):
-	pass
-
-
-class UnknownMimeType(Exception):
-	pass
-
-
-class InvalidMimeType(Exception):
+class BadPostSettings(Exception):
 	pass