dashgradio_tunnel/app/services.py

from __future__ import annotations

import shlex
import socket
import subprocess
import time
from dataclasses import dataclass
from typing import Iterable, List, Optional

import psutil


@dataclass
class ServiceRecord:
    port: int
    protocol: str
    service: str
    pid: Optional[int]
    user: Optional[str]
    command: Optional[str]
    container: Optional[str]
    local_address: str


class PortScanner:
    def __init__(self, cache_ttl: float = 2.0) -> None:
        self._cache_ttl = cache_ttl
        self._last_scan = 0.0
        self._cache: List[ServiceRecord] = []

    def scan(self) -> List[ServiceRecord]:
        now = time.monotonic()
        if self._cache and now - self._last_scan < self._cache_ttl:
            return self._cache
        self._cache = self._perform_scan()
        self._last_scan = now
        return self._cache

    def _perform_scan(self) -> List[ServiceRecord]:
        connections = psutil.net_connections(kind="inet")
        container_names = _load_container_names()
        process_cache: dict[int, tuple[str, Optional[str], list[str]]] = {}
        seen_ports: set[int] = set()
        results: List[ServiceRecord] = []

        for conn in connections:
            if conn.status != psutil.CONN_LISTEN:
                continue
            if not conn.laddr:
                continue
            port = conn.laddr.port
            if port in seen_ports:
                continue
            seen_ports.add(port)
            protocol = _protocol_for_conn(conn.type)
            process_name = "unknown"
            process_user: Optional[str] = None
            process_command: Optional[str] = None
            pid = conn.pid

            if pid and pid > 0:
                cache_entry = process_cache.get(pid)
                if cache_entry is None:
                    cache_entry = _describe_process(pid)
                    process_cache[pid] = cache_entry
                process_name, process_user, raw_command = cache_entry
                if raw_command:
                    process_command = _format_commandline(raw_command)
            local_address = f"{conn.laddr.ip}:{conn.laddr.port}"
            container = _container_name_for_pid(pid, container_names) if pid else None

            results.append(
                ServiceRecord(
                    port=port,
                    protocol=protocol,
                    service=process_name,
                    pid=pid,
                    user=process_user,
                    command=process_command,
                    container=container,
                    local_address=local_address,
                )
            )

        return results


def _protocol_for_conn(sock_type: int) -> str:
    if sock_type == socket.SOCK_STREAM:
        return "tcp"
    if sock_type == socket.SOCK_DGRAM:
        return "udp"
    return "inet"


def _describe_process(pid: int) -> tuple[str, Optional[str], list[str]]:
    try:
        process = psutil.Process(pid)
        name = process.name() or "unknown"
        user = None
        try:
            user = process.username()
        except psutil.Error:
            user = None
        command = process.cmdline()
    except psutil.Error:
        return "unknown", None, []
    return name, user, command


def _format_commandline(cmdline: Iterable[str]) -> str:
    command = " ".join(shlex.quote(part) for part in cmdline if part)
    if len(command) > 150:
        return command[:150].rstrip() + "…"
    return command


def _load_container_names() -> dict[str, str]:
    try:
        result = subprocess.run(
            [
                "docker",
                "ps",
                "--no-trunc",
                "--format",
                "{{.ID}} {{.Names}}",
            ],
            capture_output=True,
            text=True,
            check=True,
        )
    except (subprocess.SubprocessError, FileNotFoundError):
        return {}
    containers: dict[str, str] = {}
    for line in result.stdout.splitlines():
        parts = line.strip().split(maxsplit=1)
        if len(parts) != 2:
            continue
        container_id, name = parts
        containers[container_id.lower()] = name
    return containers


def _container_name_for_pid(pid: Optional[int], container_names: dict[str, str]) -> Optional[str]:
    if not pid:
        return None
    try:
        with open(f"/proc/{pid}/cgroup", "r", encoding="utf-8") as handle:
            content = handle.read()
    except OSError:
        return None
    for line in content.splitlines():
        lower = line.lower()
        if not any(keyword in lower for keyword in ("docker", "containerd", "kubepods", "podman")):
            continue
        for part in line.split("/"):
            part = part.strip()
            if not part:
                continue
            if _looks_like_container_id(part):
                for container_id, name in container_names.items():
                    if container_id.startswith(part.lower()):
                        return name
    return None


def _looks_like_container_id(value: str) -> bool:
    value = value.strip()
    if len(value) < 12:
        return False
    return all(ch in "0123456789abcdef" for ch in value[:12].lower())