Nexus/tests/test_credential_poller.py

"""Tests for SSH credential polling and add-by-ip pending queue."""

from unittest.mock import AsyncMock, MagicMock, patch

import pytest
import pytest_asyncio
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine

from server.application.services.credential_poller import (
    format_poll_errors,
    poll_credentials,
    CredentialAttemptError,
)
from server.domain.models import Base, PasswordPreset, PendingServer, SshKeyPreset
from server.infrastructure.database.crypto import encrypt_value
from server.infrastructure.database.password_preset_repo import PasswordPresetRepositoryImpl
from server.infrastructure.database.pending_server_repo import PendingServerRepositoryImpl
from server.infrastructure.database.ssh_key_preset_repo import SshKeyPresetRepositoryImpl


@pytest_asyncio.fixture
async def db_session():
    engine = create_async_engine("sqlite+aiosqlite://", echo=False)
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)
    session_factory = async_sessionmaker(engine, class_=AsyncSession)
    session = session_factory()
    try:
        yield session
    finally:
        await session.close()
        await engine.dispose()


@pytest.mark.asyncio
async def test_poll_credentials_first_password_success(db_session):
    pw_repo = PasswordPresetRepositoryImpl(db_session)
    await pw_repo.create(PasswordPreset(
        name="prod-root",
        username="root",
        encrypted_pw=encrypt_value("secret1"),
    ))
    await pw_repo.create(PasswordPreset(
        name="prod-admin",
        username="admin",
        encrypted_pw=encrypt_value("secret2"),
    ))

    call_count = 0

    async def fake_try(host, port, username, *, password=None, private_key=None, timeout=8.0):
        nonlocal call_count
        call_count += 1
        if username == "root" and password == "secret1":
            return True, ""
        return False, "auth failed"

    with patch("server.application.services.credential_poller.try_ssh_login", side_effect=fake_try):
        result = await poll_credentials(db_session, "10.0.0.1", 22)

    assert result.success is True
    assert result.match is not None
    assert result.match.auth_method == "password"
    assert result.match.username == "root"
    assert result.match.preset_name == "prod-root"
    assert call_count == 1


@pytest.mark.asyncio
async def test_poll_credentials_tries_ssh_key_after_passwords_fail(db_session):
    pw_repo = PasswordPresetRepositoryImpl(db_session)
    await pw_repo.create(PasswordPreset(
        name="bad",
        username="root",
        encrypted_pw=encrypt_value("wrong"),
    ))
    key_repo = SshKeyPresetRepositoryImpl(db_session)
    await key_repo.create(SshKeyPreset(
        name="key1",
        username="deploy",
        encrypted_private_key=encrypt_value("-----BEGIN OPENSSH PRIVATE KEY-----\nabc\n-----END OPENSSH PRIVATE KEY-----"),
        public_key="ssh-rsa AAAA",
    ))

    async def fake_try(host, port, username, *, password=None, private_key=None, timeout=8.0):
        if private_key and username == "deploy":
            return True, ""
        return False, "auth failed"

    with patch("server.application.services.credential_poller.try_ssh_login", side_effect=fake_try):
        result = await poll_credentials(db_session, "10.0.0.2", 22)

    assert result.success is True
    assert result.match is not None
    assert result.match.auth_method == "key"
    assert result.match.ssh_key_preset_id is not None
    assert result.match.username == "deploy"


@pytest.mark.asyncio
async def test_poll_credentials_all_fail(db_session):
    pw_repo = PasswordPresetRepositoryImpl(db_session)
    await pw_repo.create(PasswordPreset(
        name="p1",
        username="root",
        encrypted_pw=encrypt_value("x"),
    ))

    with patch(
        "server.application.services.credential_poller.try_ssh_login",
        new=AsyncMock(return_value=(False, "Permission denied")),
    ):
        result = await poll_credentials(db_session, "10.0.0.3", 22)

    assert result.success is False
    assert result.match is None
    assert len(result.errors) == 1
    assert "p1" in format_poll_errors(result.errors)


@pytest.mark.asyncio
async def test_poll_credentials_no_presets(db_session):
    result = await poll_credentials(db_session, "10.0.0.4", 22)
    assert result.success is False
    assert any("凭据" in e.error for e in result.errors)


@pytest.mark.asyncio
async def test_pending_server_repo_record_failure(db_session):
    repo = PendingServerRepositoryImpl(db_session)
    row = await repo.record_failure(
        name="srv",
        domain="192.168.1.5",
        port=22,
        agent_port=8601,
        target_path="",
        category=None,
        platform_id=None,
        node_id=None,
        last_error="all failed",
    )
    assert row.id
    assert row.attempts == 1

    updated = await repo.record_failure(
        name="srv",
        domain="192.168.1.5",
        port=22,
        agent_port=8601,
        target_path="",
        category=None,
        platform_id=None,
        node_id=None,
        last_error="still failed",
        existing_id=row.id,
    )
    assert updated.attempts == 2
    assert "still failed" in (updated.last_error or "")


def test_format_poll_errors_truncates():
    errors = [
        CredentialAttemptError("password", "a", "root", "e1"),
        CredentialAttemptError("ssh_key", "b", "admin", "e2"),
    ]
    text = format_poll_errors(errors)
    assert "a" in text and "b" in text


@pytest.mark.asyncio
async def test_add_by_ip_endpoint_success(monkeypatch):
    """API layer: mocked poll → server created response shape."""
    from server.api import servers as servers_api
    from server.domain.models import Server

    mock_server = Server(
        id=1,
        name="10.1.1.1",
        domain="10.1.1.1",
        port=22,
        username="root",
        auth_method="password",
        connectivity="ok",
    )

    async def fake_poll(*_a, **_k):
        return servers_api.__dict__.get("_PollResult")  # won't work

    match = type("M", (), {
        "auth_method": "password",
        "username": "root",
        "preset_id": 1,
        "ssh_key_preset_id": None,
        "password": "pass",
        "ssh_key_private": None,
        "ssh_key_public": "",
        "preset_name": "ok-pw",
    })()
    poll_result = type("R", (), {"success": True, "match": match, "errors": []})()

    monkeypatch.setattr(
        "server.application.services.credential_poller.poll_credentials",
        AsyncMock(return_value=poll_result),
    )
    monkeypatch.setattr(
        servers_api,
        "_server_exists_for_domain",
        AsyncMock(return_value=False),
    )
    monkeypatch.setattr(
        servers_api,
        "_create_server_from_poll",
        AsyncMock(return_value=mock_server),
    )
    monkeypatch.setattr(servers_api, "_server_to_dict", lambda s: {"id": s.id, "domain": s.domain})

    db = MagicMock()
    admin = MagicMock(username="admin")
    request = MagicMock()
    request.client.host = "127.0.0.1"
    service = MagicMock()

    from server.api.schemas import AddByIpRequest

    out = await servers_api.add_server_by_ip(
        request,
        AddByIpRequest(domain="10.1.1.1", port=22),
        admin,
        service,
        db,
    )
    assert out["success"] is True
    assert out["server"]["domain"] == "10.1.1.1"