wrolpi/docker/archive/main.py

#! /usr/bin/env python3
"""
This file is a simple Python/Sanic wrapper around the single-file CLI command.
"""
import asyncio
import base64
import gzip
import io
import json
import logging
import os
import os.path
import pathlib
import re
import secrets
import signal
import subprocess
import sys
import tempfile
import traceback
import zipfile
from json import JSONDecodeError

from sanic import Sanic, response
from sanic.request import Request

# Log using datetime and log level.  Log to stdout.
logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)
ch = logging.StreamHandler()
formatter = logging.Formatter('[%(asctime)s] [%(name)s] [%(levelname)s] %(message)s')
ch.setFormatter(formatter)
logger.addHandler(ch)

def find_singlefile_deno_script() -> pathlib.Path:
    """Find the Deno entry script of the single-file-cli package.

    single-file-cli is designed for Deno; the npm bin (single-file-node.js) runs under node which
    cannot reach the browser CDP on IPv6-enabled hosts and never exits after a successful save."""
    candidates = (
        '/usr/src/app/node_modules/single-file-cli/single-file',
        '/usr/local/lib/node_modules/single-file-cli/single-file',
        '/usr/lib/node_modules/single-file-cli/single-file',
        # Resolve the bin symlink back to the package directory.
        '/usr/src/.nvm/versions/node/v18.19.0/bin/single-file',
        '/usr/local/bin/single-file',
    )
    for candidate in candidates:
        path = pathlib.Path(candidate)
        if path.is_file():
            script = path.resolve().parent / 'single-file'
            try:
                with script.open('rb') as fh:
                    if b'deno run' in fh.readline():
                        return script
            except OSError:
                pass
    raise FileNotFoundError("Can't find single-file Deno script!")


SINGLEFILE_PATH = find_singlefile_deno_script()

DENO_PATH = pathlib.Path('/usr/local/bin/deno')
if not DENO_PATH.is_file():
    raise FileNotFoundError("Can't find deno executable!")

# single-file pipes the browser's stderr but never reads it; under Deno >= 2.3 that kills the
# browser during startup.  The wrapper discards the browser's stderr.
BROWSER_WRAPPER = pathlib.Path('/app/singlefile_browser_wrapper.sh')
if not BROWSER_WRAPPER.is_file():
    raise FileNotFoundError("Can't find singlefile_browser_wrapper.sh!")

BROWSER_EXEC = pathlib.Path('/usr/bin/google-chrome')
if not BROWSER_EXEC.is_file():
    print('Unable to find browser!', file=sys.stderr)
    sys.exit(1)


# Increase response timeout, archiving can take several minutes.
RESPONSE_TIMEOUT = 10 * 60
# The API sends uploaded singlefiles base64-encoded in JSON (a 33% size increase).  Sanic's default
# 100MB limit rejects large uploads with a 413 before the handler runs.
REQUEST_MAX_SIZE = 500 * 1024 * 1024
config = {
    'RESPONSE_TIMEOUT': RESPONSE_TIMEOUT,
    'REQUEST_MAX_SIZE': REQUEST_MAX_SIZE,
}

app = Sanic('archive')
app.update_config(config)

index_html = '''
<html>
<body>
<h2>This is the WROLPi archiving service.</h2>
<p>
    This is not meant to be used directly, it should be called by the
    archive module.
<p>
</body>
</html>
'''


@app.get('/')
async def index(_):
    return response.html(index_html)


async def check_output(cmd, always_log_stderr: bool = False, cwd=None, timeout: float = None):
    proc = None
    try:
        # start_new_session so the whole process group can be killed; single-file does not close
        # the browser it spawns.
        proc = await asyncio.create_subprocess_shell(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, cwd=cwd,
                                                     start_new_session=True)
        stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=timeout)
        if always_log_stderr is True or proc.returncode != 0:
            # Always log when requested, or when the call failed.
            for line in stderr.decode().splitlines():
                logger.error(line)
        return stdout, stderr, proc.returncode
    except asyncio.TimeoutError:
        proc.kill()
        await proc.wait()  # Wait for the process to fully terminate
        raise TimeoutError(f"Command '{cmd}' timed out after {timeout} seconds")
    finally:
        if proc is not None:
            # Kill anything the command left behind (e.g. the browser single-file spawns).
            try:
                os.killpg(proc.pid, signal.SIGKILL)
                logger.info(f'Killed lingering process group {proc.pid}')
            except ProcessLookupError:
                pass  # Nothing left behind; the common case.
            except PermissionError as e:
                # Do not let the kill mask the real failure (e.g. a TimeoutError in flight).
                logger.error(f'Failed to kill process group {proc.pid}', exc_info=e)


async def run_single_file(input_url: str, output_path: pathlib.Path, compress: bool):
    """Run single-file, return (return_code, output_bytes, stderr).

    single-file writes the result to a file; Deno >= 2.3 truncates large --dump-content output
    at exit, which corrupted compressed (ZIP) singlefiles."""
    cmd = f'WROLPI_BROWSER={BROWSER_EXEC}' \
          f' {DENO_PATH} run --node-modules-dir=manual' \
          ' --allow-read --allow-write --allow-net --allow-env --allow-run' \
          f' {SINGLEFILE_PATH}' \
          f' --browser-executable-path {BROWSER_WRAPPER}' \
          f'{" --compress-content " if compress else ""}' \
          f' "{input_url}"' \
          f' "{output_path}"'
    logger.debug(f'archive cmd: {cmd}')
    stdout, stderr, return_code = await check_output(cmd, always_log_stderr=True, timeout=3 * 60)
    output = output_path.read_bytes() if output_path.is_file() else b''
    return return_code, output, stderr


async def call_single_file(url, compress: bool = False) -> (bytes, bytes):
    """
    Call the CLI command for SingleFile.  Returns (singlefile, page_html): the singlefile to
    store, and the uncompressed page for readability/screenshot.

    The page is always captured uncompressed; when compress is set, a second single-file run
    converts the captured file into a compressed (SingleFileZ) self-extracting file (a local
    file:// operation, the page is not fetched again).  If the conversion fails, the
    uncompressed singlefile is returned instead.

    See https://github.com/gildas-lormeau/SingleFile
    """
    logger.info(f'archiving {url}')
    with tempfile.TemporaryDirectory(prefix='singlefile-') as tmp_dir:
        page_path = pathlib.Path(tmp_dir) / 'page.html'
        return_code, page_html, stderr = await run_single_file(url, page_path, compress=False)
        if return_code != 0 or not page_html:
            raise RuntimeError(f'Failed to single-file {url} got the following error: {stderr}')

        singlefile = page_html
        if compress:
            compressed_path = pathlib.Path(tmp_dir) / 'compressed.html'
            return_code, compressed, stderr = await run_single_file(
                f'file://{page_path}', compressed_path, compress=True)
            if return_code == 0 and zipfile.is_zipfile(io.BytesIO(compressed)):
                singlefile = compressed
            else:
                logger.error(f'Failed to compress singlefile ({return_code=}), storing it uncompressed: {url}')

    logger.debug(f'done archiving for {url}')
    return singlefile, page_html


# The prefix and payload of a base64 data URI.  Keep in sync with modules/archive/lib.py.
BASE64_DATA_URI = re.compile(rb'(data:[^,;\s"\'()]*(?:;[^,;\s"\'()]*)*;base64,)([A-Za-z0-9+/]+=*)')
DATA_URI_PLACEHOLDER = re.compile(r'WROLPIDATAURI[0-9a-f]{160}N\d+X')


def strip_data_uris(html: bytes) -> (bytes, dict):
    """Replace each base64 data URI payload with a placeholder; readability-extractor takes minutes on a
    page of large inline images (SingleFile inlines every image).  Returns the HTML and the payloads by
    placeholder.  The placeholder holds a random token, so text in the page cannot be mistaken for one."""
    # At least 133 characters: Readability drops a shorter base64 `src` as a lazy-load placeholder when another
    # attribute names an image file (e.g. alt="cat.png").
    token = secrets.token_hex(80)
    payloads = dict()

    def replace(match):
        placeholder = f'WROLPIDATAURI{token}N{len(payloads)}X'
        payloads[placeholder] = match.group(2)
        return match.group(1) + placeholder.encode()

    return BASE64_DATA_URI.sub(replace, html), payloads


def restore_data_uris(readability: dict, payloads: dict) -> dict:
    """Put the payloads removed by `strip_data_uris` back into the readability output."""
    for key, value in readability.items():
        if isinstance(value, str):
            readability[key] = DATA_URI_PLACEHOLDER.sub(
                lambda m: payloads[m.group(0)].decode() if m.group(0) in payloads else m.group(0), value)
    return readability


async def extract_readability(path: str, url: str) -> dict:
    """
    Call the CLI command for readability-extractor.

    See https://github.com/ArchiveBox/readability-extractor
    """
    logger.info(f'readability for {url}')
    html, data_uri_payloads = strip_data_uris(pathlib.Path(path).read_bytes())
    path = pathlib.Path(path).with_name('readability-input.html')
    path.write_bytes(html)
    cmd = f'readability-extractor {path} "{url}"'
    logger.debug(f'readability cmd: {cmd}')
    stdout, stderr, return_code = await check_output(cmd, timeout=3 * 60)
    try:
        output = json.loads(stdout)
    except JSONDecodeError as e:
        stderr = stderr.decode() if stderr else None
        raise RuntimeError(f'Failed to extract readability.  {stderr}') from e
    logger.debug(f'done readability for {url}')
    return restore_data_uris(output, data_uri_payloads)


async def take_screenshot(url: str) -> bytes:
    cmd = f'{BROWSER_EXEC}' \
          ' --headless' \
          ' --disable-gpu' \
          ' --no-sandbox' \
          ' --screenshot' \
          ' --window-size=1280x720' \
          f' "{url}"'
    logger.debug(f'Screenshot cmd: {cmd}')
    with tempfile.TemporaryDirectory() as tmp_dir:
        try:
            stdout, stderr, return_code = await check_output(cmd, cwd=tmp_dir, timeout=1 * 60)
            if return_code != 0:
                raise ValueError(f'Screenshot failed {return_code=}')
        except Exception as e:
            logger.error(f'Failed to screenshot {url}', exc_info=e)
            return b''

        path = pathlib.Path(f'{tmp_dir}/screenshot.png')
        if not path.is_file():
            logger.warning(f'Screenshot command did not create screenshot of {url}')
            return b''

        size = os.path.getsize(path)
        logger.info(f'Successful screenshot of {url} ({size} bytes) at {path}')
        png = path.read_bytes()
        return png


def extract_compressed_singlefile(singlefile: bytes, tmp_dir: str) -> str:
    """A compressed (SingleFileZ) singlefile is also a valid ZIP containing the uncompressed page as
    index.html (plus its resources).  Extract it so readability/screenshot can use the real HTML.

    Returns the path to the extracted index.html, or an empty string if the singlefile is not compressed.
    """
    buf = io.BytesIO(singlefile)
    if not zipfile.is_zipfile(buf):
        return ''
    try:
        with zipfile.ZipFile(buf) as zip_:
            if 'index.html' not in zip_.namelist():
                return ''
            extracted = pathlib.Path(tmp_dir) / 'extracted'
            extracted.mkdir(exist_ok=True)
            zip_.extractall(extracted)
            return str(extracted / 'index.html')
    except zipfile.BadZipFile:
        return ''


def prepare_bytes(b: bytes) -> str:
    """
    Compress and encode bytes for smaller response.
    """
    b = gzip.compress(b)
    b = base64.b64encode(b)
    b = b.decode()
    return b


@app.post('/screenshot')
async def post_screenshot(request: Request):
    """Generate a screenshot for the provided singlefile."""
    url = request.json['url']
    singlefile = request.json.get('singlefile')

    try:
        logger.info(f'Generating screenshot for {url}')

        # Decode and decompress the singlefile
        if singlefile:
            singlefile = base64.b64decode(singlefile)
            singlefile = gzip.decompress(singlefile)

        if not singlefile:
            raise ValueError(f'No singlefile provided for {url}')

        # Write singlefile to temp file and screenshot it
        # Use html suffix so chrome screenshot recognizes it as an HTML file
        with tempfile.TemporaryDirectory() as tmp_dir:
            singlefile_path = pathlib.Path(tmp_dir) / 'singlefile.html'
            singlefile_path.write_bytes(singlefile)

            # A compressed (SingleFileZ) singlefile must be extracted before it can be screenshot.
            extractable = extract_compressed_singlefile(singlefile, tmp_dir) or str(singlefile_path)

            screenshot = None
            try:
                # Screenshot the local singlefile
                screenshot = await take_screenshot(f'file://{extractable}')
            except Exception as e:
                logger.error(f'Failed to take screenshot of {extractable}', exc_info=e)

            # Fall back to URL if local screenshot failed
            if not screenshot:
                logger.warning(f'Failed to screenshot local singlefile, attempting to screenshot URL: {url}')
                try:
                    screenshot = await take_screenshot(url)
                except Exception as e:
                    logger.error(f'Failed to take screenshot of {url}', exc_info=e)

        if not screenshot:
            raise ValueError(f'Failed to generate screenshot for {url}')

        # Compress for smaller response
        screenshot = prepare_bytes(screenshot)

        ret = dict(
            url=url,
            screenshot=screenshot,
        )
        return response.json(ret)
    except Exception as e:
        logger.error(f'Failed to generate screenshot for {url}', exc_info=e)
        error = str(traceback.format_exc())
        return response.json({'error': f'Failed to generate screenshot for {url} traceback is below... \n\n {error}'})


@app.post('/json')
async def post_archive(request: Request):
    url = request.json['url']

    try:
        # May have been passed the singlefile contents, do the extractions and return.
        singlefile = request.json.get('singlefile')
        singlefile = base64.b64decode(singlefile) if singlefile else None
        compress = bool(request.json.get('compress'))

        if singlefile:
            # A provided singlefile may be compressed (SingleFileZ); readability and screenshot
            # must use the extracted index.html.
            page_html = None
        else:
            # Fetch and create from the internet; readability/screenshot use the uncompressed page.
            singlefile, page_html = await call_single_file(url, compress=compress)

        # Use html suffix so chrome screenshot recognizes it as an HTML file.
        with tempfile.TemporaryDirectory() as tmp_dir:
            page_path = pathlib.Path(tmp_dir) / 'page.html'
            if page_html:
                page_path.write_bytes(page_html)
                extractable = str(page_path)
            else:
                page_path.write_bytes(singlefile)
                extractable = extract_compressed_singlefile(singlefile, tmp_dir) or str(page_path)

            readability = None
            try:
                readability = await extract_readability(extractable, url)
            except Exception as e:
                # Readability had error, but its is not required.
                logger.error(f'Failed to get readability', exc_info=e)

            screenshot = None
            try:
                # Screenshot the local singlefile, if that fails try the URL.
                screenshot = await take_screenshot(f'file://{extractable}')
            except Exception as e:
                logger.error(f'Failed to take screenshot of {extractable}', exc_info=e)

            if not screenshot:
                logger.warning(f'Failed to screenshot local singlefile attempting to screenshot: {url}')
                try:
                    screenshot = await take_screenshot(url)
                except Exception as e:
                    logger.error(f'Failed to take screenshot of {url}', exc_info=e)

        # Compress for smaller response.
        singlefile = prepare_bytes(singlefile)
        screenshot = prepare_bytes(screenshot)

        ret = dict(
            url=url,
            singlefile=singlefile,
            readability=readability,
            screenshot=screenshot,
        )
        return response.json(ret)
    except Exception as e:
        logger.error(f'Failed to archive {url}', exc_info=e)
        error = str(traceback.format_exc())
        return response.json({'error': f'Failed to archive {url} traceback is below... \n\n {error}'})


if __name__ == '__main__':
    app.run('0.0.0.0', 8080, workers=4, auto_reload=True)