diff --git a/aws_mcp_proxy/server.py b/aws_mcp_proxy/server.py index 98c06fe..fb62242 100644 --- a/aws_mcp_proxy/server.py +++ b/aws_mcp_proxy/server.py @@ -30,6 +30,7 @@ from aws_mcp_proxy.mcp_proxy_manager import McpProxyManager from aws_mcp_proxy.utils import ( create_transport_with_sigv4, + determine_aws_region, determine_service_name, ) from fastmcp.server.server import FastMCP @@ -46,9 +47,11 @@ async def setup_mcp_mode(local_mcp: FastMCP, args) -> None: # Validate and determine service service = determine_service_name(args.endpoint, args.service) - # Get profile and region - profile = args.profile or os.getenv('AWS_PROFILE') - region = os.getenv('AWS_REGION', 'us-west-2') + # Validate and determine region + region = determine_aws_region(args.endpoint, args.region) + + # Get profile + profile = args.profile # Log server configuration logger.info( @@ -57,7 +60,7 @@ async def setup_mcp_mode(local_mcp: FastMCP, args) -> None: logger.info('Running in MCP mode') # Create transport with SigV4 authentication - transport = create_transport_with_sigv4(args.endpoint, service, profile) + transport = create_transport_with_sigv4(args.endpoint, service, region, profile) # Create proxy with the transport proxy = FastMCP.as_proxy(transport) @@ -99,6 +102,13 @@ def parse_args(): parser.add_argument( '--profile', help='AWS profile to use (uses AWS_PROFILE environment variable if not provided)', + default=os.getenv('AWS_PROFILE', 'default'), + ) + + parser.add_argument( + '--region', + help='AWS region to use (uses AWS_REGION if not provided)', + default=os.getenv('AWS_REGION', 'default'), ) parser.add_argument( diff --git a/aws_mcp_proxy/sigv4_helper.py b/aws_mcp_proxy/sigv4_helper.py index 61ef5bf..e6bab65 100644 --- a/aws_mcp_proxy/sigv4_helper.py +++ b/aws_mcp_proxy/sigv4_helper.py @@ -17,7 +17,6 @@ import boto3 import httpx import logging -import os from botocore.auth import SigV4Auth from botocore.awsrequest import AWSRequest from botocore.credentials import Credentials @@ -150,15 +149,13 @@ def create_aws_session(profile: Optional[str] = None) -> boto3.Session: return session -def create_sigv4_auth( - service: str, profile: Optional[str] = None, region: Optional[str] = None -) -> SigV4HTTPXAuth: +def create_sigv4_auth(service: str, region: str, profile: Optional[str] = None) -> SigV4HTTPXAuth: """Create SigV4 authentication for AWS requests. Args: service: AWS service name for SigV4 signing profile: AWS profile to use (optional) - region: AWS region (defaults to AWS_REGION env var or us-west-2) + region: AWS region (defaults to AWS_REGION env var or us-east-1) Returns: SigV4HTTPXAuth instance @@ -170,10 +167,6 @@ def create_sigv4_auth( session = create_aws_session(profile) credentials = session.get_credentials() - # Get region from parameter, environment variable, or default - if not region: - region = os.environ.get('AWS_REGION', 'us-west-2') - # Create SigV4Auth with explicit credentials sigv4_auth = SigV4HTTPXAuth( credentials=credentials, @@ -187,8 +180,8 @@ def create_sigv4_auth( def create_sigv4_client( service: str, + region: str, profile: Optional[str] = None, - region: Optional[str] = None, headers: Optional[Dict[str, str]] = None, auth: Optional[httpx.Auth] = None, **kwargs: Any, @@ -198,7 +191,7 @@ def create_sigv4_client( Args: service: AWS service name for SigV4 signing profile: AWS profile to use (optional) - region: AWS region (optional, defaults to AWS_REGION env var or us-west-2) + region: AWS region (optional, defaults to AWS_REGION env var or us-east-1) headers: Headers to include in requests auth: Auth parameter (ignored as we provide our own) **kwargs: Additional arguments to pass to httpx.AsyncClient @@ -225,7 +218,7 @@ def create_sigv4_client( ) # Create SigV4 auth - sigv4_auth = create_sigv4_auth(service, profile, region) + sigv4_auth = create_sigv4_auth(service, region, profile) # Create the client with SigV4 auth and error handling event hook logger.info("Creating httpx.AsyncClient with SigV4 authentication for service '%s'", service) diff --git a/aws_mcp_proxy/utils.py b/aws_mcp_proxy/utils.py index 7ab322c..907c7ea 100644 --- a/aws_mcp_proxy/utils.py +++ b/aws_mcp_proxy/utils.py @@ -23,7 +23,10 @@ def create_transport_with_sigv4( - url: str, service: str, profile: Optional[str] = None + url: str, + service: str, + region: str, + profile: Optional[str] = None, ) -> StreamableHttpTransport: """Create a StreamableHttpTransport with SigV4 authentication. @@ -31,6 +34,7 @@ def create_transport_with_sigv4( url: The endpoint URL service: AWS service name for SigV4 signing profile: AWS profile to use (optional) + region: AWS region to use (Optional) Returns: StreamableHttpTransport instance with SigV4 authentication @@ -42,7 +46,12 @@ def client_factory( auth: Optional[httpx.Auth] = None, ) -> httpx.AsyncClient: return create_sigv4_client( - service=service, profile=profile, headers=headers, timeout=timeout, auth=auth + service=service, + profile=profile, + region=region, + headers=headers, + timeout=timeout, + auth=auth, ) return StreamableHttpTransport( @@ -81,3 +90,35 @@ def determine_service_name(endpoint: str, service: Optional[str] = None) -> str: 'Please provide the service name explicitly using --service argument.' ) return determined_service + + +def determine_aws_region(endpoint: str, region: Optional[str] = None) -> str: + """Validate and determine the AWS region. + + Args: + endpoint: The endpoint URL + region: Optional region name + + Returns: + Validated AWS region + + Raises: + ValueError: If region cannot be determined + """ + if region: + return region + + # Parse AWS region from endpoint URL + parsed = urlparse(endpoint) + hostname = parsed.hostname or '' + + # Extract region name (pattern: service.region.api.aws or service-name.region.api.aws) + region_match = re.search(r'\.([a-z0-9-]+)\.api\.aws', hostname) + determined_region = region_match.group(1) if region_match else None + + if not determined_region: + raise ValueError( + f"Could not determine AWS region from endpoint '{endpoint}'. " + 'Please provide the region explicitly using --region argument.' + ) + return determined_region diff --git a/tests/unit/test_server.py b/tests/unit/test_server.py index 76e10c1..62bd1dc 100644 --- a/tests/unit/test_server.py +++ b/tests/unit/test_server.py @@ -36,6 +36,7 @@ async def test_setup_mcp_mode( mock_args = Mock() mock_args.endpoint = 'https://test.example.com' mock_args.service = 'test-service' + mock_args.region = 'us-east-1' mock_args.profile = None mock_args.read_only = True @@ -71,6 +72,7 @@ async def test_setup_mcp_mode_with_tools( mock_args = Mock() mock_args.endpoint = 'https://test.example.com' mock_args.service = 'test-service' + mock_args.region = 'us-east-1' mock_args.profile = None mock_args.read_only = True @@ -106,6 +108,7 @@ async def test_setup_mcp_mode_tool_registration_error( mock_args = Mock() mock_args.endpoint = 'https://test.example.com' mock_args.service = 'test-service' + mock_args.region = 'us-east-1' mock_args.profile = None mock_args.read_only = True @@ -188,8 +191,7 @@ def test_create_sigv4_client(self, mock_sigv4_auth, mock_async_client, mock_sess mock_session.return_value = mock_session_instance # Act - with patch.dict('os.environ', {'AWS_REGION': 'us-west-2'}): - create_sigv4_client(service='test-service', profile='test-profile') + create_sigv4_client(service='test-service', region='us-west-2', profile='test-profile') # Assert mock_session.assert_called_once_with(profile_name='test-profile') @@ -206,7 +208,7 @@ def test_create_sigv4_client_no_credentials(self, mock_session): # Act & Assert with pytest.raises(ValueError) as exc_info: - create_sigv4_client(service='test-service') + create_sigv4_client(service='test-service', region='test-region') assert 'No AWS credentials found' in str(exc_info.value) def test_main_module_execution(self): diff --git a/tests/unit/test_sigv4_helper.py b/tests/unit/test_sigv4_helper.py index c802fcc..67abb36 100644 --- a/tests/unit/test_sigv4_helper.py +++ b/tests/unit/test_sigv4_helper.py @@ -16,7 +16,6 @@ import httpx import json -import os import pytest from aws_mcp_proxy.sigv4_helper import ( SigV4HTTPXAuth, @@ -238,34 +237,12 @@ def test_create_sigv4_auth_default(self, mock_create_session): mock_create_session.return_value = mock_session # Test auth creation - result = create_sigv4_auth('test-service') + result = create_sigv4_auth('test-service', 'test-region') # Verify auth was created correctly assert isinstance(result, SigV4HTTPXAuth) assert result.service == 'test-service' - assert result.region == 'us-west-2' # default region - assert result.credentials == mock_credentials - - @patch('aws_mcp_proxy.sigv4_helper.create_aws_session') - @patch.dict(os.environ, {'AWS_REGION': 'eu-west-1'}) - def test_create_sigv4_auth_with_env_region(self, mock_create_session): - """Test creating SigV4 auth with region from environment variable.""" - # Mock session and credentials - mock_session = Mock() - mock_credentials = Mock() - mock_credentials.access_key = 'test_access_key' - mock_credentials.secret_key = 'test_secret_key' - mock_credentials.token = None - mock_session.get_credentials.return_value = mock_credentials - mock_create_session.return_value = mock_session - - # Test auth creation - result = create_sigv4_auth('test-service', profile='test-profile') - - # Verify auth was created with environment region - assert isinstance(result, SigV4HTTPXAuth) - assert result.service == 'test-service' - assert result.region == 'eu-west-1' # from environment + assert result.region == 'test-region' # default region assert result.credentials == mock_credentials @patch('aws_mcp_proxy.sigv4_helper.create_aws_session') @@ -304,10 +281,10 @@ def test_create_sigv4_client_default(self, mock_client_class, mock_create_auth): mock_client_class.return_value = mock_client # Test client creation - result = create_sigv4_client(service='test-service') + result = create_sigv4_client(service='test-service', region='test-region') # Verify client was created correctly - mock_create_auth.assert_called_once_with('test-service', None, None) + mock_create_auth.assert_called_once_with('test-service', 'test-region', None) # Check that AsyncClient was called with correct parameters call_args = mock_client_class.call_args @@ -330,7 +307,9 @@ def test_create_sigv4_client_with_custom_headers(self, mock_client_class, mock_c # Test client creation with custom headers custom_headers = {'Custom-Header': 'custom-value'} - result = create_sigv4_client(service='test-service', headers=custom_headers) + result = create_sigv4_client( + service='test-service', region='test-region', headers=custom_headers + ) # Verify client was created with merged headers call_args = mock_client_class.call_args @@ -359,7 +338,7 @@ def test_create_sigv4_client_with_custom_service_and_region( ) # Verify auth was created with custom parameters - mock_create_auth.assert_called_once_with('custom-service', 'test-profile', 'us-east-1') + mock_create_auth.assert_called_once_with('custom-service', 'us-east-1', 'test-profile') assert result == mock_client @patch('aws_mcp_proxy.sigv4_helper.create_sigv4_auth') @@ -374,7 +353,10 @@ def test_create_sigv4_client_with_kwargs(self, mock_client_class, mock_create_au # Test client creation with additional kwargs result = create_sigv4_client( - service='test-service', verify=False, proxies={'http': 'http://proxy:8080'} + service='test-service', + region='test-region', + verify=False, + proxies={'http': 'http://proxy:8080'}, ) # Verify client was created with additional kwargs @@ -409,7 +391,7 @@ def test_create_sigv4_client_with_prompt_context(self, mock_client_class, mock_c ) # Verify client was created correctly with prompt context - mock_create_auth.assert_called_once_with('test-service', None, 'us-west-2') + mock_create_auth.assert_called_once_with('test-service', 'us-west-2', None) # Check that AsyncClient was called with correct parameters including prompt headers call_args = mock_client_class.call_args diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 2d35f57..4e2f440 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -17,6 +17,7 @@ import pytest from aws_mcp_proxy.utils import ( create_transport_with_sigv4, + determine_aws_region, determine_service_name, ) from fastmcp.client.transports import StreamableHttpTransport @@ -35,8 +36,9 @@ def test_create_transport_with_sigv4(self, mock_create_sigv4_client): url = 'https://test-service.us-west-2.api.aws/mcp' service = 'test-service' profile = 'test-profile' + region = 'us-east-1' - result = create_transport_with_sigv4(url, service, profile) + result = create_transport_with_sigv4(url, service, region, profile) # Verify result is StreamableHttpTransport assert isinstance(result, StreamableHttpTransport) @@ -54,6 +56,7 @@ def test_create_transport_with_sigv4(self, mock_create_sigv4_client): mock_create_sigv4_client.assert_called_once_with( service=service, profile=profile, + region=region, headers={'test': 'header'}, timeout=Timeout(30.0), auth=None, @@ -67,8 +70,9 @@ def test_create_transport_with_sigv4_no_profile(self, mock_create_sigv4_client): """Test creating transport without profile.""" url = 'https://test-service.us-west-2.api.aws/mcp' service = 'test-service' + region = 'test-region' - result = create_transport_with_sigv4(url, service) + result = create_transport_with_sigv4(url, service, region) # Test that the httpx_client_factory calls create_sigv4_client correctly # We need to access the factory through the transport's internal structure @@ -77,7 +81,7 @@ def test_create_transport_with_sigv4_no_profile(self, mock_create_sigv4_client): factory(headers=None, timeout=None, auth=None) mock_create_sigv4_client.assert_called_once_with( - service=service, profile=None, headers=None, timeout=None, auth=None + service=service, region=region, profile=None, headers=None, timeout=None, auth=None ) else: # If we can't access the factory directly, just verify the transport was created @@ -144,3 +148,45 @@ def test_validate_service_name_invalid_url_failure(self): assert 'Could not determine AWS service name' in str(exc_info.value) assert endpoint in str(exc_info.value) assert '--service argument' in str(exc_info.value) + + +class TestDetermineRegion: + """Test cases for determine_aws_region function.""" + + def test_determine_region_with_region(self): + """Test determination when region is provided.""" + endpoint = 'https://mcp.us-east-1.api.aws/mcp' + region = 'custom-region' + + result = determine_aws_region(endpoint, region) + + assert result == region + + def test_determine_region_without_region_success(self): + """Test determination when region is not provided but can be parsed.""" + endpoint = 'https://mcp.us-east-1.api.aws/mcp' + expected_region = 'us-east-1' + + result = determine_aws_region(endpoint) + + assert result == expected_region + + def test_determine_region_with_complex_service_name(self): + """Test parsing region from endpoint with complex service name.""" + endpoint = 'https://eks-mcp-beta.us-west-2.api.aws/mcp' + expected_region = 'us-west-2' + + result = determine_aws_region(endpoint) + + assert result == expected_region + + def test_determine_region_without_region_failure(self): + """Test determination when region cannot be determined.""" + endpoint = 'https://service.example.com' + + with pytest.raises(ValueError) as exc_info: + determine_aws_region(endpoint) + + assert 'Could not determine AWS region' in str(exc_info.value) + assert endpoint in str(exc_info.value) + assert '--region argument' in str(exc_info.value)