Added support for banned tags

This commit is contained in:
Anon 2022-08-29 01:40:06 -07:00
parent 50c9f9da51
commit 1d9b183e08
3 changed files with 44 additions and 18 deletions

View File

@ -58,6 +58,8 @@ settings_credentials = {
} }
} }
settings_banned = tuple()
# Basic settings to configure Yandere Bot's behavior # Basic settings to configure Yandere Bot's behavior
settings_behavior = { settings_behavior = {
"max_size": 15*1024*1024, "max_size": 15*1024*1024,

View File

@ -24,13 +24,21 @@ class downloader:
username = None username = None
password = None password = None
tmp = 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.username = username
self.password = password self.password = password
self.tmp = tmp 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): def download(self, profile):
# Search ratings: s=safe, e=explicit # Search ratings: s=safe, e=explicit
@ -42,6 +50,7 @@ class downloader:
search_url = "{}&tags={}".format(search_url, search_tags) search_url = "{}&tags={}".format(search_url, search_tags)
search_request = None search_request = None
while True:
if self.username and self.password: if self.username and self.password:
search_request = requests.get(search_url, search_request = requests.get(search_url,
auth=(self.username, self.password) auth=(self.username, self.password)
@ -49,13 +58,25 @@ class downloader:
else: else:
search_request = requests.get(search_url) search_request = requests.get(search_url)
# Return None if the request failed
if search_request.status_code != 200: if search_request.status_code != 200:
print("Search request returned:", search_request.status_code) print("Search request returned:", search_request.status_code)
return None continue
# Check to make sure large_file_url key exists in the response
elif "large_file_url" not in search_request.json()[0]: elif "large_file_url" not in search_request.json()[0]:
return None 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
import pprint
pprint.pprint(search_request.json()[0]["tag_string"])
large_file_url = search_request.json()[0]["file_url"] large_file_url = search_request.json()[0]["file_url"]
@ -85,7 +106,7 @@ class downloader:
"search_url": search_url, "search_url": search_url,
"large_file_url": large_file_url, "large_file_url": large_file_url,
"full_path": [full_path], "full_path": [full_path],
"explicit": explicit "explicit": explicit,
} }

View File

@ -42,6 +42,7 @@ class YandereBot:
settings_post_default = None settings_post_default = None
settings_encrypt = None settings_encrypt = None
settings_credentials = None settings_credentials = None
settings_banned = None
# Class variables # Class variables
mastodon_api = None mastodon_api = None
@ -84,7 +85,8 @@ class YandereBot:
"settings_post", "settings_post",
"settings_post_default", "settings_post_default",
"settings_encrypt", "settings_encrypt",
"settings_credentials" "settings_credentials",
"settings_banned"
) )
_settings = settings or default_settings _settings = settings or default_settings
for ele in _settings: for ele in _settings:
@ -141,11 +143,12 @@ class YandereBot:
# Returns a list of media paths (without the hashes) # Returns a list of media paths (without the hashes)
def get_media_list(self, picked): def get_media_list(self, picked):
try: try:
banned = self.settings_banned
backend_s = picked["backend"] backend_s = picked["backend"]
backend = importlib.import_module(backend_s) backend = importlib.import_module(backend_s)
username = self.settings_credentials[backend_s]["username"] username = self.settings_credentials[backend_s]["username"]
password = self.settings_credentials[backend_s]["password"] 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) img = downloader.download(picked)
return img return img
except ImportError: except ImportError: