Nexus/tests/test_server_file_transfer.py

"""Unit tests for cross-server file transfer."""

from __future__ import annotations

import json
from unittest.mock import AsyncMock, MagicMock

import pytest
from pydantic import ValidationError

from server.api.remote_path_validation import RemotePathError, assert_clipboard_transfer_safe
from server.application.services.server_file_transfer_service import (
    MAX_RELAY_FILE_BYTES,
    ServerTransferError,
    _assert_distinct_servers,
    _download_url_to_remote_path,
    _extract_archive_on_server,
    _relay_file,
    cleanup_source_archive,
    deliver_transfer_package,
    pull_transfer_package,
    relay_files,
)
from server.domain.models import Server


pytestmark = pytest.mark.unit


# ── validation helpers ──


def test_assert_distinct_servers_rejects_same():
    with pytest.raises(ServerTransferError, match="不能相同"):
        _assert_distinct_servers(1, 1)


def test_assert_distinct_servers_ok():
    _assert_distinct_servers(1, 2)


def test_clipboard_transfer_safe_rejects_dest_inside_source():
    with pytest.raises(RemotePathError):
        assert_clipboard_transfer_safe(["/www/a"], "/www/a/sub")


def test_relay_file_size_limit_constant():
    assert MAX_RELAY_FILE_BYTES == 524_288_000


# ── schemas ──


def test_server_file_transfer_relay_schema_validation():
    from server.api.schemas import ServerFileTransferRelay

    payload = ServerFileTransferRelay(
        source_server_id=1,
        dest_server_id=2,
        sources=["/www/a.txt"],
        dest_dir="/www/backup",
    )
    assert payload.dest_dir == "/www/backup"
    assert payload.sources == ["/www/a.txt"]


def test_server_file_transfer_relay_schema_rejects_empty_sources():
    from server.api.schemas import ServerFileTransferRelay

    with pytest.raises(ValidationError):
        ServerFileTransferRelay(
            source_server_id=1,
            dest_server_id=2,
            sources=[],
            dest_dir="/www/backup",
        )


def test_server_file_transfer_deliver_schema_validation():
    from server.api.schemas import ServerFileTransferDeliver

    payload = ServerFileTransferDeliver(
        source_server_id=1,
        dest_server_id=3,
        paths=["/www/wwwroot"],
        dest_path="/tmp/site.tar.gz",
        extract=True,
    )
    assert payload.format == "tar.gz"
    assert payload.extract is True


# ── JWT public path ──


def test_transfer_download_is_public_get():
    from server.api.auth_jwt import _is_public_path

    assert _is_public_path("/api/sync/transfer-download/abc123", "GET") is True
    # Prefix 白名单对所有 method 生效；路由仅注册 GET
    assert _is_public_path("/api/sync/transfer-download/abc123", "POST") is True
    assert _is_public_path("/api/sync/server-transfer", "GET") is False


# ── Redis transfer package store ──


class _FakeRedis:
    def __init__(self) -> None:
        self.data: dict[str, str] = {}

    async def set(self, key: str, value: str, ex: int | None = None) -> None:
        self.data[key] = value

    async def get(self, key: str) -> str | None:
        return self.data.get(key)

    async def delete(self, key: str) -> None:
        self.data.pop(key, None)


@pytest.mark.asyncio
async def test_transfer_package_store_roundtrip(monkeypatch):
    from server.infrastructure.redis import transfer_package_store as store

    fake = _FakeRedis()
    monkeypatch.setattr(store, "get_redis", lambda: fake)

    meta = {"source_server_id": 1, "archive_path": "/tmp/x.tar.gz", "format": "tar.gz"}
    await store.save_package_meta("tok-a", meta, ttl=60)
    loaded = await store.load_package_meta("tok-a")
    assert loaded == meta

    await store.delete_package_meta("tok-a")
    assert await store.load_package_meta("tok-a") is None


@pytest.mark.asyncio
async def test_transfer_package_store_invalid_json(monkeypatch):
    from server.infrastructure.redis import transfer_package_store as store

    fake = _FakeRedis()
    fake.data["transfer:pkg:bad"] = "not-json"
    monkeypatch.setattr(store, "get_redis", lambda: fake)

    assert await store.load_package_meta("bad") is None


