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 include safe/nsfw filter
import magic
import requests
import os
@ -12,15 +13,46 @@ def random_tag(*tags):
return len(tags) == 1 and tags[0].lower() == "random"
def get_most_sever_rating(rating):
return rating in ("q", "e")
def collect_tags(post):
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:
username = None
password = None
tmp = None
unallowed_extensions = (".zip",)
limit=50
def __init__(self, username=None, password=None, tmp="/tmp"):
self.username = username
@ -48,7 +80,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=1"
search_url = "https://danbooru.donmai.us/posts.json?random=true&limit={}".format(self.limit)
if tags and not random_tag(*tags):
search_tags = "+".join(tags)
search_url = "{}&tags={}".format(search_url, search_tags)
@ -62,34 +94,32 @@ 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]
tag_response = []
if "file_url" not in response:
print("file_url is not in response")
selected = select_from_response(search_request.json(), profile)
if selected is None:
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())
nsfw = response["rating"]
nsfw = get_most_sever_rating(nsfw)
file_url = response["file_url"]
tag_response = collect_tags(selected)
nsfw = get_nsfw(selected)
file_url = selected["file_url"]
basename = file_url.rsplit("/", 1)[1]
full_path = os.path.join(self.tmp, basename)
r = {
# Add profile to dictioanry
"name": profile["name"],
"backend": profile["backend"],
"tags": profile["tags"],
"message": profile["message"],
"message_nsfw": profile["message_nsfw"],
"name": profile["name"],
"backend": profile["backend"],
"tags": profile["tags"],
"message": profile["message"],
"message_nsfw": profile["message_nsfw"],
"force_nsfw": profile["force_nsfw"],
"banned_tags": profile["banned_tags"],
# Query results
"search_url": search_url,

View File

@ -44,7 +44,6 @@ class YandereBot:
settings_post_default = None
settings_encrypt = None
settings_credentials = None
settings_banned = None
# Class variables
mastodon_api = None
@ -89,7 +88,6 @@ class YandereBot:
"settings_post_default",
"settings_encrypt",
"settings_credentials",
"settings_banned"
)
_settings = settings or default_settings
for ele in _settings:
@ -199,15 +197,6 @@ class YandereBot:
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)
@ -220,21 +209,19 @@ class YandereBot:
return mime_category in ("image", "video")
def _post(self, picked):
# Validate picked
if picked 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):
raise InvalidMimeType("Invalid mime type")
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])
media_list = self.upload_media_list([picked["full_path"]])
content_type, message = self.get_post_text(picked, media_list)
if self.debug_mode:
return picked
@ -274,8 +261,6 @@ class YandereBot:
try:
# Post
picked_profile = self.pick_profile()
print("Posting...", picked_profile["name"])
picked = self.download_media(picked_profile)
self._post(picked)
@ -293,7 +278,7 @@ class YandereBot:
print("Invalid post:", e)
# Failed post
except (BannedTag, UnknownMimeType, InvalidMimeType, FileNotFoundError) as e:
except (InvalidMimeType, FileNotFoundError) as e:
# Decrement currentIndexCount to repost from the same profile
self.currentIndexCount -= 1
print("Posting error:", e)
@ -304,6 +289,8 @@ 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)
@ -423,14 +410,6 @@ class InvalidPost(Exception):
pass
class BannedTag(Exception):
pass
class UnknownMimeType(Exception):
pass
class InvalidMimeType(Exception):
pass