dashgradio_tunnel/app/services.py

from __future__ import annotations

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

import psutil

from app.policy import TunnelPolicy
from app.storage import Storage


@dataclass
class ServiceRecord:
    port: int
    protocol: str
    service: str
    pid: Optional[int]
    user: Optional[str]
    command: Optional[str]
    container: Optional[str]
    local_address: str
    allowed: bool = True
    policy_reason: Optional[str] = None


class PortScanner:
    def __init__(
        self,
        cache_ttl: float = 2.0,
        policy: TunnelPolicy | None = None,
        storage: Storage | None = None,
    ) -> None:
        self._cache_ttl = cache_ttl
        self._last_scan = 0.0
        self._cache: List[ServiceRecord] = []
        self._policy = policy
        self._storage = storage

    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
            allowed = True
            reason: Optional[str] = None
            if self._policy:
                allowed, reason = self._policy.is_port_allowed(port)

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

        self._persist_services(results)
        return results

    def _persist_services(self, records: List[ServiceRecord]) -> None:
        if not self._storage:
            return
        payload = [asdict(record) for record in records]
        try:
            self._storage.record_services(payload)
        except Exception:
            pass


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())