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 BatchIpParseResult:
    """Parsed batch IP input after blank-skip and dedupe."""
    domains: list[str]
    input_lines: int = 0
    duplicates_removed: int = 0


def parse_batch_ip_lines(text: str, *, max_lines: int = 200) -> BatchIpParseResult:
    """Parse multi-line IP/domain input: one per line, ignore blanks, dedupe in order (case-insensitive)."""
    seen: set[str] = set()
    result: list[str] = []
    input_lines = 0
    for raw_line in text.splitlines():
        domain = raw_line.strip()
        if not domain:
            continue
        input_lines += 1
        key = domain.casefold()
        if key in seen:
            continue
        seen.add(key)
        result.append(domain)
        if len(result) > max_lines:
            raise ValueError(f"单次最多添加 {max_lines} 个地址")
    return BatchIpParseResult(
        domains=result,
        input_lines=input_lines,
        duplicates_removed=max(0, input_lines - len(result)),
    )