# This file is part of AOMO, "Archive Of My Own", a collection of Python and PHP # scripts designed act as an improved local backup system # for works published on https://archiveofourown.org # # Copyright (c) 2025 Cyberpro123, except where otherwise noted. # # This project is available under the GNU General Public License (GPL) v2.0, # available at https://www.gnu.org/licenses/old-licenses/gpl-2.0.en.html # or in the 'LICENSE.md' file that should be distributed alongside this one. # # Report issues to https://codeberg.org/Cyberpro123/AOMO # Contact author at cyberpro123@posteo.com import argparse import configparser import logging import os.path from _io import TextIOWrapper import sys ################################################################ # Constants ################################################################ version = "1.3.0+unstable" loopRetries = 100 # How many times to retry doing an action that sometimes fails (ie downloading a work) before giving up and throwing an exception loopErrorTemplate = "Experienced {2} while {0}, sleeping for {1} seconds before retrying. Arguments: {3!r}" ao3WorksPerSeriesPage = 20 threadNameBulk = "worker" commentHeader = "AOMO Metadata Header" commentVersion = "Version: " commentTimestampDownloaded = "Date Downloaded: " commentTimestampEdited = "Date Edited: " imagePrefix = "\uea03" imagePrefixSubstitute = "\uea04" ################################################################ # Exceptions ################################################################ class UnforseenBehaviorError(Exception): def __init__(self, message): super().__init__(message) self.message = message class NeedLoginError(Exception): def __init__(self, message): super().__init__(message) self.message = message class InvalidConfigurationError(Exception): def __init__(self, message): super().__init__(message) self.message = message class LoopFailError(Exception): def __init__(self, message): super().__init__(message) self.message = message class HTTPError(Exception): def __init__(self, description, status_code, url=None, content=None): self.desciption = description self.status_code = status_code self.content = content self.url = url super().__init__(description, status_code, url, content) def __str__(self): string = f"(Error {self.status_code}): {self.desciption}" if self.url: string += f"\nURL: {self.url}" if self.content: string += f"\nContent:\n{self.content}" return string class UnavailablePageError(Exception): def __init__(self, message): super().__init__(message) self.message = message ################################################################ # Functions ################################################################ def parseInfile(infile) -> set: out = set(()) for lineRaw in infile: line = str(lineRaw).split("#")[0].strip() if line: out.add(line) return out def getLogger( level: int = logging.INFO, includeThreadName: bool = True, mode: str = "stdout", stream: TextIOWrapper = None, ) -> logging.Logger: if not stream: match mode: case "stdout": stream = sys.stdout case "stderr": stream = sys.stderr logger = logging.getLogger(__name__) if includeThreadName: formatStr = "[%(asctime)s] [%(levelname)-8s] [%(threadName)-8s]: %(message)s" else: formatStr = "[%(asctime)s] [%(levelname)-8s]: %(message)s" logging.basicConfig( stream=stream, level=level, format=formatStr, datefmt="%H:%M:%S", ) if stream != sys.stderr: logger.info("Logger Initialized") return logger def tail( filepath: str, lines: int = 10, ) -> list: filesize = os.stat(filepath).st_size with open(filepath, errors="surrogateescape") as file: file.seek(max(filesize - 8192, 0)) buffer = file.read().split("\n") while len(buffer) > lines: buffer.pop(0) return buffer def getMetadataFromRawCommentByID( id: int, config: dict, ) -> tuple[str, int, int]: filepath = os.path.join(config["dirRaws"], f"{id}.html") try: lines = tail(filepath=filepath) except UnicodeDecodeError: lines = [] version = "" edited = 0 downloaded = 0 if commentHeader in str(lines): for line in lines: if commentVersion in line: version = line[len(commentVersion) :] elif commentTimestampEdited in line: edited = int(line[len(commentTimestampEdited) :]) elif commentTimestampDownloaded in line: downloaded = int(line[len(commentTimestampDownloaded) :]) elif not lines: logging.getLogger(__name__).info(f"Raw [{id}] couldn't be read by common.tail()") else: logging.getLogger(__name__).info(f"Raw [{id}] doesn't seem to have a metadata comment.") return version, edited, downloaded def getWorkIdFromLink( url: str, ) -> int | None: if "archiveofourown.org/works" in url: splitList = url.split("/") out = "" done = False for char in splitList[splitList.index("works") + 1]: try: int(char) except ValueError: done = True else: if not done: out += str(int(char)) if out: return out else: return None else: return None def init( description: str = None, json: str = "", dryRun: bool = False, all: str = "", auto: str = "", secondOut: str = "", main: tuple = None, since: str = None, defaultOutfile=sys.stdout, chapters: str = "", fileType: dict = {}, ) -> None: global args global config global logger parser = argparse.ArgumentParser(description=description) parser.add_argument( "--infile", "-i", nargs="?", type=argparse.FileType("r"), default=sys.stdin, ) parser.add_argument( "--outfile", "-o", nargs="?", type=argparse.FileType("w"), default=defaultOutfile, ) if secondOut: parser.add_argument( "--batch-out", "-b", nargs="?", type=argparse.FileType("w"), default=sys.stdout, help=secondOut, ) if json: parser.add_argument( "--json", "-j", nargs=1, type=argparse.FileType(json), help="JSON file containing a whitelist and/or blacklist of works to include.", ) if dryRun: parser.add_argument( "--dry-run", "-d", action="store_true", help="Perform a 'dry run' without downloading any files or saving anything to local storage.", ) if all or auto: autoGroup = parser.add_mutually_exclusive_group() if all: autoGroup.add_argument( "--all", action="store_true", help=all, ) if auto: autoGroup.add_argument( "--auto", action="store_true", help=auto, ) if main: parser.add_argument( main[0], nargs=1, type=str, help=main[1], ) if since: parser.add_argument( "--since", "-s", type=str, help=since, ) if chapters: parser.add_argument( "--chapters", type=int, nargs="*", default=[0], help=chapters, ) if fileType: parser.add_argument( "--fileType", type=str, choices=fileType["choices"], default=fileType["default"], help=fileType["help"], ) parser.add_argument("--log-level", "-l", default="") parser.add_argument("--config", "-c", default="") args = parser.parse_args() ini = configparser.ConfigParser() configArgFile = os.path.abspath(os.path.expanduser(args.config)) if args.config and os.path.exists(configArgFile) and os.path.getsize(configArgFile): ini.read(configArgFile) else: ini.read( [ "/etc/aomo/config.ini", os.path.expanduser("~/.config/aomo/config.ini"), os.path.join(os.path.dirname(__file__), "config.ini"), os.path.join(os.path.dirname(__file__), ".config.ini"), ] ) config = {} config["dirRaws"] = os.path.expanduser(ini.get("dir", "raws", fallback="~/Documents/AOMO/raws/")) config["dirAo3Css"] = os.path.expanduser(ini.get("dir", "ao3css", fallback="~/Documents/AOMO/ao3css/")) config["dirImg"] = os.path.expanduser(ini.get("dir", "images", fallback="~/Documents/AOMO/images/")) config["dirWebUi"] = os.path.expanduser(ini.get("dir", "webui", fallback="~/Documents/AOMO/webui/")) config["dirExport"] = os.path.expanduser(ini.get("dir", "export", fallback="~/Documents/AOMO/export/")) config["sqlType"] = ini.get("db", "type", fallback="sqlite") config["sqlLocation"] = ini.get("db", "location", fallback="main.sqlite") config["sqlUsername"] = ini.get("db", "username", fallback="") config["sqlPassword"] = ini.get("db", "password", fallback="") config["ao3UsernameFile"] = os.path.expanduser(ini.get("ao3", "usernameFile", fallback="~/Documents/AOMO/secrets/username.secret")) config["ao3PasswordFile"] = os.path.expanduser(ini.get("ao3", "passwordFile", fallback="~/Documents/AOMO/secrets/password.secret")) config["ao3SessionPickle"] = os.path.expanduser(ini.get("ao3", "pickle", fallback="~/Documents/AOMO/secrets/session.pickle")) config["ao3DoLogin"] = ini.getboolean("ao3", "login", fallback=False) config["ao3DoLoginAlways"] = ini.getboolean("ao3", "loginAlways", fallback=False) logLevelRaw = ini.get("logs", "level", fallback="") logLevel = logging.INFO if args.log_level: logLevelRaw = args.log_level if logLevelRaw: try: logLevel = int(logLevelRaw) except ValueError: for i in ( ("d", logging.DEBUG), ("i", logging.INFO), ("w", logging.WARNING), ("e", logging.ERROR), ("c", logging.CRITICAL), ): if str(logLevelRaw)[:1].lower() == i[0]: logLevel = i[1] logger = getLogger(level=logLevel, stream=args.outfile)