# ── SFTP relay file size guard ──


@pytest.mark.asyncio
async def test_relay_file_rejects_oversized():
    src_sftp = AsyncMock()
    dest = Server(id=2, name="b", domain="2.2.2.2")
    stat = MagicMock()
    stat.size = MAX_RELAY_FILE_BYTES + 1
    src_sftp.stat = AsyncMock(return_value=stat)

    with pytest.raises(ServerTransferError, match="超过直传上限"):
        await _relay_file(src_sftp, AsyncMock(), dest, "/www/big.bin", "/tmp/big.bin")


@pytest.mark.asyncio
async def test_relay_file_uses_remote_write_bytes(monkeypatch):
    src_sftp = AsyncMock()
    dest = Server(id=2, name="b", domain="2.2.2.2")
    conn_dst = AsyncMock()
    stat = MagicMock()
    stat.size = 3
    src_sftp.stat = AsyncMock(return_value=stat)

    read_chunks = [b"ab", b"c", b""]

    class _FakeReader:
        async def read(self, _n: int) -> bytes:
            return read_chunks.pop(0)

        async def __aenter__(self):
            return self

        async def __aexit__(self, *_a):
            return None

    src_sftp.open = MagicMock(return_value=_FakeReader())
    write_mock = AsyncMock()
    monkeypatch.setattr(
        "server.application.services.server_file_transfer_service.remote_write_bytes",
        write_mock,
    )

    nbytes = await _relay_file(src_sftp, conn_dst, dest, "/www/a.txt", "/tmp/a.txt")
    assert nbytes == 3
    write_mock.assert_awaited_once()
    assert write_mock.await_args.args[2] == b"abc"


# ── deliver orchestration (mocked) ──


@pytest.mark.asyncio
async def test_extract_archive_uses_validating_decompress(monkeypatch):
    from server.application.services import server_file_transfer_service as svc

    server = Server(name="t", domain="1.2.3.4")
    decompress = AsyncMock(return_value={"exit_code": 0})
    monkeypatch.setattr(svc, "run_remote_decompress", decompress)

    await svc._extract_archive_on_server(
        server, "/www/site/pkg.tar.gz", "/www/site", "tar.gz",
    )

    decompress.assert_awaited_once_with(
        server, "/www/site/pkg.tar.gz", "/www/site", "tar.gz", timeout=600,
    )


@pytest.mark.asyncio
async def test_deliver_uses_btpanel_when_source_configured(monkeypatch):
    from server.application.services import server_file_transfer_service as svc

    source = Server(
        id=1,
        name="a",
        domain="1.1.1.1",
        extra_attrs={
            "bt_panel": {
                "base_url": "https://1.1.1.1:8888",
                "api_key_enc": "enc",
            },
        },
    )
    bt_result = {
        "mode": "deliver",
        "via": "btpanel_download",
        "dest_path": "/www/site/pkg.tar.gz",
    }
    bt_mock = AsyncMock(return_value=bt_result)
    package_mock = AsyncMock()
    monkeypatch.setattr(svc, "_get_server", AsyncMock(return_value=source))
    monkeypatch.setattr(svc, "bt_panel_configured", lambda _s: True)
    monkeypatch.setattr(svc, "read_bt_panel_credentials", lambda _s: object())
    monkeypatch.setattr(svc, "_deliver_via_btpanel_download", bt_mock)
    monkeypatch.setattr(svc, "create_transfer_package", package_mock)

    session = MagicMock()
    result = await deliver_transfer_package(
        session,
        source_server_id=1,
        paths=["/www/wwwroot/index.html"],
        dest_server_id=2,
        dest_path="/www/site/pkg.tar.gz",
    )

    assert result["via"] == "btpanel_download"
    bt_mock.assert_awaited_once()
    package_mock.assert_not_called()


