AirPlan-V2/lib/air_runtime/xdb_capture.py

"""AirXDB screenshot capture backends with automatic fallback chain.

Three capture methods:
  - KmsGrabCapture: DRM/KMS native screenshot via ffmpeg (no sudo)
  - XvfbCapture: Xvfb virtual framebuffer screenshot
  - FallbackCapture: text placeholder when no GUI capture is available

CaptureManager tries them in order (kms -> xvfb -> fallback) unless
a specific method is requested.
"""

from __future__ import annotations

import shutil
import subprocess
from dataclasses import dataclass
from pathlib import Path


@dataclass
class CaptureResult:
    success: bool
    output_path: Path | None
    method: str  # kms, xvfb, fallback
    error: str | None = None


class KmsGrabCapture:
    """KMS/DRM screenshot via ffmpeg -f kmsgrab.

    No sudo is ever used.  The caller must have read access to
    /dev/dri/cardX for this to work.
    """

    def capture(self, output_path: Path) -> CaptureResult:
        # 1. Detect whether the current user can access a DRI device
        dri_dev = Path("/dev/dri/card0")
        if dri_dev.exists():
            try:
                # Test read permission (no sudo)
                open(dri_dev).close()
            except PermissionError:
                return CaptureResult(False, None, "kms", "no /dev/dri permission")
        else:
            return CaptureResult(False, None, "kms", "no GPU")

        # 2. Use ffmpeg directly (no sudo)
        try:
            result = subprocess.run(
                [
                    "ffmpeg", "-y", "-f", "kmsgrab", "-i", "-",
                    "-frames:v", "1",
                    "-vf", "hwdownload,format=bgr0",
                    "-f", "image2", str(output_path),
                ],
                capture_output=True,
                text=True,
                timeout=30,
            )
            if result.returncode == 0 and output_path.exists():
                return CaptureResult(True, output_path, "kms", None)
            return CaptureResult(False, None, "kms", result.stderr)
        except FileNotFoundError:
            return CaptureResult(False, None, "kms", "ffmpeg not found")
        except subprocess.TimeoutExpired:
            return CaptureResult(False, None, "kms", "ffmpeg timed out")
        except Exception as e:
            return CaptureResult(False, None, "kms", str(e))


class XvfbCapture:
    """Xvfb virtual display screenshot via xvfb-run + scrot."""

    def capture(
        self,
        output_path: Path,
        width: int = 1920,
        height: int = 1080,
    ) -> CaptureResult:
        # Check xvfb-run availability
        if not shutil.which("xvfb-run"):
            return CaptureResult(False, None, "xvfb", "xvfb-run not found")

        # Use scrot as the screenshot tool inside Xvfb
        try:
            result = subprocess.run(
                [
                    "xvfb-run",
                    "--auto-servernum",
                    "--server-args",
                    f"-screen 0 {width}x{height}x24",
                    "--",
                    "scrot",
                    str(output_path),
                ],
                capture_output=True,
                text=True,
                timeout=60,
            )
            if result.returncode == 0 and output_path.exists():
                return CaptureResult(True, output_path, "xvfb", None)
            return CaptureResult(False, None, "xvfb", result.stderr)
        except FileNotFoundError:
            return CaptureResult(False, None, "xvfb", "xvfb-run not found")
        except subprocess.TimeoutExpired:
            return CaptureResult(False, None, "xvfb", "xvfb-run timed out")
        except Exception as e:
            return CaptureResult(False, None, "xvfb", str(e))


class FallbackCapture:
    """Fallback: write a text placeholder when no GUI capture is available."""

    def capture(
        self,
        output_path: Path,
        description: str = "",
    ) -> CaptureResult:
        output_path.write_text(
            f"[GUI capture not available]\n{description}\n"
        )
        return CaptureResult(True, output_path, "fallback", None)


class CaptureManager:
    """Auto-select the best capture method with fallback chain.

    Order: kms -> xvfb -> fallback.
    Use ``prefer`` to force a specific method or "auto" for the chain.
    """

    def __init__(self) -> None:
        self.kms = KmsGrabCapture()
        self.xvfb = XvfbCapture()
        self.fallback = FallbackCapture()

    def capture(
        self,
        output_path: Path,
        prefer: str = "auto",
    ) -> CaptureResult:
        """
        prefer: "kms" | "xvfb" | "fallback" | "auto"
        auto mode tries kms -> xvfb -> fallback in order.
        """
        if prefer in ("kms", "auto"):
            result = self.kms.capture(output_path)
            if result.success:
                return result
            if prefer == "kms":
                return result  # forced kms failed, return failure

        if prefer in ("xvfb", "auto"):
            result = self.xvfb.capture(output_path)
            if result.success:
                return result
            if prefer == "xvfb":
                return result  # forced xvfb failed, return failure

        # fallback
        return self.fallback.capture(output_path, "all capture methods failed")