import os
import types
import pytest
from typing import Optional
from pentest_mcp_server.server import PentestMCPServer
class DummySFTP:
def __init__(self):
self.operations = []
self.closed = False
async def put(self, local_path, remote_path):
self.operations.append(("put", local_path, remote_path))
async def get(self, remote_path, local_path):
os.makedirs(os.path.dirname(local_path) or ".", exist_ok=True)
with open(local_path, "wb") as f:
f.write(b"data")
self.operations.append(("get", remote_path, local_path))
async def exit(self):
self.closed = True
class DummySSHManager:
def __init__(self, sftp: DummySFTP):
self._sftp = sftp
self.run_commands = []
async def ensure_connected(self) -> bool:
return True
async def start_sftp_client(self):
return self._sftp
async def run_command(self, cmd: str, timeout: Optional[int] = None):
self.run_commands.append(cmd)
return types.SimpleNamespace(exit_status=0, stdout="", stderr="")
@pytest.mark.asyncio
async def test_upload_success(tmp_path):
server = PentestMCPServer()
sftp = DummySFTP()
server.ssh_manager = DummySSHManager(sftp)
server.tmux_manager = object()
async def ok_init():
return True
server._ensure_initialized = ok_init # type: ignore[attr-defined]
local_file = tmp_path / "sample.txt"
local_file.write_text("hello")
result = await server._handle_upload_file({
"local_path": str(local_file),
"remote_path": "/tmp/sample.txt"
})
assert result["status"] == "success"
assert ("put", str(local_file), "/tmp/sample.txt") in sftp.operations
assert sftp.closed is True
@pytest.mark.asyncio
async def test_upload_missing_local():
server = PentestMCPServer()
sftp = DummySFTP()
server.ssh_manager = DummySSHManager(sftp)
server.tmux_manager = object()
async def ok_init():
return True
server._ensure_initialized = ok_init # type: ignore[attr-defined]
result = await server._handle_upload_file({
"local_path": "C:/does/not/exist.txt",
"remote_path": "/tmp/sample.txt"
})
assert result["status"] == "error"
assert "Local file not found" in result["message"]
@pytest.mark.asyncio
async def test_download_success(tmp_path):
server = PentestMCPServer()
sftp = DummySFTP()
server.ssh_manager = DummySSHManager(sftp)
server.tmux_manager = object()
async def ok_init():
return True
server._ensure_initialized = ok_init # type: ignore[attr-defined]
local_dest = tmp_path / "nested" / "file.txt"
result = await server._handle_download_file({
"remote_path": "/tmp/remote.txt",
"local_path": str(local_dest)
})
assert result["status"] == "success"
assert local_dest.exists()
assert ("get", "/tmp/remote.txt", str(local_dest)) in sftp.operations
assert sftp.closed is True
@pytest.mark.asyncio
async def test_parse_handler_sftp_exit(tmp_path):
server = PentestMCPServer()
sftp = DummySFTP()
server.ssh_manager = DummySSHManager(sftp)
server.tmux_manager = object()
async def ok_init():
return True
server._ensure_initialized = ok_init # type: ignore[attr-defined]
async def fake_parse(tool, fmt, path):
return {"ok": True, "size": os.path.getsize(path)}
server._parse_output_file = fake_parse # type: ignore[attr-defined]
result = await server._handle_parse_tool_output({
"tool": "custom",
"file_path": "/tmp/remote.any",
"format": "txt"
})
assert result["status"] == "success"
assert sftp.closed is True
