webdrop-bridge/src/webdrop_bridge/ui/restricted_web_view.py

"""Restricted web view with URL whitelist enforcement for Kiosk-mode."""

import fnmatch
import logging
from pathlib import Path
from typing import List, Optional, Union

from PySide6.QtCore import QStandardPaths, QUrl
from PySide6.QtGui import QDesktopServices
from PySide6.QtWebEngineCore import QWebEngineNavigationRequest, QWebEnginePage, QWebEngineProfile
from PySide6.QtWebEngineWidgets import QWebEngineView

logger = logging.getLogger(__name__)


logger = logging.getLogger(__name__)


class CustomWebEnginePage(QWebEnginePage):
    """Custom page that handles new window requests for downloads."""

    def acceptNavigationRequest(
        self, url: Union[QUrl, str], nav_type: QWebEnginePage.NavigationType, is_main_frame: bool
    ) -> bool:
        """Handle navigation requests, including download links.

        Args:
            url: Target URL (QUrl or string)
            nav_type: Type of navigation (link click, form submit, etc.)
            is_main_frame: Whether this is the main frame

        Returns:
            True to accept navigation, False to reject
        """
        # Convert to string if QUrl
        url_str = url.toString() if isinstance(url, QUrl) else url

        # Log all navigation attempts for debugging
        logger.debug(f"Navigation request: {url_str} (type={nav_type}, main_frame={is_main_frame})")

        # Check if this might be a download (common file extensions)
        download_extensions = [
            ".pdf",
            ".zip",
            ".rar",
            ".7z",
            ".tar",
            ".gz",
            ".doc",
            ".docx",
            ".xls",
            ".xlsx",
            ".ppt",
            ".pptx",
            ".jpg",
            ".jpeg",
            ".png",
            ".gif",
            ".bmp",
            ".svg",
            ".mp4",
            ".mp3",
            ".avi",
            ".mov",
            ".wav",
            ".exe",
            ".msi",
            ".dmg",
            ".pkg",
            ".csv",
            ".txt",
            ".json",
            ".xml",
        ]

        if any(url_str.lower().endswith(ext) for ext in download_extensions):
            logger.debug(f"Detected potential download URL: {url_str}")
            # This will trigger downloadRequested if it's a download

        return super().acceptNavigationRequest(url, nav_type, is_main_frame)

    def createWindow(self, window_type: QWebEnginePage.WebWindowType) -> QWebEnginePage:
        """Handle new window requests (target=_blank, window.open, etc.).

        Many downloads are triggered via target="_blank" links.

        Args:
            window_type: Type of window being created

        Returns:
            New page instance for the window
        """
        logger.debug(f"New window requested, type: {window_type}")

        # Create a temporary page to handle the download
        # This page will never be displayed but allows downloads to work
        download_page = QWebEnginePage(self.profile(), self)

        logger.debug("Created temporary page for download/popup")

        # Return the temporary page - it will trigger downloadRequested if it's a download
        return download_page


class RestrictedWebEngineView(QWebEngineView):
    """Web view that enforces URL whitelist for Kiosk-mode security.

    If allowed_urls is empty, no restrictions are applied.
    If allowed_urls is not empty, only matching URLs are loaded in the view.
    Non-matching URLs open in the system default browser.
    """

    def __init__(self, allowed_urls: Optional[List[str]] = None):
        """Initialize the restricted web view.

        Args:
            allowed_urls: List of allowed URL patterns (empty = no restriction)
                         Patterns support wildcards: *.example.com, localhost, etc.
        """
        super().__init__()
        self.allowed_urls = allowed_urls or []

        # Create persistent profile for cookie and session storage
        self.profile = self._create_persistent_profile()

        # Use custom page for better download handling with persistent profile
        custom_page = CustomWebEnginePage(self.profile, self)
        self.setPage(custom_page)

        logger.debug(
            "RestrictedWebEngineView initialized with CustomWebEnginePage and persistent profile"
        )

        # Connect to navigation request handler
        self.page().navigationRequested.connect(self._on_navigation_requested)

    def _create_persistent_profile(self) -> QWebEngineProfile:
        """Create and configure a persistent web engine profile.

        This enables persistent cookies and cache storage, allowing
        authentication sessions (e.g., Microsoft login) to persist
        across application restarts.

        Returns:
            Configured QWebEngineProfile with persistent storage
        """
        # Get application data directory
        app_data_dir = QStandardPaths.writableLocation(
            QStandardPaths.StandardLocation.AppDataLocation
        )

        # Create profile directory path
        profile_path = Path(app_data_dir) / "WebEngineProfile"
        profile_path.mkdir(parents=True, exist_ok=True)

        # Create persistent profile with custom storage location
        # Using "WebDropBridge" as the profile name
        # Note: No parent specified so we control the lifecycle
        profile = QWebEngineProfile("WebDropBridge")
        profile.setPersistentStoragePath(str(profile_path))

        # Configure persistent cookies (critical for authentication)
        profile.setPersistentCookiesPolicy(
            QWebEngineProfile.PersistentCookiesPolicy.ForcePersistentCookies
        )

        # Enable HTTP cache for better performance
        profile.setHttpCacheType(QWebEngineProfile.HttpCacheType.DiskHttpCache)

        # Set cache size to 100 MB
        profile.setHttpCacheMaximumSize(100 * 1024 * 1024)

        logger.debug(f"Created persistent profile at: {profile_path}")
        logger.debug("Cookies policy: ForcePersistentCookies")
        logger.debug("HTTP cache: DiskHttpCache (100 MB)")

        return profile

    def _on_navigation_requested(self, request: QWebEngineNavigationRequest) -> None:
        """Handle navigation requests and enforce URL whitelist.

        Args:
            request: Navigation request to process
        """
        url = request.url()

        # If no restrictions, allow all URLs
        if not self.allowed_urls:
            return

        # Check if URL matches whitelist
        if self._is_url_allowed(url):
            # Allow the navigation (default behavior)
            return

        # URL not whitelisted - open in system browser
        request.reject()
        QDesktopServices.openUrl(url)

    def _is_url_allowed(self, url: QUrl) -> bool:
        """Check if a URL matches the whitelist.

        Supports:
        - Exact domain matches: example.com
        - Wildcard patterns: *.example.com
        - Localhost variations: localhost, 127.0.0.1
        - File URLs: file://...

        Args:
            url: QUrl to check

        Returns:
            True if URL is allowed, False otherwise
        """
        url_str = url.toString()
        host = url.host()
        scheme = url.scheme()

        # Allow file:// URLs (local webapp)
        if scheme == "file":
            return True

        # If no whitelist, allow all URLs
        if not self.allowed_urls:
            return True

        # Check against whitelist patterns
        for pattern in self.allowed_urls:
            # Check full URL pattern
            if fnmatch.fnmatch(url_str, pattern):
                return True

            # Check host pattern (for http/https URLs)
            if scheme in ("http", "https") and host:
                if fnmatch.fnmatch(host, pattern):
                    return True

                # Also check with www prefix variants
                if fnmatch.fnmatch(f"www.{host}", pattern):
                    return True
                if fnmatch.fnmatch(host.replace("www.", ""), pattern):
                    return True

        return False