diff --git a/src/bedrock_agentcore_starter_toolkit/operations/runtime/launch.py b/src/bedrock_agentcore_starter_toolkit/operations/runtime/launch.py index 5e29db81..4b1c5d14 100644 --- a/src/bedrock_agentcore_starter_toolkit/operations/runtime/launch.py +++ b/src/bedrock_agentcore_starter_toolkit/operations/runtime/launch.py @@ -433,18 +433,46 @@ def _execute_codebuild_workflow( if not ecr_only: _ensure_execution_role(agent_config, project_config, config_path, agent_name, region, account_id) - # Prepare CodeBuild + # Enhanced CodeBuild preparation with cross-account support log.info("Preparing CodeBuild project and uploading source...") - codebuild_service = CodeBuildService(session) - # Use cached CodeBuild role from config if available + # Get CodeBuild execution role from config + codebuild_execution_role = None if hasattr(agent_config, "codebuild") and agent_config.codebuild.execution_role: - log.info("Using CodeBuild role from config: %s", agent_config.codebuild.execution_role) codebuild_execution_role = agent_config.codebuild.execution_role - else: + log.info("Using CodeBuild role from config: %s", codebuild_execution_role) + + # Detect cross-account scenario + build_account = None + deployment_account = session.client("sts").get_caller_identity()["Account"] + + if codebuild_execution_role: + try: + # Extract account from role ARN: arn:aws:iam::ACCOUNT-ID:role/ROLE-NAME + build_account = codebuild_execution_role.split(":")[4] + + if build_account != deployment_account: + log.info("Cross-account CodeBuild detected:") + log.info(" Deployment account: %s", deployment_account) + log.info(" Build account: %s", build_account) + log.info(" CodeBuild role: %s", codebuild_execution_role) + else: + log.info("Same-account CodeBuild (role in deployment account)") + build_account = None # Treat as same-account + except (IndexError, AttributeError): + log.warning("Invalid CodeBuild role ARN format: %s", codebuild_execution_role) + build_account = None + + # Initialize CodeBuild service with role information + codebuild_service = CodeBuildService(session, codebuild_execution_role) + + # Create or use existing CodeBuild execution role + if not codebuild_execution_role: + # No role specified - create one in deployment account codebuild_execution_role = codebuild_service.create_codebuild_execution_role( account_id=account_id, ecr_repository_arn=ecr_repository_arn, agent_name=agent_name ) + log.info("Created CodeBuild role in deployment account: %s", codebuild_execution_role) source_location = codebuild_service.upload_source(agent_name=agent_name) diff --git a/src/bedrock_agentcore_starter_toolkit/services/codebuild.py b/src/bedrock_agentcore_starter_toolkit/services/codebuild.py index 328aad09..0c82fc43 100644 --- a/src/bedrock_agentcore_starter_toolkit/services/codebuild.py +++ b/src/bedrock_agentcore_starter_toolkit/services/codebuild.py @@ -7,7 +7,7 @@ import time import zipfile from pathlib import Path -from typing import List +from typing import List, Optional import boto3 from botocore.exceptions import ClientError @@ -18,19 +18,76 @@ class CodeBuildService: """Service for managing CodeBuild projects and builds for ARM64.""" - def __init__(self, session: boto3.Session): - """Initialize CodeBuild service with AWS session.""" + def __init__(self, session: boto3.Session, codebuild_role_arn: Optional[str] = None): + """Initialize CodeBuild service with AWS session. + + Args: + session: Primary AWS session for deployment account + codebuild_role_arn: Optional CodeBuild execution role ARN (for cross-account) + """ self.session = session - self.client = session.client("codebuild") - self.s3_client = session.client("s3") - self.iam_client = session.client("iam") + self.codebuild_role_arn = codebuild_role_arn self.logger = logging.getLogger(__name__) + + # Determine if this is cross-account CodeBuild + self.build_account = self._extract_build_account() + self.deployment_account = session.client("sts").get_caller_identity()["Account"] + self.is_cross_account_codebuild = ( + self.build_account is not None and self.build_account != self.deployment_account + ) + + # Create appropriate session for CodeBuild operations + if self.is_cross_account_codebuild: + self.build_session = self._create_build_session() + self.client = self.build_session.client("codebuild") + self.s3_client = self.build_session.client("s3") + self.iam_client = self.build_session.client("iam") + self.logger.info("CodeBuild initialized in cross-account mode (account: %s)", self.build_account) + else: + self.build_session = session + self.client = session.client("codebuild") + self.s3_client = session.client("s3") + self.iam_client = session.client("iam") + self.logger.info("CodeBuildService initialized in same-account mode") + self.source_bucket = None - self.account_id = session.client("sts").get_caller_identity()["Account"] + + def _extract_build_account(self) -> Optional[str]: + """Extract build account ID from CodeBuild role ARN.""" + if not self.codebuild_role_arn: + return None + try: + # ARN format: arn:aws:iam::ACCOUNT-ID:role/ROLE-NAME + return self.codebuild_role_arn.split(":")[4] + except (IndexError, AttributeError): + self.logger.warning("Invalid CodeBuild role ARN format: %s", self.codebuild_role_arn) + return None + + def _create_build_session(self) -> boto3.Session: + """Create AWS session by assuming the CodeBuild role.""" + if not self.codebuild_role_arn: + return self.session + + try: + sts_client = self.session.client("sts") + response = sts_client.assume_role( + RoleArn=self.codebuild_role_arn, RoleSessionName=f"bedrock-agentcore-build-{int(time.time())}" + ) + + credentials = response["Credentials"] + return boto3.Session( + aws_access_key_id=credentials["AccessKeyId"], + aws_secret_access_key=credentials["SecretAccessKey"], + aws_session_token=credentials["SessionToken"], + region_name=self.session.region_name, + ) + except Exception as e: + self.logger.error("Failed to assume CodeBuild role %s: %s", self.codebuild_role_arn, e) + raise RuntimeError(f"Failed to assume CodeBuild role {self.codebuild_role_arn}: {e}") from e def get_source_bucket_name(self, account_id: str) -> str: """Get S3 bucket name for CodeBuild sources.""" - region = self.session.region_name + region = self.build_session.region_name return f"bedrock-agentcore-codebuild-sources-{account_id}-{region}" def ensure_source_bucket(self, account_id: str) -> str: @@ -50,7 +107,7 @@ def ensure_source_bucket(self, account_id: str) -> str: ) from e # Create bucket (no ExpectedBucketOwner needed for create_bucket) - region = self.session.region_name + region = self.build_session.region_name if region == "us-east-1": self.s3_client.create_bucket(Bucket=bucket_name) else: @@ -72,7 +129,9 @@ def ensure_source_bucket(self, account_id: str) -> str: def upload_source(self, agent_name: str) -> str: """Upload current directory to S3, respecting .dockerignore patterns.""" - account_id = self.account_id + # Use build account for S3 bucket (cross-account) or deployment account (same-account) + account_id = self.build_account or self.deployment_account + self.logger.info("Using account %s for S3 bucket", account_id) bucket_name = self.ensure_source_bucket(account_id) self.source_bucket = bucket_name diff --git a/tests/operations/runtime/test_launch_cross_account.py b/tests/operations/runtime/test_launch_cross_account.py new file mode 100644 index 00000000..63eccc1f --- /dev/null +++ b/tests/operations/runtime/test_launch_cross_account.py @@ -0,0 +1,155 @@ +"""Simple tests for cross-account launch functionality.""" + +from unittest.mock import Mock, patch +import pytest + +from bedrock_agentcore_starter_toolkit.operations.runtime.launch import _execute_codebuild_workflow +from bedrock_agentcore_starter_toolkit.utils.runtime.schema import ( + AWSConfig, + BedrockAgentCoreAgentSchema, + BedrockAgentCoreConfigSchema, + NetworkConfiguration, + ObservabilityConfig, + BedrockAgentCoreDeploymentInfo, + CodeBuildConfig, +) + + +class TestLaunchCrossAccount: + """Test cross-account functionality in launch operations.""" + + def test_codebuild_service_initialization_cross_account(self, tmp_path): + """Test CodeBuildService is initialized with cross-account role.""" + # Create agent config with cross-account CodeBuild role + aws_config = AWSConfig( + account="123456789012", + region="us-west-2", + execution_role="arn:aws:iam::123456789012:role/ExecutionRole", + ecr_repository="123456789012.dkr.ecr.us-west-2.amazonaws.com/test-repo", + network_configuration=NetworkConfiguration(), + observability=ObservabilityConfig(), + ) + + codebuild_config = CodeBuildConfig() + codebuild_config.execution_role = "arn:aws:iam::987654321098:role/BuildRole" + + agent_config = BedrockAgentCoreAgentSchema( + name="test-agent", + entrypoint="test.py", + aws=aws_config, + bedrock_agentcore=BedrockAgentCoreDeploymentInfo(), + codebuild=codebuild_config, + ) + + project_config = BedrockAgentCoreConfigSchema( + default_agent="test-agent", + agents={"test-agent": agent_config} + ) + + config_path = tmp_path / "config.yaml" + + with patch('bedrock_agentcore_starter_toolkit.operations.runtime.launch.CodeBuildService') as mock_cb_service, \ + patch('bedrock_agentcore_starter_toolkit.operations.runtime.launch._ensure_ecr_repository') as mock_ecr, \ + patch('bedrock_agentcore_starter_toolkit.operations.runtime.launch._ensure_execution_role') as mock_role, \ + patch("boto3.Session") as mock_session: + + # Setup mocks + mock_ecr.return_value = "123456789012.dkr.ecr.us-west-2.amazonaws.com/test-repo" + mock_role.return_value = "arn:aws:iam::123456789012:role/ExecutionRole" + + mock_service_instance = Mock() + mock_service_instance.upload_source.return_value = "s3://bucket/source.zip" + mock_service_instance.create_or_update_project.return_value = "test-project" + mock_service_instance.start_build.return_value = "build-123" + mock_service_instance.wait_for_completion.return_value = None + mock_service_instance.source_bucket = "test-bucket" + mock_cb_service.return_value = mock_service_instance + + deployment_session = Mock() + # Mock STS client for account detection + mock_sts = Mock() + mock_sts.get_caller_identity.return_value = {"Account": "123456789012"} + deployment_session.client.return_value = mock_sts + mock_session.return_value = deployment_session + + # Execute + _execute_codebuild_workflow( + config_path=config_path, + agent_name="test-agent", + agent_config=agent_config, + project_config=project_config, + ecr_only=False + ) + + # Verify CodeBuildService was called with cross-account role + mock_cb_service.assert_called_once_with( + deployment_session, + "arn:aws:iam::987654321098:role/BuildRole" + ) + + def test_codebuild_service_initialization_same_account(self, tmp_path): + """Test CodeBuildService is initialized without cross-account role.""" + # Create agent config without cross-account CodeBuild role + aws_config = AWSConfig( + account="123456789012", + region="us-west-2", + execution_role="arn:aws:iam::123456789012:role/ExecutionRole", + ecr_repository="123456789012.dkr.ecr.us-west-2.amazonaws.com/test-repo", + network_configuration=NetworkConfiguration(), + observability=ObservabilityConfig(), + ) + + codebuild_config = CodeBuildConfig() + # No execution_role set - same account scenario + + agent_config = BedrockAgentCoreAgentSchema( + name="test-agent", + entrypoint="test.py", + aws=aws_config, + bedrock_agentcore=BedrockAgentCoreDeploymentInfo(), + codebuild=codebuild_config, + ) + + project_config = BedrockAgentCoreConfigSchema( + default_agent="test-agent", + agents={"test-agent": agent_config} + ) + + config_path = tmp_path / "config.yaml" + + with patch('bedrock_agentcore_starter_toolkit.operations.runtime.launch.CodeBuildService') as mock_cb_service, \ + patch('bedrock_agentcore_starter_toolkit.operations.runtime.launch._ensure_ecr_repository') as mock_ecr, \ + patch('bedrock_agentcore_starter_toolkit.operations.runtime.launch._ensure_execution_role') as mock_role, \ + patch("boto3.Session") as mock_session: + + # Setup mocks + mock_ecr.return_value = "123456789012.dkr.ecr.us-west-2.amazonaws.com/test-repo" + mock_role.return_value = "arn:aws:iam::123456789012:role/ExecutionRole" + + mock_service_instance = Mock() + mock_service_instance.create_codebuild_execution_role.return_value = "arn:aws:iam::123456789012:role/CodeBuildRole" + mock_service_instance.upload_source.return_value = "s3://bucket/source.zip" + mock_service_instance.create_or_update_project.return_value = "test-project" + mock_service_instance.start_build.return_value = "build-123" + mock_service_instance.wait_for_completion.return_value = None + mock_service_instance.source_bucket = "test-bucket" + mock_cb_service.return_value = mock_service_instance + + deployment_session = Mock() + # Mock STS client for account detection + mock_sts = Mock() + mock_sts.get_caller_identity.return_value = {"Account": "123456789012"} + deployment_session.client.return_value = mock_sts + mock_session.return_value = deployment_session + + # Execute + _execute_codebuild_workflow( + config_path=config_path, + agent_name="test-agent", + agent_config=agent_config, + project_config=project_config, + ecr_only=False + ) + + # Verify CodeBuildService was called without cross-account role + mock_cb_service.assert_called_once_with(deployment_session, None) diff --git a/tests/services/test_codebuild.py b/tests/services/test_codebuild.py index 9ecd9f1e..899544a0 100644 --- a/tests/services/test_codebuild.py +++ b/tests/services/test_codebuild.py @@ -84,7 +84,9 @@ def client_factory(service_name): assert service.s3_client == mock_s3 assert service.iam_client == mock_iam assert service.source_bucket is None - assert service.account_id == "123456789012" # Verify account_id is stored + assert service.deployment_account == "123456789012" # Verify deployment account is stored + assert service.build_account is None # No cross-account role provided + assert service.is_cross_account_codebuild is False def test_get_source_bucket_name(self, codebuild_service): """Test S3 bucket name generation.""" @@ -188,6 +190,7 @@ def test_upload_source_success( result = codebuild_service.upload_source("test-agent") expected_key = "test-agent/source.zip" + # Use deployment account since no cross-account role provided expected_s3_url = f"s3://bedrock-agentcore-codebuild-sources-123456789012-us-west-2/{expected_key}" assert result == expected_s3_url diff --git a/tests/services/test_codebuild_cross_account.py b/tests/services/test_codebuild_cross_account.py new file mode 100644 index 00000000..e27dd729 --- /dev/null +++ b/tests/services/test_codebuild_cross_account.py @@ -0,0 +1,110 @@ +"""Simple tests for cross-account CodeBuild functionality.""" + +from unittest.mock import Mock, patch +import pytest + +from bedrock_agentcore_starter_toolkit.services.codebuild import CodeBuildService + + +class TestCodeBuildCrossAccount: + """Test cross-account CodeBuild functionality.""" + + def test_init_same_account(self): + """Test initialization for same-account scenario.""" + mock_session = Mock() + mock_session.region_name = "us-west-2" + + mock_sts = Mock() + mock_sts.get_caller_identity.return_value = {"Account": "123456789012"} + mock_session.client.return_value = mock_sts + + service = CodeBuildService(mock_session) + + assert service.deployment_account == "123456789012" + assert service.build_account is None + assert service.is_cross_account_codebuild is False + + def test_init_cross_account(self): + """Test initialization for cross-account scenario.""" + mock_session = Mock() + mock_session.region_name = "us-west-2" + + mock_sts = Mock() + mock_sts.get_caller_identity.return_value = {"Account": "123456789012"} + + mock_build_session = Mock() + mock_build_session.region_name = "us-west-2" + + def client_factory(service_name): + if service_name == "sts": + return mock_sts + return Mock() + + mock_session.client = client_factory + + with patch.object(CodeBuildService, '_create_build_session', return_value=mock_build_session): + service = CodeBuildService(mock_session, "arn:aws:iam::987654321098:role/BuildRole") + + assert service.deployment_account == "123456789012" + assert service.build_account == "987654321098" + assert service.is_cross_account_codebuild is True + + def test_extract_build_account(self): + """Test build account extraction from role ARN.""" + mock_session = Mock() + mock_session.region_name = "us-west-2" + + mock_sts = Mock() + mock_sts.get_caller_identity.return_value = {"Account": "123456789012"} + mock_session.client.return_value = mock_sts + + # Test valid ARN + with patch.object(CodeBuildService, '_create_build_session'): + service = CodeBuildService(mock_session, "arn:aws:iam::987654321098:role/BuildRole") + assert service._extract_build_account() == "987654321098" + + # Test invalid ARN + with patch.object(CodeBuildService, '_create_build_session'): + service = CodeBuildService(mock_session, "invalid-arn") + assert service._extract_build_account() is None + + # Test no ARN + service = CodeBuildService(mock_session) + assert service._extract_build_account() is None + + def test_upload_source_account_selection(self): + """Test that upload_source uses correct account.""" + mock_session = Mock() + mock_session.region_name = "us-west-2" + + mock_sts = Mock() + mock_sts.get_caller_identity.return_value = {"Account": "123456789012"} + + mock_s3 = Mock() + mock_s3.head_bucket.return_value = {} + + def client_factory(service_name): + if service_name == "sts": + return mock_sts + elif service_name == "s3": + return mock_s3 + return Mock() + + mock_session.client = client_factory + + # Same account scenario + service = CodeBuildService(mock_session) + + with patch('os.walk', return_value=[(".", [], ["test.py"])]), \ + patch('zipfile.ZipFile'), \ + patch('tempfile.NamedTemporaryFile'), \ + patch('os.unlink'): + + service.upload_source("test-agent") + + # Should use deployment account + expected_bucket = "bedrock-agentcore-codebuild-sources-123456789012-us-west-2" + mock_s3.head_bucket.assert_called_with( + Bucket=expected_bucket, + ExpectedBucketOwner="123456789012" + )