mirror of
https://codeberg.org/Cyberpro123/AO3-DL.git
synced 2026-07-23 06:35:58 +00:00
236 lines
7.9 KiB
Python
Executable file
236 lines
7.9 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 os.path
|
|
import traceback
|
|
import urllib.parse
|
|
|
|
import bs4
|
|
|
|
import common
|
|
import download
|
|
from process import makeStr
|
|
import sql
|
|
|
|
numToKey = {1: "summary", 2: "notes_start", 3: "notes_end", 4: "body"}
|
|
|
|
|
|
def soupToHrefs(soup: bs4.BeautifulSoup) -> tuple[set, set]:
|
|
output = set(())
|
|
prev = set(())
|
|
for img in soup.find_all("img"):
|
|
if "src" in img.attrs:
|
|
if common.imagePrefix in img.attrs["src"]:
|
|
prev.add(img.attrs["src"])
|
|
else:
|
|
output.add(img.attrs["src"])
|
|
return output, prev
|
|
|
|
|
|
def idToHrefs(
|
|
id: int,
|
|
config: dict,
|
|
logger: common.logging.Logger,
|
|
) -> tuple[set, set]:
|
|
sqlStr = (
|
|
f"SELECT chapter, summary, notes_start, notes_end, body FROM text WHERE id={id}"
|
|
)
|
|
result = sql.execute(
|
|
commands=[(sqlStr, ())],
|
|
config=config,
|
|
logger=logger,
|
|
)
|
|
hrefs = set(())
|
|
prev = set(())
|
|
for i in result[0]:
|
|
for num, j in enumerate(i):
|
|
if num == 0:
|
|
chapter = j
|
|
else:
|
|
new, old = soupToHrefs(soup=bs4.BeautifulSoup(j, "lxml"))
|
|
hrefs.update(new)
|
|
prev.update(old)
|
|
logger.info(
|
|
f"Finished searching for links in [{numToKey[num]}] of Work [{id}] chapter [{chapter}]"
|
|
)
|
|
return hrefs, prev
|
|
|
|
|
|
def downloadImage(
|
|
url: str,
|
|
config: dict,
|
|
logger: common.logging.Logger,
|
|
) -> str:
|
|
parsed = urllib.parse.urlparse(url=url)
|
|
path = str(parsed.path).strip("/")
|
|
filePathFullEncoded = os.path.join(
|
|
config["dirImg"],
|
|
parsed.netloc,
|
|
path,
|
|
)
|
|
filePathFullDecoded = urllib.parse.unquote(filePathFullEncoded)
|
|
os.makedirs(os.path.split(filePathFullDecoded)[0], exist_ok=True)
|
|
if not (
|
|
os.path.exists(filePathFullDecoded) and os.path.getsize(filePathFullDecoded)
|
|
):
|
|
rawImg = download.getPage(
|
|
url=url, logger=logger, retries=min(10, common.loopRetries)
|
|
)
|
|
with open(filePathFullDecoded, "wb") as file:
|
|
file.write(rawImg)
|
|
return f"{common.imagePrefix}{os.path.relpath(path=filePathFullEncoded, start=config['dirImg'])}"
|
|
|
|
|
|
def editText(
|
|
id: int,
|
|
hrefs: dict,
|
|
config: dict,
|
|
logger: common.logging.Logger,
|
|
) -> None:
|
|
sqlStr = (
|
|
f"SELECT chapter, summary, notes_start, notes_end, body FROM text WHERE id={id}"
|
|
)
|
|
result = sql.execute(
|
|
commands=[(sqlStr, ())],
|
|
config=config,
|
|
logger=logger,
|
|
)
|
|
commands = []
|
|
for i in result[0]:
|
|
for num, j in enumerate(i):
|
|
if num == 0:
|
|
chapter = j
|
|
elif not (chapter == 0 and numToKey[num] == "body"):
|
|
soup = bs4.BeautifulSoup(f"<html><body>{j}</body></html>", "lxml")
|
|
for img in soup.find_all("img"):
|
|
if (
|
|
"src" in img.attrs
|
|
and img.attrs["src"] in hrefs.keys()
|
|
and hrefs[img.attrs["src"]]
|
|
):
|
|
img.attrs["src"] = hrefs[img.attrs["src"]]
|
|
commands.append(
|
|
(
|
|
f"UPDATE text SET {numToKey[num]} = ? WHERE id = {id} AND chapter = {chapter}",
|
|
(makeStr(soup.body.contents),),
|
|
)
|
|
)
|
|
del soup
|
|
sql.execute(commands=commands, config=config, logger=logger)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
common.init(
|
|
description="Searches text in the database for image links, downloads the linked images, and redirects the links to the new local copies.",
|
|
all="Download all images in every work",
|
|
secondOut="URLs of images that couldn't be downloaded",
|
|
)
|
|
if common.args.all:
|
|
workIDs = sql.getWorkIDs(config=common.config, logger=common.logger)
|
|
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")
|
|
hrefs = {}
|
|
previousDownloads = set(())
|
|
futures1 = {}
|
|
with concurrent.futures.ThreadPoolExecutor(
|
|
max_workers=10, thread_name_prefix=common.threadNameBulk
|
|
) as pool1:
|
|
for id in workIDs:
|
|
futures1[id] = pool1.submit(
|
|
idToHrefs,
|
|
id=id,
|
|
config=common.config,
|
|
logger=common.logger,
|
|
)
|
|
common.logger.info("Multitasking Pool 1 complete, getting results...")
|
|
for id in futures1:
|
|
try:
|
|
result = futures1[id].result()
|
|
except Exception as ex:
|
|
common.logger.error(
|
|
f"Work [{id}] raised [{type(ex).__name__}] while looking for image links"
|
|
)
|
|
traceback.print_exception(ex, file=common.args.outfile)
|
|
else:
|
|
for href in result[0]:
|
|
hrefs[href] = ""
|
|
previousDownloads.update(result[1])
|
|
del futures1
|
|
common.logger.info(
|
|
f"Found [{len(hrefs)}] undownloaded image links and [{len(previousDownloads)}] previously downloaded image links"
|
|
)
|
|
futures2 = {}
|
|
with concurrent.futures.ThreadPoolExecutor(
|
|
max_workers=10, thread_name_prefix=common.threadNameBulk
|
|
) as pool2:
|
|
for href in hrefs:
|
|
futures2[href] = pool2.submit(
|
|
downloadImage,
|
|
url=href,
|
|
config=common.config,
|
|
logger=common.logger,
|
|
)
|
|
common.logger.info("Multitasking Pool 2 complete, getting results...")
|
|
successfulDownloads = 0
|
|
for href in futures2:
|
|
try:
|
|
result = futures2[href].result()
|
|
except (common.LoopFailError, common.UnavailablePageError):
|
|
common.logger.debug(f"Image [{href}] failed to download")
|
|
common.args.batch_out.write(f"{href}\n")
|
|
except Exception as ex:
|
|
common.logger.error(
|
|
f"Image [{href}] raised [{type(ex).__name__}] while trying to download"
|
|
)
|
|
traceback.print_exception(ex, file=common.args.outfile)
|
|
else:
|
|
hrefs[href] = result
|
|
successfulDownloads += 1
|
|
del futures2
|
|
common.logger.info(
|
|
f"Successfully downloaded [{successfulDownloads}] images (out of [{len(hrefs)}])"
|
|
)
|
|
futures3 = {}
|
|
with concurrent.futures.ThreadPoolExecutor(
|
|
max_workers=10, thread_name_prefix=common.threadNameBulk
|
|
) as pool3:
|
|
for id in workIDs:
|
|
futures3[id] = pool3.submit(
|
|
editText,
|
|
id=id,
|
|
hrefs=hrefs,
|
|
config=common.config,
|
|
logger=common.logger,
|
|
)
|
|
common.logger.info("Multitasking Pool 3 complete, getting results...")
|
|
errors = 0
|
|
for id in futures3:
|
|
try:
|
|
result = futures3[id].result()
|
|
except Exception as ex:
|
|
errors += 1
|
|
common.logger.error(
|
|
f"Work [{id}] raised [{type(ex).__name__}] while editing links"
|
|
)
|
|
traceback.print_exception(ex, file=common.args.outfile)
|
|
del futures3
|
|
common.logger.info(
|
|
f"Got [{errors}] errors while replacing external image links with local image links"
|
|
)
|