Bjorn/web_utils/package_utils.py

"""package_utils.py - Package installation, listing, and removal endpoints."""
from __future__ import annotations
import json
import logging
import os
import re
import subprocess
import time
from typing import Any, Dict

from logger import Logger

logger = Logger(name="package_utils.py", level=logging.DEBUG)

# Regex: alphanumeric, hyphens, underscores, dots, brackets (for extras like pkg[extra])
_VALID_PACKAGE_NAME = re.compile(r'^[a-zA-Z0-9_\-\.]+(\[[a-zA-Z0-9_\-\.,]+\])?$')


class PackageUtils:
    """Utilities for pip package management."""

    def __init__(self, shared_data):
        self.logger = logger
        self.shared_data = shared_data

    # =========================================================================
    # JSON ENDPOINTS
    # =========================================================================

    def list_packages_json(self, data: Dict) -> Dict:
        """Return all tracked packages."""
        try:
            packages = self.shared_data.db.list_packages()
            return {"status": "success", "data": packages}
        except Exception as e:
            self.logger.error(f"list_packages error: {e}")
            return {"status": "error", "message": str(e)}

    def uninstall_package(self, data: Dict) -> Dict:
        """Uninstall a pip package and remove from DB."""
        try:
            name = data.get("name")
            if not name:
                return {"status": "error", "message": "name is required"}
            if not _VALID_PACKAGE_NAME.match(name):
                return {"status": "error", "message": "Invalid package name"}

            result = subprocess.run(
                ["pip", "uninstall", "-y", name],
                capture_output=True, text=True, timeout=120,
            )
            if result.returncode != 0:
                return {"status": "error", "message": result.stderr.strip() or "Uninstall failed"}

            self.shared_data.db.remove_package(name)
            return {"status": "success", "message": f"Package '{name}' uninstalled"}
        except subprocess.TimeoutExpired:
            return {"status": "error", "message": "Uninstall timed out"}
        except Exception as e:
            self.logger.error(f"uninstall_package error: {e}")
            return {"status": "error", "message": str(e)}

    # =========================================================================
    # SSE ENDPOINT
    # =========================================================================

    def install_package(self, handler):
        """Stream pip install output as SSE events (GET endpoint)."""
        from urllib.parse import parse_qs, urlparse
        query = parse_qs(urlparse(handler.path).query)
        name = query.get("name", [""])[0].strip()

        # Validate
        if not name:
            handler.send_response(400)
            handler.send_header("Content-Type", "application/json")
            handler.end_headers()
            handler.wfile.write(json.dumps({"status": "error", "message": "name is required"}).encode("utf-8"))
            return
        if not _VALID_PACKAGE_NAME.match(name):
            handler.send_response(400)
            handler.send_header("Content-Type", "application/json")
            handler.end_headers()
            handler.wfile.write(json.dumps({"status": "error", "message": "Invalid package name"}).encode("utf-8"))
            return

        max_lifetime = 300  # 5 minutes maximum
        start_time = time.time()
        process = None
        try:
            handler.send_response(200)
            handler.send_header("Content-Type", "text/event-stream")
            handler.send_header("Cache-Control", "no-cache")
            handler.send_header("Connection", "keep-alive")
            handler.send_header("Access-Control-Allow-Origin", "*")
            handler.end_headers()

            process = subprocess.Popen(
                ["pip", "install", "--break-system-packages", name],
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                text=True,
                bufsize=1,
            )

            for line in process.stdout:
                if time.time() - start_time > max_lifetime:
                    self.logger.warning("install_package SSE stream reached max lifetime")
                    break

                payload = json.dumps({"line": line.rstrip(), "done": False})
                try:
                    handler.wfile.write(f"data: {payload}\n\n".encode("utf-8"))
                    handler.wfile.flush()
                except (ConnectionResetError, ConnectionAbortedError, BrokenPipeError, OSError):
                    self.logger.info("Client disconnected during package install")
                    break

            process.wait(timeout=30)
            success = process.returncode == 0

            # Get version on success
            version = ""
            if success:
                try:
                    show = subprocess.run(
                        ["pip", "show", name],
                        capture_output=True, text=True, timeout=15,
                    )
                    for show_line in show.stdout.splitlines():
                        if show_line.startswith("Version:"):
                            version = show_line.split(":", 1)[1].strip()
                            break
                except Exception:
                    pass

                # Record in DB
                self.shared_data.db.add_package(name, version)

            payload = json.dumps({"line": "", "done": True, "success": success, "version": version})
            try:
                handler.wfile.write(f"data: {payload}\n\n".encode("utf-8"))
                handler.wfile.flush()
            except (ConnectionResetError, ConnectionAbortedError, BrokenPipeError, OSError):
                pass

        except (ConnectionResetError, ConnectionAbortedError, BrokenPipeError):
            self.logger.info("Client disconnected from package install SSE stream")
        except Exception as e:
            self.logger.error(f"install_package SSE error: {e}")
            try:
                payload = json.dumps({"line": f"Error: {e}", "done": True, "success": False, "version": ""})
                handler.wfile.write(f"data: {payload}\n\n".encode("utf-8"))
                handler.wfile.flush()
            except Exception:
                pass
        finally:
            if process:
                try:
                    if process.stdout and not process.stdout.closed:
                        process.stdout.close()
                    if process.poll() is None:
                        process.terminate()
                        try:
                            process.wait(timeout=5)
                        except subprocess.TimeoutExpired:
                            process.kill()
                            process.wait()
                except Exception:
                    pass
            self.logger.info("Package install SSE stream closed")