""" Copyright 2024-2026 ChatterMate Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. """ from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest from fastapi import HTTPException from app.core.config import settings from app.core.outbound_policy import ( OutboundTargetBlocked, check_outbound_target, require_public_target, ) from app.knowledge import url_safety from app.models.mcp_tool import MCPTransportType from app.services import db_connector_service from app.services.db_connector_service import DBConnectorConfig, first_hop from app.tools.mcp_manager import MCPToolsManager from app.utils.agno_utils import create_model INTERNAL_TARGETS = [ "http://127.0.0.1:11434/v1", "http://localhost:11434/v1", "localhost:11434/v1", # no scheme: must not parse as hostless "http://169.254.169.254/latest/meta-data/", "https://10.1.2.3/mcp", "192.168.0.10", # bare database host "::1", # bare IPv6 host "http://[fd00::1]/mcp", ] @pytest.fixture def public_only(monkeypatch): monkeypatch.setattr(settings, "OUTBOUND_URL_POLICY", "public_only") @pytest.fixture def policy_any(monkeypatch): monkeypatch.setattr(settings, "OUTBOUND_URL_POLICY", "any") def _resolves_to(address): return patch.object(url_safety.socket, "getaddrinfo", return_value=[(2, 1, 6, "", (address, 0))]) class TestPolicy: def test_default_is_any(self): assert settings.OUTBOUND_URL_POLICY == "any" @pytest.mark.parametrize("target", INTERNAL_TARGETS) def test_any_allows_internal_targets(self, policy_any, target): check_outbound_target(target, "Target") @pytest.mark.parametrize("target", INTERNAL_TARGETS) def test_public_only_refuses_internal_targets(self, public_only, target): with pytest.raises(OutboundTargetBlocked, match="Target must be a public address"): check_outbound_target(target, "Target") def test_public_only_refuses_a_name_resolving_inside(self, public_only): with _resolves_to("10.0.0.9"), pytest.raises(OutboundTargetBlocked): check_outbound_target("db.internal.example.com", "Database host") def test_public_only_allows_a_public_target(self, public_only): with _resolves_to("93.184.216.34"): check_outbound_target("https://api.example.com/v1", "Base URL") check_outbound_target("db.example.com", "Database host") def test_an_empty_target_is_left_to_field_validation(self, public_only): check_outbound_target(None, "Base URL") check_outbound_target("", "Base URL") def test_require_public_target_is_a_400(self, public_only): with pytest.raises(HTTPException) as exc: require_public_target("http://127.0.0.1", "MCP server URL") assert exc.value.status_code == 400 assert "MCP server URL must be a public address" in exc.value.detail class TestCheckedAgainBeforeConnecting: """Saving is not enough: a hostname can be re-pointed after it was saved.""" def test_ai_base_url(self, public_only): with pytest.raises(HTTPException) as exc: create_model("OPENAI_COMPATIBLE", "key", "m", base_url="http://127.0.0.1:11434/v1") assert "must be a public address" in exc.value.detail def test_ai_base_url_allowed_under_any(self, policy_any): model = create_model("OPENAI_COMPATIBLE", "key", "m", base_url="http://127.0.0.1:11434/v1") assert model.base_url == "http://127.0.0.1:11434/v1" @pytest.mark.asyncio async def test_mcp_server_url(self, public_only): config = SimpleNamespace( name="Internal", transport_type=MCPTransportType.HTTP, url="http://169.254.169.254/mcp", headers=None, timeout=None, sse_read_timeout=None, terminate_on_close=True, ) with patch("app.tools.mcp_manager.MCPTools") as mcp_tools: result = await MCPToolsManager(connect_budget=5).test_tool_config(config) assert result["success"] is False assert "MCP server URL must be a public address" in result["error"] mcp_tools.assert_not_called() def _db_config(self, **overrides): values = dict( id=None, organization_id=None, name="c", engine="postgresql", host="10.0.0.5", port=5432, database="d", username="u", password="p", ) values.update(overrides) return DBConnectorConfig(**values) def test_database_host(self, public_only): with patch.object(db_connector_service, "_driver_connect") as driver: with pytest.raises(OutboundTargetBlocked, match="Database host"): with db_connector_service._connection(self._db_config()): pass driver.assert_not_called() def test_ssh_host_is_the_first_hop_when_tunnelling(self, public_only): config = self._db_config(ssh_enabled=True, ssh_host="127.0.0.1", ssh_username="u") with patch.object(db_connector_service, "_ssh_tunnel") as tunnel: with pytest.raises(OutboundTargetBlocked, match="SSH host"): with db_connector_service._connection(config): pass tunnel.assert_not_called() def test_private_database_behind_a_public_bastion_is_allowed(self, public_only): config = self._db_config(ssh_enabled=True, ssh_host="bastion.example.com", ssh_username="u") tunnel = MagicMock() tunnel.return_value.__enter__.return_value = ("127.0.0.1", 40000) with _resolves_to("93.184.216.34"), \ patch.object(db_connector_service, "_ssh_tunnel", tunnel), \ patch.object(db_connector_service, "_driver_connect") as driver: with db_connector_service._connection(config): pass driver.assert_called_once_with(config, "127.0.0.1", 40000) def test_first_hop_picks_the_bastion_only_when_tunnelling(): assert first_hop(False, "bastion", "db") == ("db", "Database host") assert first_hop(True, "bastion", "db") == ("bastion", "SSH host")