Compare commits

...

2 Commits

2 changed files with 64 additions and 55 deletions

View File

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

View File

@ -44,7 +44,6 @@ 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
@ -89,7 +88,6 @@ class YandereBot:
"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:
@ -199,15 +197,6 @@ class YandereBot:
return content_type, string_post 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): def valid_mimetype(self, picked):
full_path = picked["full_path"] full_path = picked["full_path"]
mime = magic.from_file(full_path, mime=True) mime = magic.from_file(full_path, mime=True)
@ -220,21 +209,19 @@ class YandereBot:
return mime_category in ("image", "video") return mime_category in ("image", "video")
def _post(self, picked): def _post(self, picked):
# Validate picked # Validate picked
if picked is None: if picked is None:
raise InvalidPost("Picked post is None") raise InvalidPost("Picked post is None")
elif self.is_banned(picked):
raise BannedTag("Tag is banned") if not os.path.isfile(picked["full_path"]):
raise FileNotFoundError("File not found: {}".format(picked))
elif not self.valid_mimetype(picked): elif not self.valid_mimetype(picked):
raise InvalidMimeType("Invalid mime type") raise InvalidMimeType("Invalid mime type")
full_path = picked["full_path"] media_list = self.upload_media_list([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) content_type, message = self.get_post_text(picked, media_list)
if self.debug_mode: if self.debug_mode:
return picked return picked
@ -274,8 +261,6 @@ class YandereBot:
try: try:
# Post # Post
picked_profile = self.pick_profile() picked_profile = self.pick_profile()
print("Posting...", picked_profile["name"])
picked = self.download_media(picked_profile) picked = self.download_media(picked_profile)
self._post(picked) self._post(picked)
@ -293,7 +278,7 @@ class YandereBot:
print("Invalid post:", e) print("Invalid post:", e)
# Failed post # Failed post
except (BannedTag, UnknownMimeType, InvalidMimeType, FileNotFoundError) as e: except (InvalidMimeType, FileNotFoundError) as e:
# Decrement currentIndexCount to repost from the same profile # Decrement currentIndexCount to repost from the same profile
self.currentIndexCount -= 1 self.currentIndexCount -= 1
print("Posting error:", e) print("Posting error:", e)
@ -304,6 +289,8 @@ class YandereBot:
file_limit_reached = False file_limit_reached = False
with contextlib.suppress(IndexError): with contextlib.suppress(IndexError):
file_limit_reached = (e.args[1] == 413) file_limit_reached = (e.args[1] == 413)
if file_limit_reached:
self.currentIndexCount -= 1
print("API Error:", e) print("API Error:", e)
@ -423,14 +410,6 @@ class InvalidPost(Exception):
pass pass
class BannedTag(Exception):
pass
class UnknownMimeType(Exception):
pass
class InvalidMimeType(Exception): class InvalidMimeType(Exception):
pass pass