AO3-DL/python/common.py

# 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.2.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,
) -> 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=sys.stdout,
    )
    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,
        )
    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["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)