From 1d9b183e081767303555e952ff100033d1a0e0c4 Mon Sep 17 00:00:00 2001
From: Anon <yanderefedi@proton.me>
Date: Mon, 29 Aug 2022 01:40:06 -0700
Subject: [PATCH] Added support for banned tags

---
 default/cfg.py          |  2 ++
 src/danbooru_backend.py | 51 +++++++++++++++++++++++++++++------------
 src/yandere_bot.py      |  9 +++++---
 3 files changed, 44 insertions(+), 18 deletions(-)

diff --git a/default/cfg.py b/default/cfg.py
index 2feaf68..aba7261 100644
--- a/default/cfg.py
+++ b/default/cfg.py
@@ -58,6 +58,8 @@ settings_credentials = {
 			}
 		}
 
+settings_banned = tuple()
+
 # Basic settings to configure Yandere Bot's behavior
 settings_behavior = {
 	"max_size":             15*1024*1024,
diff --git a/src/danbooru_backend.py b/src/danbooru_backend.py
index b9d1a61..8e3d011 100644
--- a/src/danbooru_backend.py
+++ b/src/danbooru_backend.py
@@ -24,12 +24,20 @@ class downloader:
 	username = None
 	password = None
 	tmp = None
+	banned = None
 
-	def __init__(self, username=None, password=None, tmp="/tmp"):
+	def __init__(self, banned = tuple(), username=None, password=None, tmp="/tmp"):
 		self.username = username
 		self.password = password
 		self.tmp = tmp
+		self.banned = banned
+	
 
+	def is_banned(self, tag_list):
+		for tag in self.banned:
+			if tag in tag_list:
+				return True
+		return False
 
 
 	def download(self, profile):
@@ -42,20 +50,33 @@ class downloader:
 			search_url = "{}&tags={}".format(search_url, search_tags)
 		
 		search_request = None
-		if self.username and self.password:
-			search_request = requests.get(search_url,
-					auth=(self.username, self.password)
-					)
-		else:
-			search_request = requests.get(search_url)
+		while True:
+			if self.username and self.password:
+				search_request = requests.get(search_url,
+						auth=(self.username, self.password)
+						)
+			else:
+				search_request = requests.get(search_url)
+	
+			if search_request.status_code != 200:
+				print("Search request returned:", search_request.status_code)
+				continue
+			elif "large_file_url" not in search_request.json()[0]:
+				continue
+			elif "tag_string" not in search_request.json()[0]:
+				continue
+			elif "tag_string_general" not in search_request.json()[0]:
+				continue
+			elif self.is_banned(search_request.json()[0]["tag_string"]):
+				print("Banned Tag1:", search_request.json()[0]["tag_string"])
+				continue
+			elif self.is_banned(search_request.json()[0]["tag_string_general"]):
+				print("Banned Tag2",search_request.json()[0]["tag_string_general"])
+				continue
+			break
 
-		# Return None if the request failed
-		if search_request.status_code != 200:
-			print("Search request returned:", search_request.status_code)
-			return None
-		# Check to make sure large_file_url key exists in the response
-		elif "large_file_url" not in search_request.json()[0]:
-			return None
+		import pprint
+		pprint.pprint(search_request.json()[0]["tag_string"])
 
 		large_file_url = search_request.json()[0]["file_url"]
 
@@ -85,7 +106,7 @@ class downloader:
 			"search_url":		search_url,
 			"large_file_url":	large_file_url,
 			"full_path":		[full_path],
-			"explicit":			explicit
+			"explicit":			explicit,
 			}
 
 
diff --git a/src/yandere_bot.py b/src/yandere_bot.py
index 5b28b42..5958b28 100644
--- a/src/yandere_bot.py
+++ b/src/yandere_bot.py
@@ -42,6 +42,7 @@ class YandereBot:
 	settings_post_default = None
 	settings_encrypt = None
 	settings_credentials = None
+	settings_banned = None
 
 	# Class variables
 	mastodon_api = None
@@ -84,7 +85,8 @@ class YandereBot:
 				"settings_post",
 				"settings_post_default",
 				"settings_encrypt",
-				"settings_credentials"
+				"settings_credentials",
+				"settings_banned"
 			)
 			_settings = settings or default_settings
 			for ele in _settings:
@@ -137,15 +139,16 @@ class YandereBot:
 			date_next_selection.strftime(self.settings_time["long_date_format"]), next_selection_seconds)	)
 		print("[ {} Selected during session | {} Failed ]\n".format(
 			self.currentSessionCount, self.failed_uploads)	)
-
+	
 	# Returns a list of media paths (without the hashes)
 	def get_media_list(self, picked):
 		try:
+			banned = self.settings_banned
 			backend_s = picked["backend"]
 			backend = importlib.import_module(backend_s)
 			username = self.settings_credentials[backend_s]["username"]
 			password = self.settings_credentials[backend_s]["password"]
-			downloader = backend.downloader(username, password, tmp=self.settings_behavior["tmp_dir"])
+			downloader = backend.downloader(banned, username, password, tmp=self.settings_behavior["tmp_dir"])
 			img = downloader.download(picked)
 			return img
 		except ImportError: