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 json
import logging
import os.path
import pathlib
import subprocess
import tempfile

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)

SINGLEFILE_PATH = pathlib.Path('/usr/src/app/node_modules/single-file-cli/single-file')
if not SINGLEFILE_PATH.is_file():
    SINGLEFILE_PATH = pathlib.Path('/usr/src/app/node_modules/single-file/cli/single-file')
if not SINGLEFILE_PATH.is_file():
    raise FileNotFoundError("Can't find single-file executable!")

# Increase response timeout, archiving can take several minutes.
RESPONSE_TIMEOUT = 10 * 60
config = {
    'RESPONSE_TIMEOUT': RESPONSE_TIMEOUT,
}

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):
    proc = await asyncio.create_subprocess_shell(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, cwd=cwd)
    stdout, stderr = await proc.communicate()
    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


async def call_single_file(url) -> bytes:
    """
    Call the CLI command for SingleFile.

    See https://github.com/gildas-lormeau/SingleFile
    """
    logger.info(f'archiving {url}')
    cmd = f'{SINGLEFILE_PATH}' \
          r' --browser-executable-path /usr/bin/chromium-browser' \
          r' --browser-args [\"--no-sandbox\"]' \
          r' --dump-content ' \
          f' "{url}"'
    logger.debug(f'archive cmd: {cmd}')
    stdout, stderr, return_code = await check_output(cmd, always_log_stderr=True)
    if return_code != 0 or not stdout:
        raise ValueError(f'Failed to single-file {url}')
    logger.debug(f'done archiving for {url}')
    return stdout


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}')
    cmd = f'readability-extractor {path} "{url}"'
    logger.debug(f'readability cmd: {cmd}')
    stdout, stderr, return_code = await check_output(cmd)
    output = json.loads(stdout)
    logger.debug(f'done readability for {url}')
    return output


async def take_screenshot(url: str) -> bytes:
    logger.info(f'Screenshot: {url}')
    cmd = '/usr/bin/chromium-browser' \
          ' --headless' \
          ' --disable-gpu' \
          ' --no-sandbox' \
          ' --screenshot' \
          ' --window-size=1280,720' \
          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)
            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():
            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 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('/json')
async def post_archive(request: Request):
    url = request.json['url']
    try:
        singlefile, screenshot = await asyncio.gather(call_single_file(url), take_screenshot(url))
        with tempfile.NamedTemporaryFile('wb') as fh:
            fh.write(singlefile)
            readability = await extract_readability(fh.name, url)

        # 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)
        return response.json({'error': f'Failed to archive {url} error was... \n\n {e}'})


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