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: