"""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_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