@pytest.mark.asyncio
async def test_pull_downloads_via_tmp_then_mv(monkeypatch):
    from server.application.services import server_file_transfer_service as svc

    dest = Server(id=2, name="b", domain="2.2.2.2")
    meta = {
        "source_server_id": 1,
        "archive_path": "/tmp/pkg.tar.gz",
        "filename": "pkg.tar.gz",
        "size": 201,
        "format": "tar.gz",
        "expires_at": "2026-06-21T12:00:00+00:00",
    }
    monkeypatch.setattr(svc, "load_package_meta", AsyncMock(return_value=meta))
    monkeypatch.setattr(svc, "_get_server", AsyncMock(return_value=dest))
    monkeypatch.setattr(
        svc,
        "_apply_dest_permissions",
        AsyncMock(return_value=[{"applied": True}]),
    )

    exec_calls: list[str] = []

    async def fake_exec(_server, cmd, timeout=120):
        exec_calls.append(cmd)
        if cmd.startswith("tar "):
            return {"exit_code": 0}
        return {"exit_code": 0}

    monkeypatch.setattr(svc, "exec_ssh_command_with_fallback", fake_exec)
    monkeypatch.setattr(svc, "_api_base_url", lambda: "https://api.example.com")

    session = MagicMock()
    result = await svc.pull_transfer_package(
        session,
        token="tok123",
        dest_server_id=2,
        dest_path="/www/site/pkg.tar.gz",
        extract=False,
    )

    assert result["dest_path"] == "/www/site/pkg.tar.gz"
    curl_cmds = [c for c in exec_calls if "curl " in c]
    mv_cmds = [c for c in exec_calls if c.startswith("mv -f")]
    assert len(curl_cmds) == 1
    assert "/tmp/nexus-xfer-" in curl_cmds[0]
    assert "/www/site/pkg.tar.gz" not in curl_cmds[0]
    assert len(mv_cmds) == 1
    assert "/www/site/pkg.tar.gz" in mv_cmds[0]


@pytest.mark.asyncio
async def test_download_url_uses_curl_insecure_for_btpanel(monkeypatch):
    from server.application.services import server_file_transfer_service as svc

    exec_calls: list[str] = []

    async def fake_exec(_server, cmd, timeout=120):
        exec_calls.append(cmd)
        return {"exit_code": 0}

    monkeypatch.setattr(svc, "exec_ssh_command_with_fallback", fake_exec)
    monkeypatch.setattr(svc, "_ensure_remote_dir", AsyncMock())
    server = Server(name="b", domain="2.2.3.4")
    await svc._download_url_to_remote_path(
        server,
        download_url="https://1.1.1.1:8888/download?filename=%2Ftmp%2Fa.tar.gz",
        dest_path="/www/site/a.tar.gz",
        curl_insecure=True,
    )
    curl_cmd = next(c for c in exec_calls if "curl " in c)
    assert " -k " in curl_cmd or curl_cmd.startswith("curl -fsSL -k")


def _sample_package(*, size: int = 128) -> dict:
    return {
        "mode": "package",
        "token": "pkg-token-xyz",
        "download_url": "https://api.example.com/api/sync/transfer-download/pkg-token-xyz",
        "archive_path": "/tmp/nexus-xfer-pkg.tar.gz",
        "size": size,
        "expires_at": "2026-06-21T12:00:00+00:00",
        "source_server_id": 1,
    }


@pytest.mark.asyncio
async def test_deliver_success_cleans_archive_and_deletes_token(monkeypatch):
    from server.application.services import server_file_transfer_service as svc

    package = _sample_package()
    pull_result = {"dest_path": "/www/site.tar.gz", "extracted_to": "/www"}

    monkeypatch.setattr(svc, "create_transfer_package", AsyncMock(return_value=package))
    monkeypatch.setattr(svc, "pull_transfer_package", AsyncMock(return_value=pull_result))
    monkeypatch.setattr(
        svc,
        "_get_server",
        AsyncMock(return_value=Server(name="src", domain="1.2.3.4")),
    )
    cleanup_mock = AsyncMock()
    delete_mock = AsyncMock()
    monkeypatch.setattr(svc, "cleanup_source_archive", cleanup_mock)
    monkeypatch.setattr(svc, "delete_package_meta", delete_mock)

    session = MagicMock()
    result = await deliver_transfer_package(
        session,
        source_server_id=1,
        paths=["/www/wwwroot"],
        dest_server_id=2,
        dest_path="/www/site.tar.gz",
    )

    assert result["mode"] == "deliver"
    assert result["dest_path"] == "/www/site.tar.gz"
    cleanup_mock.assert_awaited_once()
    delete_mock.assert_awaited_once_with("pkg-token-xyz")


