Nexus/server/application/services/credential_poller.py

"""Poll password/SSH-key presets against a host until one SSH login succeeds."""

from __future__ import annotations

import logging
from dataclasses import dataclass, field

from sqlalchemy.ext.asyncio import AsyncSession

from server.infrastructure.database.crypto import decrypt_value
from server.infrastructure.database.password_preset_repo import PasswordPresetRepositoryImpl
from server.infrastructure.database.ssh_key_preset_repo import SshKeyPresetRepositoryImpl
from server.infrastructure.ssh.asyncssh_pool import try_ssh_login

logger = logging.getLogger("nexus.credential_poller")


@dataclass
class CredentialAttemptError:
    preset_type: str
    preset_name: str
    username: str
    error: str


@dataclass
class PollMatch:
    auth_method: str
    username: str
    preset_id: int | None = None
    ssh_key_preset_id: int | None = None
    password: str | None = None
    ssh_key_private: str | None = None
    ssh_key_public: str | None = None
    preset_name: str = ""


@dataclass
class PollResult:
    success: bool
    match: PollMatch | None = None
    errors: list[CredentialAttemptError] = field(default_factory=list)


_PRESET_TYPE_LABEL = {
    "password": "密码",
    "ssh_key": "密钥",
    "system": "系统",
}


def _preset_type_label(preset_type: str) -> str:
    return _PRESET_TYPE_LABEL.get(preset_type, preset_type)


def format_poll_errors(errors: list[CredentialAttemptError]) -> str:
    if not errors:
        return "无可用凭据预设或全部登录失败"
    lines = [
        f"[{_preset_type_label(e.preset_type)}] {e.preset_name} ({e.username}): {e.error}"
        for e in errors
    ]
    return "\n".join(lines)[:4000]


async def poll_credentials(
    db: AsyncSession,
    host: str,
    port: int = 22,
    *,
    timeout: float = 8.0,
) -> PollResult:
    """Try each password preset then each SSH key preset; first success wins."""
    pw_repo = PasswordPresetRepositoryImpl(db)
    key_repo = SshKeyPresetRepositoryImpl(db)
    password_presets = await pw_repo.get_all()
    ssh_key_presets = await key_repo.get_all()

    errors: list[CredentialAttemptError] = []

    if not password_presets and not ssh_key_presets:
        return PollResult(success=False, errors=[
            CredentialAttemptError("system", "—", "—", "请先在凭据管理中添加密码或 SSH 密钥预设"),
        ])

    for preset in password_presets:
        username = (preset.username or "root").strip() or "root"
        try:
            password = decrypt_value(preset.encrypted_pw)
        except ValueError as e:
            err = CredentialAttemptError("password", preset.name, username, f"解密失败: {e}")
            errors.append(err)
            logger.warning("Password preset %s decrypt failed: %s", preset.id, e)
            continue

        ok, msg = await try_ssh_login(
            host, port, username, password=password, timeout=timeout,
        )
        if ok:
            return PollResult(
                success=True,
                match=PollMatch(
                    auth_method="password",
                    username=username,
                    preset_id=preset.id,
                    password=password,
                    preset_name=preset.name,
                ),
            )
        errors.append(CredentialAttemptError("password", preset.name, username, msg or "认证失败"))

    for preset in ssh_key_presets:
        username = (preset.username or "root").strip() or "root"
        try:
            private_key = decrypt_value(preset.encrypted_private_key)
        except ValueError as e:
            err = CredentialAttemptError("ssh_key", preset.name, username, f"解密失败: {e}")
            errors.append(err)
            logger.warning("SSH key preset %s decrypt failed: %s", preset.id, e)
            continue

        ok, msg = await try_ssh_login(
            host, port, username, private_key=private_key, timeout=timeout,
        )
        if ok:
            return PollResult(
                success=True,
                match=PollMatch(
                    auth_method="key",
                    username=username,
                    ssh_key_preset_id=preset.id,
                    ssh_key_private=private_key,
                    ssh_key_public=preset.public_key or "",
                    preset_name=preset.name,
                ),
            )
        errors.append(CredentialAttemptError("ssh_key", preset.name, username, msg or "认证失败"))

    return PollResult(success=False, errors=errors)


@dataclass(frozen=True)
class BatchServerEntry:
    """One host from batch add input."""
    domain: str
    name: str | None = None


@dataclass(frozen=True)
class BatchIpParseResult:
    """Parsed batch IP input after blank-skip and dedupe."""
    entries: list[BatchServerEntry]
    input_lines: int = 0
    duplicates_removed: int = 0

    @property
    def domains(self) -> list[str]:
        return [e.domain for e in self.entries]


def _looks_like_host(value: str) -> bool:
    value = value.strip()
    if not value:
        return False
    if value.replace(".", "").isdigit():
        parts = value.split(".")
        return len(parts) == 4 and all(p.isdigit() and 0 <= int(p) <= 255 for p in parts)
    return "." in value or ":" in value


def _parse_batch_line(raw_line: str) -> BatchServerEntry | None:
    """Parse one line: IP/domain only, or ``名称<TAB>IP`` / ``名称,IP``."""
    line = raw_line.strip()
    if not line:
        return None

    if "\t" in line:
        name_part, domain = line.split("\t", 1)
        domain = domain.strip()
        name = name_part.strip()
        if domain and name:
            return BatchServerEntry(domain=domain, name=name)

    if "," in line:
        name_part, _, domain_part = line.rpartition(",")
        domain = domain_part.strip()
        name = name_part.strip()
        if domain and name and _looks_like_host(domain):
            return BatchServerEntry(domain=domain, name=name)

    return BatchServerEntry(domain=line, name=None)


def parse_batch_ip_lines(text: str, *, max_lines: int = 200) -> BatchIpParseResult:
    """Parse multi-line batch input: one host per line, optional name prefix, dedupe by domain."""
    seen: set[str] = set()
    result: list[BatchServerEntry] = []
    input_lines = 0
    for raw_line in text.splitlines():
        entry = _parse_batch_line(raw_line)
        if entry is None:
            continue
        input_lines += 1
        key = entry.domain.casefold()
        if key in seen:
            continue
        seen.add(key)
        result.append(entry)
        if len(result) > max_lines:
            raise ValueError(f"单次最多添加 {max_lines} 个地址")
    return BatchIpParseResult(
        entries=result,
        input_lines=input_lines,
        duplicates_removed=max(0, input_lines - len(result)),
    )