Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 14 additions & 4 deletions aws_mcp_proxy/server.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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(
Expand All @@ -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)
Expand Down Expand Up @@ -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)',
Comment thread
acmlau marked this conversation as resolved.
default=os.getenv('AWS_REGION', 'default'),
)

parser.add_argument(
Expand Down
17 changes: 5 additions & 12 deletions aws_mcp_proxy/sigv4_helper.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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')

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We should pass the region here, from the parsed args rather than fetching it again.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As a standalone function, it makes sense to check whether region is not set.
I did notice that we never passed the region as a parameter when creating the sigv4_auth https://github.com/aws/aws-mcp-proxy/blob/main/aws_mcp_proxy/server.py#L60, will create new PR to address that + removing some unused parameters

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

By the time we create sigv4 auth, region should have been resolved, and region should not be None


# Create SigV4Auth with explicit credentials
sigv4_auth = SigV4HTTPXAuth(
credentials=credentials,
Expand All @@ -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,
Expand All @@ -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
Expand All @@ -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)
Expand Down
45 changes: 43 additions & 2 deletions aws_mcp_proxy/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,14 +23,18 @@


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.

Args:
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
Expand All @@ -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(
Expand Down Expand Up @@ -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
8 changes: 5 additions & 3 deletions tests/unit/test_server.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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')
Expand All @@ -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):
Expand Down
44 changes: 13 additions & 31 deletions tests/unit/test_sigv4_helper.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@

import httpx
import json
import os
import pytest
from aws_mcp_proxy.sigv4_helper import (
SigV4HTTPXAuth,
Expand Down Expand Up @@ -238,34 +237,12 @@
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
Comment thread Fixed

Check notice

Code scanning / Bandit

Use of assert detected. The enclosed code will be removed when compiling to optimised byte code. Note test

Use of assert detected. The enclosed code will be removed when compiling to optimised byte code.
assert result.credentials == mock_credentials

@patch('aws_mcp_proxy.sigv4_helper.create_aws_session')
Expand Down Expand Up @@ -304,10 +281,10 @@
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
Expand All @@ -330,7 +307,9 @@

# 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
Expand Down Expand Up @@ -359,7 +338,7 @@
)

# 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')
Expand All @@ -374,7 +353,10 @@

# 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
Expand Down Expand Up @@ -409,7 +391,7 @@
)

# 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
Expand Down
Loading