utp_service/backend/app/services/device_discovery.py

"""
SSH-based device discovery: connects to each rack PC, reads its caps config,
extracts plugin URIs and classifies discovered network devices.

Discovery chain:
  Object DB → racks table (host IPs) → SSH per rack → parse config → upsert devices

Limitations (devices not discoverable automatically):
  - Serial/USB devices (no network URI)
  - RODOS-8 relays, Moxa switches (not represented as plugins)
  - Physical MikroTik router (not in rack configs)
"""

import re
from configparser import RawConfigParser
from dataclasses import dataclass, field
from io import StringIO
from urllib.parse import urlparse

import asyncssh
import psycopg
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.models.device import Device
from app.models.object import Object
from app.services.crypto import decrypt

# ── URI extraction ────────────────────────────────────────────────────────────

_SERIAL_SCHEMES = {"serial", "com"}
_SERVICE_HOSTS  = {"localhost", "127.0.0.1"}

def _extract_ip(uri: str) -> str | None:
    """Return IP from a URI string, or None if not a network device."""
    if not uri:
        return None
    try:
        parsed = urlparse(uri)
        scheme = (parsed.scheme or "").lower()
        if scheme in _SERIAL_SCHEMES:
            return None
        hostname = parsed.hostname
        if not hostname or hostname in _SERVICE_HOSTS:
            return None
        # Validate it looks like an IP (simple check)
        if re.match(r"^\d{1,3}(\.\d{1,3}){3}$", hostname):
            return hostname
    except Exception:
        pass
    return None


# ── Plugin class → device category ───────────────────────────────────────────

def _cls_to_category(cls: str) -> str:
    cls_lower = cls.lower()
    if "camera" in cls_lower or "facedetect" in cls_lower:
        return "camera"
    if "controller" in cls_lower or "ovenmk" in cls_lower or "rodos" in cls_lower:
        return "io_board"
    if "payment" in cls_lower or "sberpilot" in cls_lower or "payx" in cls_lower:
        return "bank_terminal"
    if "payonline" in cls_lower or "kkt" in cls_lower or "fiscal" in cls_lower:
        return "cash_register"
    return "other"


def _cls_to_role(cls: str, plugin_name: str) -> str:
    cls_lower = cls.lower()
    name_lower = plugin_name.lower()
    if "camera" in cls_lower or "camera" in name_lower:
        if "grz" in name_lower or "plate" in name_lower or "lpr" in name_lower or "frame" in name_lower:
            return "plate"
        if "face" in name_lower or "lico" in name_lower:
            return "face"
        return "overview"
    return "other"


# Plugins whose URIs point at the CAPS server itself, not at separate devices
_SKIP_PLUGIN_PREFIXES = ("mnemo", "processor", "recognition", "photo", "voice")


# ── Config parsing ────────────────────────────────────────────────────────────

@dataclass
class DiscoveredDevice:
    ip: str
    category: str
    role: str
    plugin_name: str
    cls: str


def _parse_config(config_text: str, rack_host: str) -> list[DiscoveredDevice]:
    """Extract discovered network devices from a caps config file."""
    parser = RawConfigParser()
    parser.read_string(config_text)

    discovered: list[DiscoveredDevice] = []
    seen_ips: set[str] = set()

    for section in parser.sections():
        if not section.startswith("plugin:"):
            continue

        plugin_name = section[len("plugin:"):]
        if any(plugin_name.startswith(prefix) for prefix in _SKIP_PLUGIN_PREFIXES):
            continue

        cls = parser.get(section, "cls", fallback="")
        if not cls:
            continue

        # Collect all URI-like values in this section
        uri_fields = ("uri", "frame_uri", "host")
        for field_name in uri_fields:
            raw = parser.get(section, field_name, fallback=None)
            if not raw:
                continue
            ip = _extract_ip(raw)
            if not ip or ip == rack_host or ip in seen_ips:
                continue
            seen_ips.add(ip)
            discovered.append(DiscoveredDevice(
                ip=ip,
                category=_cls_to_category(cls),
                role=_cls_to_role(cls, plugin_name),
                plugin_name=plugin_name,
                cls=cls,
            ))

    return discovered


# ── SSH operations ────────────────────────────────────────────────────────────

async def _read_remote_config(
    host: str,
    port: int,
    username: str,
    password: str,
    config_path: str,
) -> str:
    async with asyncssh.connect(
        host,
        port=port,
        username=username,
        password=password,
        known_hosts=None,
        connect_timeout=10,
    ) as conn:
        result = await conn.run(f"cat {config_path}", check=True)
        return result.stdout


