security-alert-center/backend/app/services/winrm_connect.py

"""WinRM connectivity test for Windows hosts using domain admin credentials."""

from __future__ import annotations

import socket
from dataclasses import dataclass

from app.models import Host


class WinAdminNotConfiguredError(Exception):
    pass


class HostNotWindowsError(Exception):
    pass


class HostTargetMissingError(Exception):
    pass


@dataclass(frozen=True)
class WinRmTestResult:
    ok: bool
    message: str
    target: str
    hostname: str | None = None


@dataclass(frozen=True)
class WinRmCmdResult:
    ok: bool
    message: str
    target: str
    stdout: str = ""
    stderr: str = ""
    exit_code: int | None = None


def _winrm_session(
    target: str,
    user: str,
    password: str,
    *,
    timeout_sec: int = 15,
):
    from app.services.win_admin_settings import normalize_win_admin_user

    user = normalize_win_admin_user(user.strip())
    target = target.strip()
    if not target or not user or not password:
        raise WinAdminNotConfiguredError("Windows admin credentials or target missing")

    try:
        import winrm
    except ImportError as exc:
        raise RuntimeError("pywinrm is not installed on SAC server") from exc

    operation_timeout_sec = max(5, timeout_sec)
    read_timeout_sec = operation_timeout_sec + 15
    endpoint = f"http://{target}:5985/wsman"
    session = winrm.Session(
        endpoint,
        auth=(user, password),
        transport="ntlm",
        read_timeout_sec=read_timeout_sec,
        operation_timeout_sec=operation_timeout_sec,
    )
    return session, winrm


def run_winrm_cmd(
    *,
    target: str,
    user: str,
    password: str,
    remote_cmd: str,
    timeout_sec: int = 30,
) -> WinRmCmdResult:
    target = target.strip()
    try:
        session, winrm = _winrm_session(target, user, password, timeout_sec=timeout_sec)
        result = session.run_cmd(remote_cmd)
    except winrm.exceptions.WinRMTransportError as exc:
        detail = str(exc)
        hint = ""
        if "credentials were rejected" in detail.lower():
            hint = " Проверьте логин (B26\\user), пароль и WinRM на ПК."
        return WinRmCmdResult(
            ok=False,
            message=f"WinRM transport error ({target}): {detail}{hint}",
            target=target,
        )
    except (socket.timeout, TimeoutError):
        return WinRmCmdResult(
            ok=False,
            message=f"WinRM timeout ({timeout_sec}s) to {target}:5985",
            target=target,
        )
    except WinAdminNotConfiguredError as exc:
        return WinRmCmdResult(ok=False, message=str(exc), target=target)
    except RuntimeError as exc:
        return WinRmCmdResult(ok=False, message=str(exc), target=target)
    except Exception as exc:
        return WinRmCmdResult(ok=False, message=f"WinRM error ({target}): {exc}", target=target)

    stdout = (result.std_out or b"").decode("utf-8", errors="replace")
    stderr = (result.std_err or b"").decode("utf-8", errors="replace")
    exit_code = int(result.status_code)
    ok = exit_code == 0
    if ok:
        message = f"WinRM OK ({target}), exit 0"
    else:
        detail = stderr.strip() or stdout.strip() or f"exit code {exit_code}"
        message = f"WinRM command failed ({target}): {detail[:500]}"
    return WinRmCmdResult(
        ok=ok,
        message=message,
        target=target,
        stdout=stdout,
        stderr=stderr,
        exit_code=exit_code,
    )


def run_winrm_qwinsta(*, target: str, user: str, password: str) -> WinRmCmdResult:
    return run_winrm_cmd(target=target, user=user, password=password, remote_cmd="qwinsta", timeout_sec=45)


def run_winrm_logoff(*, target: str, user: str, password: str, session_id: int) -> WinRmCmdResult:
    if session_id < 0:
        return WinRmCmdResult(ok=False, message="Invalid session_id", target=target)
    return run_winrm_cmd(
        target=target,
        user=user,
        password=password,
        remote_cmd=f"logoff {session_id} /v",
        timeout_sec=45,
    )


def run_winrm_on_host_targets(
    host: Host,
    *,
    user: str,
    password: str,
    action,
) -> tuple[WinRmCmdResult | None, list[str]]:
    """Try WinRM targets in hostname-first order; action(target) -> WinRmCmdResult."""
    attempts: list[str] = []
    last: WinRmCmdResult | None = None
    for target in iter_winrm_targets(host):
        result = action(target=target)
        last = result
        attempts.append(f"{target}={'OK' if result.ok else 'fail'}")
        if result.ok:
            return result, attempts
    return last, attempts


def resolve_windows_host_target(host: Host) -> str:
    targets = iter_winrm_targets(host)
    if not targets:
        raise HostTargetMissingError("Host has no hostname or IPv4 for WinRM test")
    return targets[0]


def _netbios_name_variants(name: str | None) -> list[str]:
    """Короткое имя ПК (NetBIOS), как Enter-PSSession -ComputerName Andrisonova-PC."""
    text = (name or "").strip()
    if not text:
        return []
    short = text.split(".")[0].split("\\")[0].strip()
    if not short:
        return []
    if short.casefold() == text.casefold():
        return [short]
    return [short, text]


def _looks_like_computer_name(value: str) -> bool:
    text = value.strip()
    if not text or " " in text:
        return False
    return len(text) <= 63


def iter_winrm_targets(host: Host) -> list[str]:
    """WinRM + NTLM: сначала имя хоста (как Enter-PSSession -ComputerName), IP — последним."""
    if not is_windows_host(host):
        raise HostNotWindowsError("Host is not Windows")

    seen: set[str] = set()
    ordered: list[str] = []

    def add(value: str | None) -> None:
        text = (value or "").strip()
        if not text or text.casefold() in seen:
            return
        seen.add(text.casefold())
        ordered.append(text)

    for variant in _netbios_name_variants(host.hostname):
        add(variant)

    inventory = host.inventory if isinstance(host.inventory, dict) else {}
    computer_name = inventory.get("computer_name")
    if isinstance(computer_name, str):
        for variant in _netbios_name_variants(computer_name):
            add(variant)

    if host.display_name and _looks_like_computer_name(host.display_name):
        for variant in _netbios_name_variants(host.display_name):
            add(variant)

    add(host.ipv4)
    return ordered


def is_windows_host(host: Host) -> bool:
    if (host.os_family or "").strip().lower() == "windows":
        return True
    return (host.product or "").strip() == "rdp-login-monitor"


def test_winrm_connection(
    *,
    target: str,
    user: str,
    password: str,
    timeout_sec: int = 15,
) -> WinRmTestResult:
    result = run_winrm_cmd(
        target=target,
        user=user,
        password=password,
        remote_cmd="hostname",
        timeout_sec=timeout_sec,
    )
    if result.ok and result.stdout.strip():
        host = result.stdout.strip().splitlines()[0]
        return WinRmTestResult(
            ok=True,
            message=f"WinRM OK, hostname={host}",
            target=result.target,
            hostname=host,
        )
    return WinRmTestResult(
        ok=result.ok,
        message=result.message,
        target=result.target,
        hostname=None,
    )