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,
    parse_batch_ip_lines,
    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_parse_batch_ip_lines_ignores_blank_and_dedupes():
    text = " 192.168.1.1 \n\n10.0.0.2\n192.168.1.1\n  \n10.0.0.3 "
    parsed = parse_batch_ip_lines(text)
    assert parsed.domains == ["192.168.1.1", "10.0.0.2", "10.0.0.3"]
    assert parsed.input_lines == 4
    assert parsed.duplicates_removed == 1


def test_parse_batch_ip_lines_case_insensitive_dedupe():
    parsed = parse_batch_ip_lines("Host.EXAMPLE.com\nhost.example.com\n10.0.0.1")
    assert parsed.domains == ["Host.EXAMPLE.com", "10.0.0.1"]
    assert parsed.duplicates_removed == 1


def test_parse_batch_ip_lines_max_limit():
    with pytest.raises(ValueError, match="最多"):
        parse_batch_ip_lines("1\n2\n3", max_lines=2)


def test_parse_batch_ip_lines_name_tab_ip():
    parsed = parse_batch_ip_lines("武汉测试公司\t192.168.1.1\n10.0.0.2")
    assert len(parsed.entries) == 2
    assert parsed.entries[0].domain == "192.168.1.1"
    assert parsed.entries[0].name == "武汉测试公司"
    assert parsed.entries[1].domain == "10.0.0.2"
    assert parsed.entries[1].name is None


def test_parse_batch_ip_lines_name_comma_ip():
    parsed = parse_batch_ip_lines("广州示例,47.107.1.2")
    assert parsed.entries[0].domain == "47.107.1.2"
    assert parsed.entries[0].name == "广州示例"


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 _try_add_by_ip → server created response shape."""
    from server.api import servers as servers_api
    from server.api.schemas import AddByIpRequest, AddByIpsBatchItemResult

    async def fake_try_add_by_ip(request, **kwargs):
        return AddByIpsBatchItemResult(
            domain=kwargs["domain"],
            status="created",
            server_id=1,
            matched_preset="ok-pw",
            matched_username="root",
        )

    monkeypatch.setattr(servers_api, "_try_add_by_ip", fake_try_add_by_ip)

    db = MagicMock()
    admin = MagicMock(username="admin")
    request = MagicMock()
    service = MagicMock()

    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"
    assert out["matched_preset"] == "ok-pw"


@pytest.mark.asyncio
async def test_add_by_ips_batch_endpoint(monkeypatch):
    from server.api import servers as servers_api
    from server.api.schemas import AddByIpsBatchRequest, AddByIpsBatchItemResult

    calls: list[str] = []

    async def fake_try_add_by_ip(request, **kwargs):
        domain = kwargs["domain"]
        calls.append(domain)
        if domain == "10.2.2.2":
            return AddByIpsBatchItemResult(domain=domain, status="created", server_id=2)
        if domain == "10.2.2.3":
            return AddByIpsBatchItemResult(domain=domain, status="pending", pending_id=9, message="fail")
        return AddByIpsBatchItemResult(domain=domain, status="skipped", message="exists")

    monkeypatch.setattr(servers_api, "_try_add_by_ip", fake_try_add_by_ip)

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

    audit_repo = MagicMock()
    audit_repo.create = AsyncMock()
    monkeypatch.setattr(servers_api, "AuditLogRepositoryImpl", lambda _db: audit_repo)

    payload = AddByIpsBatchRequest(text="10.2.2.2\n\n10.2.2.3\n10.2.2.2\n10.2.2.4")
    out = await servers_api.add_servers_by_ips_batch(request, payload, admin, service, db)

    assert calls == ["10.2.2.2", "10.2.2.3", "10.2.2.4"]
    assert out.total == 3
    assert out.created == 1
    assert out.pending == 1
    assert out.skipped == 1
    assert out.input_lines == 4
    assert out.duplicates_removed == 1