AO3-DL/python/download.py
2026-01-06 21:08:06 -08:00

561 lines
19 KiB
Python
Executable file

#!/usr/bin/env python3
# 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 concurrent.futures
import datetime
import json
import os.path
import pickle
import random
import time
import traceback
import bs4
import requests
import requests.exceptions
import common
from sql import getWorkIDs, getCompleteWorkIDs, getExcluded
def getRequest(
reqType: str,
url: str,
logger: common.logging.Logger,
params: dict = None,
session: requests.Session = None,
) -> bytes:
if session:
req = session.request(reqType, url, params=params)
else:
req = requests.request(reqType, url, params=params)
if req.history:
if req.history[0].status_code == 302 and "/users/login" in req.url:
raise common.NeedLoginError(f"Redirected to [{req.url}]")
match req.status_code:
case 200:
return req.content
case 401:
raise common.UnavailablePageError("Error 401")
case 403:
raise common.UnavailablePageError("Error 403")
case 404:
raise common.UnavailablePageError("Error 404")
case 429:
raise common.HTTPError("Rate-limiting", 429)
case 503:
raise common.HTTPError("Resource Busy", 503)
case 525:
raise common.HTTPError("SSL Error", 525)
case _:
raise common.HTTPError(
"Unknown Error", req.status_code, req.url, req.content
)
def getPage(
url: str,
logger: common.logging.Logger,
retries: int = common.loopRetries,
session: requests.Session = None,
tryAnon: bool = True,
) -> bytes:
badHost = False
for i in ("imgur", "discordapp.com", "discordapp.net"):
if i in url:
badHost = True
retries = 3
loopNo = 1
if tryAnon:
sessionInUse = None
else:
sessionInUse = session
while loopNo <= retries:
try:
logger.log(
(10 + (20 * int(loopNo > 9))),
f"Attempt [{loopNo}] getting page [{url}]",
)
req = getRequest(reqType="get", url=url, logger=logger, session=session)
except common.NeedLoginError:
if sessionInUse:
errStr = f"Page [{url}] threw PrivateWorkError despite session in use"
if common.config["ao3SessionPickle"]:
errStr += f". Consider deleting existing session file at [{common.config['ao3SessionPickle']}]."
logger.critical(errStr)
raise common.UnforseenBehaviorError(errStr)
elif session:
sessionInUse = session
logger.info(f"Page [{url}] requires login, now using provided Session")
else:
errStr = f"Page [{url}] requires login but no Session was provided"
logger.critical(errStr)
raise common.NeedLoginError(errStr)
except (
common.HTTPError,
requests.exceptions.ConnectionError,
) as ex:
loopWait(
loopNo=loopNo,
ex=ex,
errLevel=(10 + (20 * int(loopNo > 9))),
logger=logger,
activity=f"getting page [{url}]",
badHost=badHost,
)
loopNo += 1
else:
logger.debug(f"Got page [{url}]")
return req
raise common.LoopFailError(f"Failed to get page [{url}] after [{loopNo}] attempts")
def getWorkskinFromWebsiteSoup(
soup: bs4.BeautifulSoup,
logger: common.logging.Logger,
) -> str:
allStyles = soup.find_all("style", {"type": "text/css"})
if len(allStyles) != 0:
text = str(allStyles[-1].decode_contents())
if text.find("#workskin") != -1:
return text
def getRawUrlFromWebsiteSoup(
soup: bs4.BeautifulSoup,
logger: common.logging.Logger,
) -> str:
# Significant portions of this function's code are derrived from 'AO3 API',
# https://github.com/wendytg/ao3_api/
# as of commit 02e349985d927bd8693f905f440e1ef0539f1984
# available under the MIT License Copyright (c) 2019 Francisco Patrício Rodrigues
dlMenuButton = soup.find("li", {"class": "download"})
for dlButton in dlMenuButton.find_all("li"):
if dlButton.a.getText() == "HTML":
return f"https://archiveofourown.org/{dlButton.a.attrs['href']}"
def getTagsFromWebsiteSoup(
soup: bs4.BeautifulSoup,
logger: common.logging.Logger,
) -> dict:
# Significant portions of this function's code are derrived from 'AO3 API',
# https://github.com/wendytg/ao3_api/
# as of commit 02e349985d927bd8693f905f440e1ef0539f1984
# available under the MIT License Copyright (c) 2019 Francisco Patrício Rodrigues
tags = {}
################################################################
# Tags
################################################################
for i in (
("warning", "warning"),
("category", "category"),
("fandom", "fandom"),
("relationship", "ship"),
("character", "char"),
("freeform", "tag"),
):
tag = soup.find("dd", {"class": f"{i[0]} tags"})
tags[i[1]] = []
if tag:
for j in tag.find_all("li"):
tags[i[1]].append(j.a.string)
del tag
################################################################
# Stats
################################################################
for i in ("comments", "kudos", "bookmarks", "hits", "words"):
tag = soup.find("dd", {"class": i})
if tag:
tags[i] = int(tag.string.replace(",", ""))
else:
tags[i] = 0
del tag
################################################################
# Misc Strings
################################################################
tags["rating"] = soup.find("dd", {"class": "rating tags"}).a.string
tags["lang"] = soup.find("dd", {"class": "language"}).string.strip()
################################################################
# Authors
################################################################
tags["author"] = []
for tag in soup.find("h3", {"class": "byline heading"}).find_all("a"):
paraOpenLoc = tag.string.rfind("(")
paraCloseLoc = tag.string.rfind(")")
if paraOpenLoc < paraCloseLoc and paraCloseLoc != -1:
tags["author"].append(tag.string[paraOpenLoc + 1 : paraCloseLoc])
else:
tags["author"].append(tag.string)
del paraOpenLoc
del paraCloseLoc
################################################################
# Timestamps
################################################################
tags["datePublished"] = 0
tags["datePublished"] = int(
datetime.datetime.fromisoformat(
soup.find("dd", {"class": "published"}).string
).timestamp()
)
tags["dateUpdated"] = 0
tag = soup.find("dd", {"class": "status"})
if tag:
tags["dateUpdated"] = int(
datetime.datetime.fromisoformat(tag.string).timestamp()
)
else:
tags["dateUpdated"] = tags["datePublished"]
del tag
tag = soup.find("li", {"class": "download"})
if tag and tag.ul:
tags["dateEdited"] = int(tag.ul.a["href"].split("=")[-1])
del tag
################################################################
# Chapter Counts
################################################################
tags["nChap"] = 0
tags["xChap"] = 0
tags["nChap"], tags["xChap"] = soup.find("dd", {"class": "chapters"}).string.split(
"/"
)
tags["complete"] = bool(tags["nChap"] == tags["xChap"])
return tags
def loopWait(
loopNo: int,
ex: Exception,
errLevel: int,
logger: common.logging.Logger,
activity: str,
badHost: bool = False,
) -> None:
random.seed()
pauseMult = 0.5 + random.random()
if badHost:
loopNo = 0.1
pauseTime = int(loopNo * 5 * pauseMult * (1 + 2 * int("429" in str(ex))))
logger.log(
errLevel,
common.loopErrorTemplate.format(
activity,
pauseTime,
type(ex).__name__,
ex.args,
),
)
time.sleep(pauseTime)
def login(
config: dict,
logger: common.logging.Logger,
) -> requests.Session | None:
# Significant portions of this function's code are derrived from 'AO3 API',
# https://github.com/wendytg/ao3_api/
# as of commit 02e349985d927bd8693f905f440e1ef0539f1984
# available under the MIT License Copyright (c) 2019 Francisco Patrício Rodrigues
if config["ao3DoLogin"]:
pickleFilepath = config["ao3SessionPickle"]
usernameFilepath = config["ao3UsernameFile"]
passwordFilepath = config["ao3PasswordFile"]
if (
pickleFilepath
and os.path.exists(pickleFilepath)
and os.path.getsize(pickleFilepath)
and (
os.path.getmtime(pickleFilepath) - datetime.datetime.now().timestamp()
< datetime.timedelta(days=30).total_seconds()
)
):
with open(pickleFilepath, "rb") as file:
session = pickle.load(file)
return session
usernameExists = bool(
os.path.exists(usernameFilepath) and os.path.getsize(usernameFilepath)
)
passwordExists = bool(
os.path.exists(passwordFilepath) and os.path.getsize(passwordFilepath)
)
errStr = ""
if (not usernameExists) and (not passwordExists):
errStr = f"Username & password files [{usernameFilepath}] and [{passwordFilepath}] either don't exist or are 0 bytes long."
if not usernameExists:
errStr = f"Username file [{usernameFilepath}] either doesn't exist or is 0 bytes long."
if not passwordExists:
errStr = f"Password file [{passwordFilepath}] either doesn't exist or is 0 bytes long."
if errStr:
logger.critical(errStr)
raise FileNotFoundError(errStr)
del errStr
with open(usernameFilepath) as file:
usernameStr = file.read().strip()
with open(passwordFilepath) as file:
passwordStr = file.read().strip()
session = requests.Session()
tempSoup = bs4.BeautifulSoup(
getRequest(
reqType="get",
url="https://archiveofourown.org/users/login",
logger=logger,
session=session,
),
"lxml",
)
getRequest(
reqType="post",
url="https://archiveofourown.org/users/login",
logger=logger,
params={
"user[login]": usernameStr,
"user[password]": passwordStr,
"authenticity_token": tempSoup.find(
"input", {"name": "authenticity_token"}
)["value"],
},
session=session,
)
if pickleFilepath:
os.makedirs(os.path.dirname(pickleFilepath), exist_ok=True)
with open(pickleFilepath, "wb") as file:
pickle.dump(session, file)
return session
else:
session = None
return session
def matchTag(
tagDict: dict,
category: str,
tag,
) -> bool:
if category in ("rating", "lang"):
return bool(tag == tagDict[category])
if category in ("warning", "category", "fandom", "ship", "char", "tag", "author"):
for i in tagDict[category]:
if tag in i:
return True
return False
if category == "complete":
return bool(str(tag)[:1].lower() == str(tagDict["complete"])[:1].lower())
def matchStat(
tagDict: dict,
stat: str,
value: int,
mode: str,
):
if str(mode)[:1].lower() in (1, "1", "g"):
return bool(tagDict[stat] < value)
else:
return bool(tagDict[stat] > value)
def processNode(
tagDict: dict,
node: dict,
) -> bool:
if node:
match node.get("op", "and"): # 'op' = 'operator'
case "matchTag":
return matchTag(
tagDict=tagDict, category=node["category"], tag=node["tag"]
)
case "matchStat":
return matchStat(
tagDict=tagDict,
stat=node["stat"],
value=node["value"],
mode=node["greater"],
)
case "not":
return not processNode(tagDict=tagDict, node=node["children"][0])
case "or":
valid = False
counter = 0
while (not valid) and (counter < len(node["children"])):
valid = processNode(tagDict=tagDict, node=node["children"][counter])
counter += 1
del counter
return valid
case "and":
valid = True
for child in node["children"]:
if not processNode(tagDict=tagDict, node=child):
valid = False
return valid
else:
return True
def mainSingle(
workID: int,
session: requests.Session,
config: dict,
logger: common.logging.Logger,
dryRun: bool = False,
whitelist: dict = {},
) -> bool:
global countWorksTotal
global countWorksStarted
global countWorksComplete
global countWorksNoChange
global countWorkskins
countWorksStarted += 1
logger.info(f"Starting work [{workID}] - [{countWorksStarted} / {countWorksTotal}]")
try:
workSoup = bs4.BeautifulSoup(
getPage(
url=f"https://archiveofourown.org/works/{workID}?view_adult=true",
logger=logger,
session=session,
tryAnon=(not config["ao3DoLoginAlways"]),
),
"lxml",
)
except Exception as ex:
logger.error(f"Work webpage [{workID}] failed with error [{type(ex).__name__}]")
raise ex
else:
logger.debug(f"Work [{workID}] got work webpage successfully")
tagDict = getTagsFromWebsiteSoup(soup=workSoup, logger=logger)
rawFullFilename = os.path.join(config["dirRaws"], f"{workID}.html")
if (
os.path.exists(rawFullFilename)
and common.getMetadataFromRawCommentByID(workID, config)[1]
== tagDict["dateEdited"]
):
logger.info(
f"Work [{workID}]'s previously downloaded raw is still up-to-date [{tagDict['dateEdited']}]"
)
countWorksNoChange += 1
return False
if processNode(tagDict, whitelist):
logger.debug(f"Work [{workID}] passed the blacklist")
else:
logger.warning(f"Work [{workID}] failed the blacklist")
common.args.batch_out.write(f"{workID}\n")
return False
workskin = getWorkskinFromWebsiteSoup(soup=workSoup, logger=logger)
if workskin:
countWorkskins += 1
if not dryRun:
with open(os.path.join(config["dirRaws"], f"{workID}.css"), "w") as file:
file.write(workskin)
raw = getPage(
url=getRawUrlFromWebsiteSoup(soup=workSoup, logger=logger),
logger=logger,
session=session,
tryAnon=(not config["ao3DoLoginAlways"],),
)
if not dryRun:
with open(rawFullFilename, "wb") as file:
file.write(raw)
headerStr = f"""
<!--
{common.commentHeader}
{common.commentVersion}{common.version}
{common.commentTimestampEdited}{tagDict['dateEdited']}
{common.commentTimestampDownloaded}{int(datetime.datetime.now().timestamp())}
-->
"""
file.write(headerStr.encode("utf-8"))
countWorksComplete += 1
logger.info(
f"Completed work [{workID}] - [{countWorksComplete + countWorksNoChange} / {countWorksTotal}]"
)
del tagDict
del workskin
del raw
return True
if __name__ == "__main__":
common.init(
description="Download 'Raw' HTML files from AO3 to be processed later.",
json="r",
dryRun=True,
secondOut="WorkIDs of Works that fail the blacklist.",
all="Check for updates for all Raws and all works in database",
auto="Check for updates for all Raws and all works in database not marked as complete",
)
if common.args.all or common.args.auto:
if common.args.all:
workIDs = getWorkIDs(common.config, common.logger)
else:
workIDs = getCompleteWorkIDs(common.config, common.logger)
for filename in os.listdir(common.config["dirRaws"]):
if filename[-5:] == ".html" and filename[:-5].isdecimal():
workID = int(filename[:-5])
else:
workIDs = set(())
inRaw = common.parseInfile(common.args.infile)
for line in inRaw:
try:
workIDs.add(int(line))
except ValueError:
common.logger.error(f"Could not convert input [{line}] to int")
workIDs.difference_update(getExcluded())
countWorksTotal = len(workIDs)
countWorksStarted = 0
countWorksComplete = 0
countWorksNoChange = 0
countWorkskins = 0
common.logger.info(f"Got [{countWorksTotal}] workIDs")
session = login(
config=common.config,
logger=common.logger,
)
if common.args.json:
jsonFile = common.args.json[0]
whitelist = json.load(common.args.json[0])
else:
whitelist = {}
futuresNEO = {}
with concurrent.futures.ThreadPoolExecutor(
max_workers=10, thread_name_prefix=common.threadNameBulk
) as pool:
for workID in workIDs:
futuresNEO[workID] = pool.submit(
mainSingle,
workID=workID,
session=session,
config=common.config,
logger=common.logger,
dryRun=common.args.dry_run,
whitelist=whitelist,
)
common.logger.log(
(20 + (10 * int(countWorksTotal != (countWorksComplete + countWorksNoChange)))),
f"Completed [{countWorksComplete}] (downloaded) + [{countWorksNoChange}] (no change since last download) (total [{
countWorksNoChange + countWorksComplete}]) of [{countWorksTotal}] works",
)
common.logger.info(f"Got [{countWorkskins}] workskins")
for workID in futuresNEO:
try:
result = futuresNEO[workID].result()
except common.UnavailablePageError:
common.logger.error(
f"Work [{workID}] raised UnavailablePageError: The work seems to have been deleted or purposefully made inaccessable by the author."
)
except Exception as ex:
common.logger.error(
f"Work [{workID}] raised [{type(ex).__name__}]: [{ex.args}]"
)
traceback.print_exception(ex, file=common.args.outfile)
common.logger.info("All operations complete, download.py exiting.")