#!/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"{j}", "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" )