Creating Generate Screenshot button for Archives which are missing a screenshot.

Author: lrnselfreliance

app/src/Events.js

@@ -105,6 +105,19 @@ function handleEvents(events) {
             eventToast('Archive Upload Failed!', message, 'error', 5000);
         }
 
+        if (event === 'screenshot_generated') {
+            eventToast(
+                'Screenshot Generated',
+                message,
+                'success',
+                5000,
+                () => window.open(url, '_self'));
+        }
+
+        if (event === 'screenshot_generation_failed') {
+            eventToast('Screenshot Generation Failed!', message, 'error', 5000);
+        }
+
         if (subject) {
             newestEvents[subject] = dt;
         }

app/src/api.js

@@ -750,6 +750,28 @@ export async function getArchive(archiveId) {
     }
 }
 
+export async function generateArchiveScreenshot(archiveId) {
+    const response = await apiPost(`${ARCHIVES_API}/${archiveId}/generate_screenshot`);
+    if (response.ok) {
+        toast({
+            type: 'success',
+            title: 'Screenshot Generation Queued',
+            description: 'Screenshot generation has been queued. This may take a moment.',
+            time: 3000,
+        });
+        return true;
+    } else {
+        const message = await getErrorMessage(response, 'Failed to queue screenshot generation.');
+        toast({
+            type: 'error',
+            title: 'Screenshot Generation Error',
+            description: message,
+            time: 5000,
+        });
+        return false;
+    }
+}
+
 export async function postDownload(downloadData) {
     if (!downloadData.downloader) {
         toast({

app/src/components/Archive.js

@@ -35,7 +35,7 @@ import {
     textEllipsis,
     useTitle
 } from "./Common";
-import {deleteArchives, postDownload, tagFileGroup, untagFileGroup} from "../api";
+import {deleteArchives, generateArchiveScreenshot, postDownload, tagFileGroup, untagFileGroup} from "../api";
 import {Link, Route, Routes, useNavigate, useParams} from "react-router-dom";
 import Message from "semantic-ui-react/dist/commonjs/collections/Message";
 import {useArchive, useDomains, useSearchArchives, useSearchOrder} from "../hooks/customHooks";
@@ -122,6 +122,14 @@ function ArchivePage() {
         }
     }
 
+    const localGenerateScreenshot = async () => {
+        const success = await generateArchiveScreenshot(data.id);
+        if (success) {
+            // Refresh the archive data after a short delay to show the new screenshot
+            setTimeout(() => fetchArchive(), 2000);
+        }
+    }
+
     const updateButton = <APIButton
         text='Update'
         color='green'
@@ -143,6 +151,17 @@ function ArchivePage() {
     >
         Delete
     </APIButton>;
+    const generateScreenshotButton = !screenshotUrl ? <APIButton
+        text='Generate Screenshot'
+        color='yellow'
+        confirmContent='Generate a screenshot for this archive?'
+        confirmButton='Generate'
+        onClick={localGenerateScreenshot}
+        obeyWROLMode={true}
+        style={{marginTop: '0.5em'}}
+    >
+        Generate Screenshot
+    </APIButton> : null;
 
     let historyList = <Loader active/>;
     if (history && history.length === 0) {
@@ -254,6 +273,7 @@ function ArchivePage() {
             {readButton}
             {updateButton}
             {deleteButton}
+            {generateScreenshotButton}
         </Segment>
 
         <Segment>

docker/archive/main.py

@@ -10,6 +10,7 @@
 import os.path
 import pathlib
 import subprocess
+import sys
 import tempfile
 import traceback
 from json import JSONDecodeError
@@ -33,6 +34,12 @@
 if not SINGLEFILE_PATH.is_file():
     raise FileNotFoundError("Can't find single-file executable!")
 
+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
 config = {
@@ -117,7 +124,7 @@ async def extract_readability(path: str, url: str) -> dict:
 
 
 async def take_screenshot(url: str) -> bytes:
-    cmd = '/usr/bin/google-chrome' \
+    cmd = f'{BROWSER_EXEC}' \
           ' --headless' \
           ' --disable-gpu' \
           ' --no-sandbox' \
@@ -155,6 +162,61 @@ def prepare_bytes(b: bytes) -> str:
     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.NamedTemporaryFile('wb', suffix='.html') as fh:
+            fh.write(singlefile)
+            fh.flush()
+
+            screenshot = None
+            try:
+                # Screenshot the local singlefile
+                screenshot = await take_screenshot(f'file://{fh.name}')
+            except Exception as e:
+                logger.error(f'Failed to take screenshot of {fh.name}', 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']

modules/archive/api.py

@@ -9,8 +9,10 @@
 
 from wrolpi.api_utils import json_response, api_app
 from wrolpi.common import logger, wrol_mode_check, api_param_limiter, TRACE_LEVEL
+from wrolpi.db import get_db_session
 from wrolpi.errors import ValidationError
 from wrolpi.events import Events
+from wrolpi.files.lib import upsert_file
 from wrolpi.schema import JSONErrorResponse
 from wrolpi.switches import register_switch_handler, ActivateSwitchMethod
 from . import lib, schema
@@ -136,6 +138,49 @@ async def singlefile_upload_switch_handler(url=None):
 singlefile_upload_switch_handler: ActivateSwitchMethod
 
 
+@register_switch_handler('generate_screenshot_switch_handler')
+async def generate_screenshot_switch_handler(archive_id=None):
+    """Used by `post_generate_screenshot` to generate screenshots in the background"""
+    q: multiprocessing.Queue = api_app.shared_ctx.archive_screenshots
+
+    trace_enabled = logger.isEnabledFor(TRACE_LEVEL)
+    if trace_enabled:
+        logger.trace(f'generate_screenshot_switch_handler called for archive_id={archive_id}')
+    try:
+        archive_id = q.get_nowait()
+    except queue.Empty:
+        if trace_enabled:
+            logger.trace(f'generate_screenshot_switch_handler called on empty queue')
+        return
+
+    try:
+        q_size = q.qsize()
+    except NotImplementedError:
+        # qsize() is not implemented on macOS
+        q_size = '?'
+    logger.info(f'generate_screenshot_switch_handler queue size: {q_size}')
+
+    try:
+        await lib.generate_archive_screenshot(archive_id)
+        # Always send success event since exceptions are raised on failure
+        from modules.archive import Archive
+        archive = lib.get_archive(archive_id=archive_id)
+        location = archive.location
+        name = archive.file_group.title or archive.file_group.url
+        logger.info(f'Generated screenshot for Archive ({q_size}): {archive_id}')
+        Events.send_screenshot_generated(f'Generated screenshot for: {name}', url=location)
+    except Exception as e:
+        logger.error(f'generate_screenshot_switch_handler failed for Archive {archive_id}', exc_info=e)
+        Events.send_screenshot_generation_failed(f'Failed to generate screenshot: {e}')
+        raise
+
+    # Call this function again so any new screenshot requests can be processed.
+    generate_screenshot_switch_handler.activate_switch()
+
+
+generate_screenshot_switch_handler: ActivateSwitchMethod
+
+
 @archive_bp.post('/upload')
 @openapi.definition(
     summary='Upload SingleFile from SingleFile browser extension and convert it to an Archive.'
@@ -153,3 +198,28 @@ async def post_upload_singlefile(request: Request):
     singlefile_upload_switch_handler.activate_switch(context=dict(url=url))
     # Return empty json response because SingleFile extension expects a JSON response.
     return json_response(dict(), status=HTTPStatus.OK)
+
+
+@archive_bp.post('/<archive_id:int>/generate_screenshot')
+@openapi.description('Generate a screenshot for an Archive that does not have one')
+@openapi.response(HTTPStatus.OK, description='Screenshot generation queued')
+@openapi.response(HTTPStatus.NOT_FOUND, JSONErrorResponse)
+@openapi.response(HTTPStatus.BAD_REQUEST, JSONErrorResponse)
+@wrol_mode_check
+async def post_generate_screenshot(_: Request, archive_id: int):
+    """Queue a screenshot generation request for an Archive."""
+    # Verify archive exists
+    try:
+        archive = lib.get_archive(archive_id=archive_id)
+    except Exception:
+        return json_response({'error': f'Archive {archive_id} not found'}, status=HTTPStatus.NOT_FOUND)
+
+    if not archive.singlefile_path:
+        return json_response({'error': 'Archive has no singlefile'}, status=HTTPStatus.BAD_REQUEST)
+
+    # Queue the screenshot generation request
+    logger.info(f'Queueing screenshot generation for Archive {archive_id}')
+    api_app.shared_ctx.archive_screenshots.put(archive_id)
+    generate_screenshot_switch_handler.activate_switch(context=dict(archive_id=archive_id))
+
+    return json_response({'message': 'Screenshot generation queued'}, status=HTTPStatus.OK)

modules/archive/lib.py

@@ -156,6 +156,42 @@ async def request_archive(url: str, singlefile: str = None) -> Tuple[str, Option
     return singlefile, readability, screenshot
 
 
+async def request_screenshot(url: str, singlefile_path: pathlib.Path) -> Optional[bytes]:
+    """Send a request to the archive service to generate a screenshot from the singlefile."""
+    logger.info(f'Sending screenshot request to archive service: {url}')
+
+    # Read, compress, and encode the singlefile
+    singlefile_contents = singlefile_path.read_bytes()
+    singlefile_compressed = gzip.compress(singlefile_contents)
+    singlefile_b64 = base64.b64encode(singlefile_compressed).decode()
+
+    data = dict(url=url, singlefile=singlefile_b64)
+    try:
+        async with aiohttp_post(f'{ARCHIVE_SERVICE}/screenshot', json_=data, timeout=ARCHIVE_TIMEOUT) as response:
+            status = response.status
+            contents = await response.json()
+        if contents and (error := contents.get('error')):
+            # Report the error from the archive service.
+            raise Exception(f'Received error from archive service: {error}')
+
+        # Compressed base64
+        screenshot = contents.get('screenshot')
+        if not screenshot:
+            logger.warning(f'Failed to get screenshot for {url=}')
+            return None
+
+        logger.debug(f'screenshot request status code {status}')
+    except Exception as e:
+        logger.error('Error when requesting screenshot', exc_info=e)
+        raise
+
+    # Decode and decompress.
+    screenshot = base64.b64decode(screenshot)
+    screenshot = gzip.decompress(screenshot)
+
+    return screenshot
+
+
 async def model_archive_result(url: str, singlefile: str, readability: dict, screenshot: bytes) -> Archive:
     """
     Convert results from ArchiveDownloader into real files.  Create Archive record.
@@ -625,3 +661,74 @@ async def singlefile_to_archive(singlefile: bytes) -> Archive:
     logger.trace(f'singlefile_to_archive modeling: {url}')
     archive: Archive = await model_archive_result(url, singlefile, readability, screenshot)
     return archive
+
+
+async def generate_archive_screenshot(archive_id: int) -> pathlib.Path:
+    """
+    Generate a screenshot for an existing Archive that doesn't have one.
+    If the Archive already has a screenshot, verify it exists and ensure it's tracked in the FileGroup.
+
+    Returns the path to the generated screenshot.
+
+    Raises:
+        ValueError: If Archive has no singlefile
+        RuntimeError: If screenshot generation fails
+    """
+    from wrolpi.db import get_db_session
+
+    with get_db_session() as session:
+        archive = Archive.find_by_id(archive_id, session=session)
+
+        if not archive.singlefile_path:
+            raise ValueError(f'Cannot generate screenshot for Archive {archive_id}: no singlefile')
+
+        # Check if screenshot already exists
+        if archive.screenshot_path:
+            # Verify the screenshot file actually exists on disk
+            if archive.screenshot_path.is_file():
+                logger.info(f'Archive {archive_id} already has a screenshot, ensuring it is tracked')
+                # Ensure the screenshot is tracked in FileGroup.files and FileGroup.data
+                with get_db_session(commit=True) as tracking_session:
+                    archive = Archive.find_by_id(archive_id, session=tracking_session)
+                    file_group = archive.file_group
+                    # append_files uses unique_by_predicate, so this is safe even if already tracked
+                    file_group.append_files(archive.screenshot_path)
+
+                    # Also update FileGroup.data (same pattern as set_screenshot)
+                    data = dict(file_group.data) if file_group.data else {}
+                    data['screenshot_path'] = str(archive.screenshot_path)
+                    file_group.data = data
+
+                    archive.validate()
+                    tracking_session.flush()
+                return archive.screenshot_path
+            else:
+                logger.warning(f'Archive {archive_id} has screenshot_path but file does not exist, regenerating')
+
+        singlefile_path = archive.singlefile_path
+        url = archive.file_group.url
+
+    # Request screenshot from Archive docker service or generate locally
+    if DOCKERIZED:
+        logger.debug(f'Requesting screenshot from archive service for Archive {archive_id}')
+        screenshot_bytes = await request_screenshot(url, singlefile_path)
+    else:
+        logger.debug(f'Generating screenshot locally for Archive {archive_id}')
+        singlefile_contents = singlefile_path.read_bytes()
+        screenshot_bytes = html_screenshot(singlefile_contents)
+
+    if not screenshot_bytes:
+        raise RuntimeError(f'Failed to generate screenshot for Archive {archive_id}')
+
+    # Save screenshot next to singlefile with same naming pattern
+    screenshot_path = singlefile_path.with_suffix('.png')
+    screenshot_path.write_bytes(screenshot_bytes)
+    logger.info(f'Generated screenshot for Archive {archive_id}: {screenshot_path}')
+
+    # Update the Archive to include the new screenshot file
+    with get_db_session(commit=True) as session:
+        archive = Archive.find_by_id(archive_id, session=session)
+        archive.set_screenshot(screenshot_path)
+        session.flush()
+
+    return screenshot_path