diff --git a/src/main.py b/src/main.py index a9703e7..1a84495 100755 --- a/src/main.py +++ b/src/main.py @@ -16,10 +16,154 @@ # You should have received a copy of the GNU General Public License # along with this program. If not, see . +import os import sys import argparse import signal -import yandere_bot +import random +import datetime +import subprocess +import FediBot +from mastodon import MastodonAPIError + +# A class that contains all of the rendering information for a single post +class PickAndRenderFrame: + # From Config File + profile_name = "" + path = "" + frames = 0 + skip = tuple() + nsfw = tuple() + output_name = "" + message = "" + message_nsfw = "" + render_script = "" + + # After pick + picked_frame = None + output_name_tr = "" + + def __init__(self, picked, datetime_format): + dt_now = datetime.datetime.now() + dt_now_str = datetime.datetime.strftime(dt_now, datetime_format) + self.profile_name = picked["profile_name"] + self.path = picked["path"] + self.frames = picked["frames"] + self.skip = picked["skip"] + self.nsfw = picked["nsfw"] + self.output_name = picked["output_name"] + self.output_name_tr = self.output_name + self.message = picked["message"] + self.message_nsfw = picked["message_nsfw"] + self.render_script = picked["render_script"] + + self.picked_frame = self._pick_frame() + + # Shell-like substitutions + translations = { + "profile_name": self.profile_name, + "frame": str(self.picked_frame), + "datetime": dt_now_str + } + + self.output_name_tr = self._translate_basename(translations) + self._render_frame() + + def __del__(self): + render_file = self.output_name_tr + if render_file and os.path.isfile(render_file): + os.remove(render_file) + + def _pick_frame(self): + picked_frame = None + while picked_frame is None: + picked_frame = random.random() * self.frames + for skip in self.skip: + begin, end = skip + if begin <= picked_frame <= end: + picked_frame = None + break + return picked_frame + + def is_nsfw(self): + for nsfw in self.nsfw: + begin, end = nsfw + if begin <= self.picked_frame <= end: + return True + return False + + def _translate_basename(self, translations): + output_name_tr = self.output_name + for k,v in translations.items(): + replace_token = "${" + k + "}" + output_name_tr = output_name_tr.replace(replace_token, v) + return output_name_tr + + # TODO: Add translation keywords to the message + def get_message(self): + return self.message_nsfw if self.is_nsfw() else self.message + + def _render_frame(self): + args = [ + self.render_script, + self.path, + self.output_name_tr, + str(self.picked_frame) + ] + subprocess.run(args) + + +class YandereBot(FediBot.YandereBot): + def __init__(self, cfg, keyfile=None, debug_mode=False): + settings = { + "settings_time": {}, + "settings_post": {}, + } + self.settings.update(settings) + super(YandereBot, self).__init__(cfg, keyfile, debug_mode) + + # Maybe I should remove this from the backend? + def print_header_stats(self, picked): + profile, frame, nsfw, path = None, None, None, None + if picked: + profile = picked.profile_name + frame = picked.picked_frame + nsfw = picked.is_nsfw() + path = picked.output_name_tr + print("Profile: {} | Frame: {} | NSFW: {} | Path: {}".format( + profile, frame, nsfw, path + )) + + + def pick(self): + dt_picked = self.settings["settings_time"]["datetime"] + picked_profile = random.choice(self.settings["settings_post"]) + + picked = PickAndRenderFrame(picked_profile, dt_picked) + media_list = [picked.output_name_tr] + spoiler = picked.is_nsfw() + message = picked.get_message() + return { + "picked": picked, + "media_list": media_list, + "spoiler": spoiler, + "message": message + } + + def after_pick(self, picked): + self.print_header_stats(picked["picked"]) + + def post(self, picked): + try: + super(YandereBot, self).post(picked) + except (FileNotFoundError, MastodonAPIError, FediBot.InvalidPost, Exception) as e: + print("Exception:", e) + + self.handle_post_exception() + + # The post failed + return None + class FailedToLoadCfg(Exception): @@ -57,7 +201,7 @@ def main(): # Flag if the bot is running in debug mode debug_mode = arguments.dry_run or arguments.debug - yandere = yandere_bot.YandereBot( + yandere = YandereBot( yandere_config, arguments.keyfile, debug_mode @@ -84,9 +228,9 @@ if __name__ == "__main__": except FailedToLoadCfg: sys.exit(10) # Exceptions raised from the bot - except yandere_bot.Debug: + except FediBot.Debug: sys.exit(6) - except yandere_bot.BadCfgFile: + except FediBot.BadCfgFile: sys.exit(4) - except yandere_bot.FailedLogin: + except FediBot.FailedLogin: sys.exit(2)