@pytest.mark.asyncio
async def test_deliver_pull_failure_rolls_back_archive_and_token(monkeypatch):
    from server.application.services import server_file_transfer_service as svc

    package = _sample_package()
    source = Server(name="src", domain="1.2.3.4")

    monkeypatch.setattr(svc, "create_transfer_package", AsyncMock(return_value=package))
    monkeypatch.setattr(
        svc,
        "pull_transfer_package",
        AsyncMock(side_effect=ServerTransferError("下载失败")),
    )
    monkeypatch.setattr(svc, "_get_server", AsyncMock(return_value=source))
    cleanup_mock = AsyncMock()
    delete_mock = AsyncMock()
    monkeypatch.setattr(svc, "cleanup_source_archive", cleanup_mock)
    monkeypatch.setattr(svc, "delete_package_meta", delete_mock)

    session = MagicMock()
    with pytest.raises(ServerTransferError, match="下载失败"):
        await deliver_transfer_package(
            session,
            source_server_id=1,
            paths=["/www/wwwroot"],
            dest_server_id=2,
            dest_path="/www/site.tar.gz",
        )

    cleanup_mock.assert_awaited_once_with(source, package["archive_path"])
    delete_mock.assert_awaited_once_with("pkg-token-xyz")


@pytest.mark.asyncio
async def test_deliver_rejects_oversized_package(monkeypatch):
    from server.application.services import server_file_transfer_service as svc

    package = _sample_package(size=MAX_RELAY_FILE_BYTES + 1)
    source = Server(name="src", domain="1.2.3.4")

    monkeypatch.setattr(svc, "create_transfer_package", AsyncMock(return_value=package))
    monkeypatch.setattr(svc, "_get_server", AsyncMock(return_value=source))
    cleanup_mock = AsyncMock()
    delete_mock = AsyncMock()
    monkeypatch.setattr(svc, "cleanup_source_archive", cleanup_mock)
    monkeypatch.setattr(svc, "delete_package_meta", delete_mock)
    pull_mock = AsyncMock()
    monkeypatch.setattr(svc, "pull_transfer_package", pull_mock)

    session = MagicMock()
    with pytest.raises(ServerTransferError, match="超过下载上限"):
        await deliver_transfer_package(
            session,
            source_server_id=1,
            paths=["/www/huge"],
            dest_server_id=2,
            dest_path="/www/huge.tar.gz",
        )

    pull_mock.assert_not_called()
    cleanup_mock.assert_awaited_once()
    delete_mock.assert_awaited_once_with("pkg-token-xyz")


# ── relay_files (mocked SSH) ──


@pytest.mark.asyncio
async def test_relay_files_rejects_same_server(monkeypatch):
    session = MagicMock()
    with pytest.raises(ServerTransferError, match="不能相同"):
        await relay_files(
            session,
            source_server_id=5,
            dest_server_id=5,
            sources=["/www/a.txt"],
            dest_dir="/www/backup",
        )


@pytest.mark.asyncio
async def test_relay_files_success(monkeypatch):
    from server.application.services import server_file_transfer_service as svc

    source = Server(id=1, name="a", domain="1.1.1.1")
    dest = Server(id=2, name="b", domain="2.2.2.2")

    async def fake_relay_tree(_src, _conn_dst, _dest, src_path, dst_path):
        assert src_path == "/www/a.txt"
        assert dst_path == "/www/backup/a.txt"
        return 42

    monkeypatch.setattr(svc, "_get_server", AsyncMock(side_effect=[source, dest]))
    monkeypatch.setattr(svc, "_ensure_remote_dir", AsyncMock())
    monkeypatch.setattr(svc, "_relay_tree", fake_relay_tree)

    class _Pool:
        async def acquire(self, _server):
            return object()

        async def release(self, _server_id):
            return None

    monkeypatch.setattr(svc, "ssh_pool", _Pool())

    class _SftpCtx:
        async def __aenter__(self):
            return AsyncMock()

        async def __aexit__(self, *_a):
            return None

    monkeypatch.setattr(svc, "sftp_session", lambda _conn: _SftpCtx())
    monkeypatch.setattr(
        svc,
        "_apply_dest_permissions",
        AsyncMock(return_value=[{"applied": True, "owner": "www", "group": "www"}]),
    )

    session = MagicMock()
    result = await relay_files(
        session,
        source_server_id=1,
        dest_server_id=2,
        sources=["/www/a.txt"],
        dest_dir="/www/backup",
    )

    assert result["mode"] == "relay"
    assert result["bytes_transferred"] == 42
    assert result["count"] == 1
    assert result["items"][0]["dest"] == "/www/backup/a.txt"


