DanbooruBot/src/main.py

#! /usr/bin/env python3

# Danbooru Bot, an image posting bot for Pleroma
# Copyright (C) 2022 Anon <yanderefedi@proton.me>
# 
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

import sys
import argparse
import signal
import yandere_bot
import datetime
import contextlib


# A class that inherits YandereBot from the yandere_bot module
# This class is used to handle command line arguments gracefully, and extend functionality quickly
# Bot specific changes should be made here (if they are minor enough).
# This will be instantiated from the main() function
class YandereBot(yandere_bot.YandereBot):
	# The below settings are required from the configuration module
	settings_time = None
	settings_reminder = None

	# Wait before running
	delay_start = 0

	def __init__(self, cfg, debug_mode, prime_bot, start_index, delay_d=None, delay_h=None, delay_s=None):
		super(YandereBot, self).__init__(cfg, debug_mode, prime_bot=False)
		self.load_settings(self.cfg, ("settings_time", "settings_reminder"))
		self.set_pretimer(delay_d, delay_h, delay_s)
		self.currentIndexCount = start_index

		if self.debug_mode:
			print("[DEBUG MODE ON - DRY RUN BEGIN]")

		# Do not perform sanity test if stdout is None.
		# input() will fail if stdout is None, which could happen with the following command
		# ./run --dry-run -h -d 'a date in the past'
		if sys.stdout is not None and not self.pass_sanity_test():
			raise FailedSanityTest

		if prime_bot:
			self.prime_bot()
	
	def print_date_time_example(self):
		print_fmt = "  {0:6} {1:10} {2}"
		time_fmt = self.settings_time["time_format"]
		date_fmt = self.settings_time["date_format"]
		current_time = self.dateSelection

		print(print_fmt.format(
			"TIME", time_fmt, current_time.strftime(time_fmt)
		))
		print(print_fmt.format(
			"DATE", date_fmt, current_time.strftime(date_fmt)
		))

	def dump_pictures(self):
		for ele in self.listPictures:
			print(ele.get_full_string())

	def set_delay_d(self, d):
		try:
			t = datetime.datetime.strptime(d, self.settings_time["date_format"])
			self.dateNextSelection = self.dateNextSelection.replace(
				year=t.year, month=t.month, day=t.day
			)
		except Exception:
			print("Invalid date format: {}\n\nCorrect date/time format examples:".format(d))
			self.print_date_time_example()
			raise DateWrongFormat

	def set_delay_h(self, h, add_24):
		try:
			t = datetime.datetime.strptime(h, self.settings_time["time_format"])
			self.dateNextSelection = self.dateNextSelection.replace(
				hour=t.hour, minute=t.minute, second=t.second, microsecond=t.microsecond
			)
			if self.dateNextSelection < self.dateSelection and add_24:
				self.dateNextSelection = yandere_bot.time_add_seconds(self.dateNextSelection, 60 * 60 * 24)
		except Exception:
			print("Invalid time format: {}\n\nCorrect date/time format examples:".format(h))
			self.print_date_time_example()
			raise TimeWrongFormat

	def set_pretimer(self, d=None, h=None, s=0):
		if d:
			self.set_delay_d(d)
		if h:
			self.set_delay_h(h, d is None)
		if s:
			self.delay_start = max(0, s)

	# Check for potential misconfigurations by the user
	def pass_sanity_test(self):
		# Calculate pre-timer value
		seconds_until_next_pos = yandere_bot.time_diff_seconds(self.dateNextSelection, self.dateSelection)

		# Possible misconfigurations that will prompt the user to continue
		pretimer_less_than_zero = seconds_until_next_pos < 0
		pretimer_greater_than_sleep = seconds_until_next_pos > self.settings_behavior["sleep_seconds"]

		# Prompt the user
		prompt_user = pretimer_less_than_zero or pretimer_greater_than_sleep

		# Remind the user to generate new OAuth tokens
		dt = datetime.datetime.strptime(self.settings_reminder, self.settings_time["long_date_format"])
		if dt < datetime.datetime.now():
			print("REMINDER: Generate new tokens!!")

		# Check if the bot is back-posting in time and make sure this is what the user wanted to avoid spamming
		if pretimer_less_than_zero:
			sleep = round(abs(seconds_until_next_pos), 2)
			images = round(sleep / (self.settings_behavior["sleep_seconds"] * self.settings_behavior["uploads_per_post"]), 2) + 1
			print("WARNING: Pre-timer is less than the current time by: {} seconds. {} images will post immediately".format(
				sleep, images
			))
		# Check if the bot will wait for longer than the default amount of sleep time configured in the cfg.py
		elif pretimer_greater_than_sleep:
			print("WARNING: Pre-timer will sleep for {} seconds. This is more than the configured amount ({} seconds)".format(
				round(seconds_until_next_pos, 2), self.settings_behavior["sleep_seconds"]
			))

		# Prompt the user if something doesn't seem right
		# This must be done before we set up our keyboard interrupts. Otherwise the below exceptions will not work.
		try:
			if prompt_user:
				# Default to 'y' if the user just presses enter
				ans = input("Do you want to continue [Y/n]? ") or "y"
				return ans.lower() in ("y", "yes")
		except (EOFError, KeyboardInterrupt):
			print()
			return False

		# Sanity test passed
		return True

	def start(self, delay=None):
		_delay = delay or self.delay_start
		return super(YandereBot, self).start(max(0, _delay))