async def _get_rack_hosts(obj: Object, db_password: str) -> list[tuple[str, str]]:
    """Return list of (external_id, host_ip) from the object's racks table."""
    if obj.db_via_tunnel and obj.server_ip and obj.ssh_user and obj.ssh_pass_enc:
        ssh_password = decrypt(obj.ssh_pass_enc)
        async with asyncssh.connect(
            obj.server_ip,
            port=obj.ssh_port or 22,
            username=obj.ssh_user,
            password=ssh_password,
            known_hosts=None,
        ) as tunnel:
            listener = await tunnel.forward_local_port(
                "127.0.0.1", 0, obj.db_host, obj.db_port or 5432
            )
            local_port = listener.get_port()
            dsn = (
                f"host=127.0.0.1 port={local_port} dbname={obj.db_name} "
                f"user={obj.db_user} password={db_password} connect_timeout=10"
            )
            return await _fetch_rack_hosts(dsn, obj.db_table)
    else:
        dsn = (
            f"host={obj.db_host} port={obj.db_port or 5432} dbname={obj.db_name} "
            f"user={obj.db_user} password={db_password} connect_timeout=10"
        )
        return await _fetch_rack_hosts(dsn, obj.db_table)


async def _fetch_rack_hosts(dsn: str, table: str) -> list[tuple[str, str]]:
    async with await psycopg.AsyncConnection.connect(dsn) as conn:
        cursor = await conn.execute(
            f"SELECT id, data->>'host' AS host FROM \"{table}\" WHERE data->>'host' IS NOT NULL"
        )
        rows = await cursor.fetchall()
    return [(str(r[0]), str(r[1])) for r in rows if r[1]]


# ── Main entry point ──────────────────────────────────────────────────────────

@dataclass
class DiscoveryResult:
    created: int = 0
    updated: int = 0
    skipped: int = 0
    errors: list[str] = field(default_factory=list)


async def discover_via_ssh(
    db: AsyncSession,
    obj: Object,
    config_path: str = "/etc/caps/caps.conf",
) -> DiscoveryResult:
    result = DiscoveryResult()

    if not all([obj.db_host, obj.db_name, obj.db_user, obj.db_pass_enc, obj.db_table]):
        result.errors.append("Object DB connection is not fully configured")
        return result
    if not obj.ssh_user or not obj.ssh_pass_enc:
        result.errors.append("SSH credentials not configured on this object")
        return result

    try:
        db_password = decrypt(obj.db_pass_enc)
        ssh_password = decrypt(obj.ssh_pass_enc)
    except Exception:
        result.errors.append("Failed to decrypt stored credentials")
        return result

    # Step 1: get rack hosts from object's DB
    try:
        rack_hosts = await _get_rack_hosts(obj, db_password)
    except Exception as exc:
        result.errors.append(f"Cannot fetch rack hosts from object DB: {exc}")
        return result

    if not rack_hosts:
        result.errors.append("No racks with host IPs found in object DB")
        return result

    # Step 2: SSH into each rack, parse config, upsert devices
    for external_id, rack_host in rack_hosts:
        # Resolve config_path override from the embedded device's connection_params
        existing_rack_pc = (await db.execute(
            select(Device).where(
                Device.object_id == obj.id,
                Device.ip == rack_host,
            )
        )).scalar_one_or_none()

        effective_config_path = config_path
        if existing_rack_pc and existing_rack_pc.connection_params:
            effective_config_path = existing_rack_pc.connection_params.get(
                "config_path", config_path
            )

        try:
            config_text = await _read_remote_config(
                host=rack_host,
                port=obj.ssh_port or 22,
                username=obj.ssh_user,
                password=ssh_password,
                config_path=effective_config_path,
            )
        except Exception as exc:
            result.errors.append(f"SSH to {rack_host} failed: {exc}")
            continue

        devices_found = _parse_config(config_text, rack_host)

        for found in devices_found:
            try:
                existing = (await db.execute(
                    select(Device).where(
                        Device.object_id == obj.id,
                        Device.ip == found.ip,
                    )
                )).scalar_one_or_none()

                if existing is not None:
                    changed = False
                    if existing.category != found.category and existing.source == "auto_ssh":
                        existing.category = found.category
                        changed = True
                    if existing.role == "other" and found.role != "other":
                        existing.role = found.role
                        changed = True
                    result.updated += 1 if changed else 0
                    result.skipped += 0 if changed else 1
                else:
                    db.add(Device(
                        object_id=obj.id,
                        ip=found.ip,
                        category=found.category,
                        role=found.role,
                        source="auto_ssh",
                        device_meta={"plugin": found.plugin_name, "cls": found.cls},
                    ))
                    result.created += 1
            except Exception as exc:
                result.errors.append(f"Error upserting {found.ip}: {exc}")

    return result