@pytest.mark.asyncio
async def test_cleanup_source_archive_uses_quoted_rm(monkeypatch):
    from server.application.services import server_file_transfer_service as svc

    server = Server(name="t", domain="1.2.3.4")
    exec_mock = AsyncMock(return_value={"exit_code": 0})
    monkeypatch.setattr(svc, "exec_ssh_command_with_fallback", exec_mock)

    await cleanup_source_archive(server, "/tmp/nexus-xfer-test.tar.gz")

    exec_mock.assert_awaited_once()
    cmd = exec_mock.await_args.args[1]
    assert "rm -f" in cmd
    assert "/tmp/nexus-xfer-test.tar.gz" in cmd


@pytest.mark.asyncio
async def test_extract_archive_on_server_uses_validating_decompress(monkeypatch):
    from server.application.services import server_file_transfer_service as svc

    server = Server(name="dest", domain="2.2.2.2")
    decompress = AsyncMock(return_value={"exit_code": 0})
    monkeypatch.setattr(svc, "run_remote_decompress", decompress)

    await _extract_archive_on_server(server, "/tmp/site.tar.gz", "/www/site", "tar.gz")

    decompress.assert_awaited_once_with(server, "/tmp/site.tar.gz", "/www/site", "tar.gz", timeout=600)


@pytest.mark.asyncio
async def test_extract_archive_on_server_raises_validation_error(monkeypatch):
    from server.application.services import server_file_transfer_service as svc

    server = Server(name="dest", domain="2.2.2.2")
    monkeypatch.setattr(
        svc,
        "run_remote_decompress",
        AsyncMock(return_value={"exit_code": -1, "stderr": "path traversal", "stdout": ""}),
    )

    with pytest.raises(ServerTransferError, match="path traversal"):
        await _extract_archive_on_server(server, "/tmp/site.tar.gz", "/www/site", "tar.gz")


@pytest.mark.asyncio
async def test_pull_transfer_package_rejects_traversal_dest(monkeypatch):
    from server.application.services import server_file_transfer_service as svc

    monkeypatch.setattr(
        svc,
        "load_package_meta",
        AsyncMock(return_value={"format": "tar.gz", "size": 10}),
    )
    monkeypatch.setattr(svc, "_get_server", AsyncMock(return_value=Server(id=2, name="dest", domain="2.2.2.2")))
    download = AsyncMock()
    monkeypatch.setattr(svc, "_download_url_to_remote_path", download)

    with pytest.raises(ServerTransferError, match=r"\.\."):
        await pull_transfer_package(
            MagicMock(),
            token="tok",
            dest_server_id=2,
            dest_path="/tmp/../etc/site.tar.gz",
        )

    download.assert_not_called()


@pytest.mark.asyncio
async def test_download_url_to_remote_path_uses_option_terminators(monkeypatch):
    from server.application.services import server_file_transfer_service as svc

    server = Server(id=2, name="dest", domain="2.2.2.2")
    calls: list[str] = []

    async def fake_exec(_server, command, *, timeout=120):
        calls.append(command)
        if command.startswith("curl "):
            return {"exit_code": 0}
        if command.startswith("mv "):
            return {"exit_code": 1, "stderr": "permission denied", "stdout": ""}
        return {"exit_code": 0}

    monkeypatch.setattr(svc, "exec_ssh_command_with_fallback", fake_exec)

    with pytest.raises(ServerTransferError, match="permission denied"):
        await _download_url_to_remote_path(
            server,
            download_url="https://nexus.example/download",
            dest_path="/tmp/-dest.tar.gz",
        )

    assert any(command.startswith("mv -f -- ") for command in calls)
    assert any(command.startswith("rm -f -- ") for command in calls)