Refactored, image validation is on the backend

src/danbooru_backend.py

@@ -4,6 +4,7 @@
 # Update to enable multiple file downloads
 # Update to include safe/nsfw filter
 
+import magic
 import requests
 import os
 
@@ -12,15 +13,46 @@ def random_tag(*tags):
 	return len(tags) == 1 and tags[0].lower() == "random"
 
 
-def get_most_sever_rating(rating):
-	return rating in ("q", "e")
+def collect_tags(post):
+	tag_response = []
+	for tag_type in "tag_string", "tag_string_general":
+		if tag_type in post:
+			tag_response.append(post[tag_type].strip())
+	return " ".join(tag_response)
+	
+
+
+def is_banned(post, profile):
+	tag_response = collect_tags(post)
+	tag_banned = profile["banned_tags"]
+	for tag in tag_banned:
+		if tag in tag_response:
+			return tag
+	return None
+
+
+
+def get_nsfw(post):
+	return post["rating"] in ("q", "e")
+
+
+def select_from_response(response, profile):
+	for post in response:
+		if is_banned(post, profile):
+			continue
+		elif "file_url" not in post:
+			continue
+		elif profile["force_nsfw"] is not None and profile["force_nsfw"] != get_nsfw(post):
+			continue
+		return post
+	return None
 
 
 class downloader:
 	username = None
 	password = None
 	tmp = None
-	unallowed_extensions = (".zip",)
+	limit=50
 
 	def __init__(self, username=None, password=None, tmp="/tmp"):
 		self.username = username
@@ -48,7 +80,7 @@ class downloader:
 		# Search ratings: s=safe, e=nsfw
 		# base_url = "https://danbooru.donmai.us/posts.json?random=true&tags={}&rating=e&limit=1"
 		tags = profile["tags"]
-		search_url = "https://danbooru.donmai.us/posts.json?random=true&limit=1"
+		search_url = "https://danbooru.donmai.us/posts.json?random=true&limit={}".format(self.limit)
 		if tags and not random_tag(*tags):
 			search_tags = "+".join(tags)
 			search_url = "{}&tags={}".format(search_url, search_tags)
@@ -62,34 +94,32 @@ class downloader:
 			search_request = requests.get(search_url)
 	
 		if search_request.status_code != 200:
+			print(search_url)
 			print("Search request returned:", search_request.status_code)
 			return None
-		response = search_request.json()[0]
-		tag_response = []
-		if "file_url" not in response:
-			print("file_url is not in response")
-			return None
 		
-		# Aggregate Tags
-		for tag_type in "tag_string", "tag_string_general":
-			if tag_type in response:
-				tag_response.append(response[tag_type].strip())
+		selected = select_from_response(search_request.json(), profile)
 
-		nsfw = response["rating"]
-		nsfw = get_most_sever_rating(nsfw)
+		if selected is None:
+			print("Could not select image based on criteria")
+			return None
 
-		file_url = response["file_url"]
+		tag_response = collect_tags(selected)
+		nsfw = get_nsfw(selected)
+		file_url = selected["file_url"]
 
 		basename = file_url.rsplit("/", 1)[1]
 		full_path = os.path.join(self.tmp, basename)
 
 		r =	{
 			# Add profile to dictioanry
-			"name":		profile["name"],
-			"backend":	profile["backend"],
-			"tags":		profile["tags"],
-			"message":	profile["message"],
-			"message_nsfw":	profile["message_nsfw"],
+			"name":				profile["name"],
+			"backend":			profile["backend"],
+			"tags":				profile["tags"],
+			"message":			profile["message"],
+			"message_nsfw":		profile["message_nsfw"],
+			"force_nsfw":		profile["force_nsfw"],
+			"banned_tags":		profile["banned_tags"],
 
 			# Query results
 			"search_url":		search_url,

src/yandere_bot.py

@@ -44,7 +44,6 @@ class YandereBot:
 	settings_post_default = None
 	settings_encrypt = None
 	settings_credentials = None
-	settings_banned = None
 
 	# Class variables
 	mastodon_api = None
@@ -89,7 +88,6 @@ class YandereBot:
 				"settings_post_default",
 				"settings_encrypt",
 				"settings_credentials",
-				"settings_banned"
 			)
 			_settings = settings or default_settings
 			for ele in _settings:
@@ -198,43 +196,32 @@ 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")
+			raise InvalidPost("Picked post is None")	
+
+		if not os.path.isfile(picked["full_path"]):
+			raise FileNotFoundError("File not found: {}".format(picked))
+		
 		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
@@ -274,8 +261,6 @@ class YandereBot:
 		try:
 			# Post
 			picked_profile = self.pick_profile()
-			print("Posting...", picked_profile["name"])
-			
 			picked = self.download_media(picked_profile)
 			self._post(picked)
 			
@@ -293,7 +278,7 @@ class YandereBot:
 			print("Invalid post:", e)
 
 		# Failed post
-		except (BannedTag, UnknownMimeType, InvalidMimeType, FileNotFoundError) as e:
+		except (InvalidMimeType, FileNotFoundError) as e:
 			# Decrement currentIndexCount to repost from the same profile
 			self.currentIndexCount -= 1
 			print("Posting error:", e)
@@ -304,6 +289,8 @@ class YandereBot:
 			file_limit_reached = False
 			with contextlib.suppress(IndexError):
 				file_limit_reached = (e.args[1] == 413)
+			if file_limit_reached:
+				self.currentIndexCount -= 1
 
 			print("API Error:", e)
 
@@ -423,14 +410,6 @@ class InvalidPost(Exception):
 	pass
 
 
-class BannedTag(Exception):
-	pass
-
-
-class UnknownMimeType(Exception):
-	pass
-
-
 class InvalidMimeType(Exception):
 	pass