src/danbooru_backend.py

@@ -4,7 +4,6 @@
 # Update to enable multiple file downloads
 # Update to include safe/nsfw filter
 
-import magic
 import requests
 import os
 
@@ -13,46 +12,15 @@ def random_tag(*tags):
 	return len(tags) == 1 and tags[0].lower() == "random"
 
 
-def collect_tags(post):
+def get_most_sever_rating(rating):
-	tag_response = []
+	return rating in ("q", "e")
-	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
-	limit=50
+	unallowed_extensions = (".zip",)
 
 	def __init__(self, username=None, password=None, tmp="/tmp"):
 		self.username = username
@@ -80,7 +48,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={}".format(self.limit)
+		search_url = "https://danbooru.donmai.us/posts.json?random=true&limit=1"
 		if tags and not random_tag(*tags):
 			search_tags = "+".join(tags)
 			search_url = "{}&tags={}".format(search_url, search_tags)
@@ -94,32 +62,34 @@ 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]
-		selected = select_from_response(search_request.json(), profile)
+		tag_response = []
-
+		if "file_url" not in response:
-		if selected is None:
+			print("file_url is not in response")
-			print("Could not select image based on criteria")
 			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())
 
-		tag_response = collect_tags(selected)
+		nsfw = response["rating"]
-		nsfw = get_nsfw(selected)
+		nsfw = get_most_sever_rating(nsfw)
-		file_url = selected["file_url"]
+
+		file_url = response["file_url"]
 
 		basename = file_url.rsplit("/", 1)[1]
 		full_path = os.path.join(self.tmp, basename)
 
 		r =	{
 			# Add profile to dictioanry
-			"name":				profile["name"],
+			"name":		profile["name"],
-			"backend":			profile["backend"],
+			"backend":	profile["backend"],
-			"tags":				profile["tags"],
+			"tags":		profile["tags"],
-			"message":			profile["message"],
+			"message":	profile["message"],
-			"message_nsfw":		profile["message_nsfw"],
+			"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,6 +44,7 @@ class YandereBot:
 	settings_post_default = None
 	settings_encrypt = None
 	settings_credentials = None
+	settings_banned = None
 
 	# Class variables
 	mastodon_api = None
@@ -88,6 +89,7 @@ class YandereBot:
 				"settings_post_default",
 				"settings_encrypt",
 				"settings_credentials",
+				"settings_banned"
 			)
 			_settings = settings or default_settings
 			for ele in _settings:
@@ -196,32 +198,43 @@ 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")
 
+		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")	
+			raise InvalidPost("Picked post is None")
-
+		elif self.is_banned(picked):
-		if not os.path.isfile(picked["full_path"]):
+			raise BannedTag("Tag is banned")
-			raise FileNotFoundError("File not found: {}".format(picked))
-		
 		elif not self.valid_mimetype(picked):
 			raise InvalidMimeType("Invalid mime type")
 
-		media_list = self.upload_media_list([picked["full_path"]])
+		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])
+
 		content_type, message = self.get_post_text(picked, media_list)
 		if self.debug_mode:
 			return picked
@@ -261,6 +274,8 @@ class YandereBot:
 		try:
 			# Post
 			picked_profile = self.pick_profile()
+			print("Posting...", picked_profile["name"])
+			
 			picked = self.download_media(picked_profile)
 			self._post(picked)
 			
@@ -278,7 +293,7 @@ class YandereBot:
 			print("Invalid post:", e)
 
 		# Failed post
-		except (InvalidMimeType, FileNotFoundError) as e:
+		except (BannedTag, UnknownMimeType, InvalidMimeType, FileNotFoundError) as e:
 			# Decrement currentIndexCount to repost from the same profile
 			self.currentIndexCount -= 1
 			print("Posting error:", e)
@@ -289,8 +304,6 @@ 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)
 
@@ -410,6 +423,14 @@ class InvalidPost(Exception):
 	pass
 
 
+class BannedTag(Exception):
+	pass
+
+
+class UnknownMimeType(Exception):
+	pass
+
+
 class InvalidMimeType(Exception):
 	pass