# Custom exceptions
class TimeWrongFormat(Exception):
	pass


class DateWrongFormat(Exception):
	pass


class FailedSanityTest(Exception):
	pass


class FailedToLoadCfg(Exception):
	pass

# Entry point if run from the command line
def main():
	# Default config file
	default_cfg = "cfg"

	# Parser
	parser = argparse.ArgumentParser(
		description="A bot for posting on Mastodon",
		# epilog="All switches can be combined for greater control",
		add_help=False)
	parser.add_argument("--dry-run", help="Will not login or post to Plemora", action="store_true")
	parser.add_argument("--debug", help="Same as --dry-run", action="store_true")
	parser.add_argument("-w", "--wait", type=int, help="Wait before posting first image (seconds)", default=0)
	parser.add_argument("-t", "--time", help="Wait for time before posting first image", default=None)
	parser.add_argument("-d", "--date", help="Wait for date before posting first image", default=None)
	parser.add_argument("-i", "--index", help="Start at index (only matters if profile select is set to sequential)", default=0)
	parser.add_argument("-c", "--config", help="Set custom config file (Default: {})".format(default_cfg), default=default_cfg)
	parser.add_argument("-o", "--output-hashes", help="Output list of hashes", action="store_true")
	parser.add_argument("-h", "--help", help="Show this help message and exit", action="store_true")
	parser.add_argument("remainder", help=argparse.SUPPRESS, nargs=argparse.REMAINDER)
	arguments = parser.parse_args()
	
	# Redirect stdout when the bot first initializes if the bot is not going to run normally
	redirect_stdout = None if arguments.output_hashes or arguments.help else sys.stdout

	# Yandere Lewd Bot
	yandere = None
	yandere_config = None

	# Configuration file for Yandere Lewd Bot
	try:
		import importlib
		yandere_config = importlib.import_module(arguments.config)
	except ImportError:
		print("Failed to Load Configuration:", arguments.config)
		raise FailedToLoadCfg

	# Flag if the bot is running in debug mode
	debug_mode = (arguments.dry_run or arguments.debug or arguments.output_hashes or arguments.help)

	with contextlib.redirect_stdout(redirect_stdout):
		prime_bot = not arguments.help

		yandere = YandereBot(
			yandere_config,
			debug_mode,
			prime_bot,
			int(arguments.index),
			arguments.date,
			arguments.time,
			arguments.wait
		)
	
	# Print Usage Information with Time and Date Formats with Examples
	if arguments.help:
		parser.print_help()
		print()
		yandere.print_date_time_example()
		print()
		return 0

	# Output all of the images in the bot's picture list
	if arguments.output_hashes:
		yandere.dump_pictures()
		return 0

	# Setup exit calls
	# Must be done after we declare our bot(s), otherwise this will be called if quitting on decrypting settings )
	def yandere_quit(signo, _frame):
		print("\nInterrupted by {}, shutting down...".format(signo))
		yandere.exit()

	# Setup our exit calls
	for sig in ("TERM", "HUP", "INT"):
		signal.signal(getattr(signal, "SIG" + sig), yandere_quit)

	return yandere.start()


if __name__ == "__main__":
	# A return value of 0 or 1 is a normal exit
	try:
		sys.exit(main())
	# Exceptions raised from the main function
	except FailedToLoadCfg:
		sys.exit(10)
	except FailedSanityTest:
		sys.exit(9)
	except DateWrongFormat:
		sys.exit(8)
	except TimeWrongFormat:
		sys.exit(7)

	# Exceptions raised from the bot
	except yandere_bot.Debug:
		sys.exit(6)
	except yandere_bot.MissingMasterList:
		sys.exit(5)
	except yandere_bot.BadCfgFile:
		sys.exit(4)
	except yandere_bot.FailedLogin:
		sys.exit(2)