feat: add support for policies and certificates (aws#371)

Author: rvkandury

src/bedrock_agentcore/tools/__init__.py

@@ -7,14 +7,20 @@
     BrowserConfiguration,
     BrowserExtension,
     BrowserSigningConfiguration,
+    Certificate,
+    CertificateLocation,
     CodeInterpreterConfiguration,
+    EnterprisePolicy,
+    EnterprisePolicyS3Location,
     ExtensionS3Location,
     ExternalProxy,
     NetworkConfiguration,
     ProfileConfiguration,
     ProxyConfiguration,
     ProxyCredentials,
     RecordingConfiguration,
+    ResourceLocation,
+    SecretsManagerLocation,
     SessionConfiguration,
     ViewportConfiguration,
     VpcConfig,
@@ -25,19 +31,25 @@
     "BasicAuth",
     "BrowserClient",
     "browser_session",
+    "Certificate",
+    "CertificateLocation",
     "CodeInterpreter",
     "code_session",
     "BrowserConfiguration",
     "BrowserExtension",
     "BrowserSigningConfiguration",
     "CodeInterpreterConfiguration",
+    "EnterprisePolicy",
+    "EnterprisePolicyS3Location",
     "ExtensionS3Location",
     "ExternalProxy",
     "NetworkConfiguration",
     "ProfileConfiguration",
     "ProxyConfiguration",
     "ProxyCredentials",
     "RecordingConfiguration",
+    "ResourceLocation",
+    "SecretsManagerLocation",
     "SessionConfiguration",
     "ViewportConfiguration",
     "VpcConfig",

src/bedrock_agentcore/tools/browser_client.py

@@ -22,7 +22,14 @@
 from bedrock_agentcore._utils.user_agent import build_user_agent_suffix
 
 from .._utils.endpoints import get_control_plane_endpoint, get_data_plane_endpoint
-from .config import BrowserExtension, ProfileConfiguration, ProxyConfiguration, ViewportConfiguration
+from .config import (
+    BrowserExtension,
+    Certificate,
+    EnterprisePolicy,
+    ProfileConfiguration,
+    ProxyConfiguration,
+    ViewportConfiguration,
+)
 
 
 def _to_dict(value):
@@ -116,6 +123,8 @@ def create_browser(
         description: Optional[str] = None,
         recording: Optional[Dict] = None,
         browser_signing: Optional[Dict] = None,
+        enterprise_policies: Optional[List[Union[EnterprisePolicy, Dict[str, Any]]]] = None,
+        certificates: Optional[List[Union[Certificate, Dict[str, Any]]]] = None,
         tags: Optional[Dict[str, str]] = None,
         client_token: Optional[str] = None,
     ) -> Dict:
@@ -148,6 +157,11 @@ def create_browser(
                 {
                     "enabled": True
                 }
+            enterprise_policies (Optional[List[Union[EnterprisePolicy, Dict]]]): Chromium
+                enterprise policies at managed enforcement level. Up to 10 policy files,
+                each .json and max 5MB, from a same-region S3 bucket.
+            certificates (Optional[List[Union[Certificate, Dict]]]): Root CA certificates
+                from Secrets Manager for the browser to trust.
             tags (Optional[Dict[str, str]]): Tags for the browser
             client_token (Optional[str]): Idempotency token
 
@@ -194,6 +208,12 @@ def create_browser(
             request_params["browserSigning"] = browser_signing
             self.logger.info("🔐 Web Bot Auth (browserSigning) enabled")
 
+        if enterprise_policies:
+            request_params["enterprisePolicies"] = [_to_dict(p) for p in enterprise_policies]
+
+        if certificates:
+            request_params["certificates"] = [_to_dict(c) for c in certificates]
+
         if tags:
             request_params["tags"] = tags
 
@@ -299,6 +319,8 @@ def start(
         proxy_configuration: Optional[Union[ProxyConfiguration, Dict[str, Any]]] = None,
         extensions: Optional[List[Union[BrowserExtension, Dict[str, Any]]]] = None,
         profile_configuration: Optional[Union[ProfileConfiguration, Dict[str, Any]]] = None,
+        enterprise_policies: Optional[List[Union[EnterprisePolicy, Dict[str, Any]]]] = None,
+        certificates: Optional[List[Union[Certificate, Dict[str, Any]]]] = None,
     ) -> str:
         """Start a browser sandbox session.
 
@@ -324,6 +346,11 @@ def start(
                 configuration for persisting browser state across sessions. Can be a
                 ProfileConfiguration dataclass or a plain dict:
                 {"profileIdentifier": "my-profile-id"}
+            enterprise_policies (Optional[List[Union[EnterprisePolicy, Dict]]]): Chromium
+                enterprise policies at recommended enforcement level. Up to 10 policy files,
+                each .json and max 5MB, from a same-region S3 bucket.
+            certificates (Optional[List[Union[Certificate, Dict]]]): Root CA certificates
+                from Secrets Manager for the browser session to trust.
 
         Returns:
             str: The session ID of the newly created session.
@@ -373,6 +400,12 @@ def start(
         if profile_configuration is not None:
             request_params["profileConfiguration"] = _to_dict(profile_configuration)
 
+        if enterprise_policies is not None:
+            request_params["enterprisePolicies"] = [_to_dict(p) for p in enterprise_policies]
+
+        if certificates is not None:
+            request_params["certificates"] = [_to_dict(c) for c in certificates]
+
         response = self.data_plane_client.start_browser_session(**request_params)
 
         self.identifier = response["browserIdentifier"]
@@ -633,6 +666,8 @@ def browser_session(
     proxy_configuration: Optional[Union[ProxyConfiguration, Dict[str, Any]]] = None,
     extensions: Optional[List[Union[BrowserExtension, Dict[str, Any]]]] = None,
     profile_configuration: Optional[Union[ProfileConfiguration, Dict[str, Any]]] = None,
+    enterprise_policies: Optional[List[Union[EnterprisePolicy, Dict[str, Any]]]] = None,
+    certificates: Optional[List[Union[Certificate, Dict[str, Any]]]] = None,
 ) -> Generator[BrowserClient, None, None]:
     """Context manager for creating and managing a browser sandbox session.
 
@@ -648,6 +683,10 @@ def browser_session(
             extensions. Each element can be a BrowserExtension dataclass or a plain dict.
         profile_configuration (Optional[Union[ProfileConfiguration, Dict[str, Any]]]): Profile
             configuration. Can be a ProfileConfiguration dataclass or a plain dict.
+        enterprise_policies (Optional[List[Union[EnterprisePolicy, Dict[str, Any]]]]): Chromium
+            enterprise policies. Each element can be an EnterprisePolicy dataclass or a plain dict.
+        certificates (Optional[List[Union[Certificate, Dict[str, Any]]]]): Root CA certificates.
+            Each element can be a Certificate dataclass or a plain dict.
 
     Yields:
         BrowserClient: An initialized and started browser client.
@@ -687,6 +726,10 @@ def browser_session(
         start_kwargs["extensions"] = extensions
     if profile_configuration is not None:
         start_kwargs["profile_configuration"] = profile_configuration
+    if enterprise_policies is not None:
+        start_kwargs["enterprise_policies"] = enterprise_policies
+    if certificates is not None:
+        start_kwargs["certificates"] = certificates
 
     client.start(**start_kwargs)
 

src/bedrock_agentcore/tools/config.py

@@ -452,6 +452,124 @@ def to_dict(self) -> Dict:
         return config
 
 
+@dataclass
+class EnterprisePolicyS3Location:
+    """S3 location of a browser enterprise policy JSON file.
+
+    Attributes:
+        bucket: S3 bucket name (must be in the same region as the API call)
+        prefix: S3 object key for the policy JSON file
+        version_id: Optional S3 object version ID
+    """
+
+    bucket: str
+    prefix: str
+    version_id: Optional[str] = None
+
+    def to_dict(self) -> Dict:
+        """Convert to API-compatible dictionary."""
+        location = {"bucket": self.bucket, "prefix": self.prefix}
+        if self.version_id:
+            location["versionId"] = self.version_id
+        return location
+
+
+@dataclass
+class ResourceLocation:
+    """Location of a resource. Currently supports S3.
+
+    Attributes:
+        s3: S3 location of the resource
+    """
+
+    s3: Optional[EnterprisePolicyS3Location] = None
+
+    def to_dict(self) -> Dict:
+        """Convert to API-compatible dictionary."""
+        if self.s3:
+            return {"s3": self.s3.to_dict()}
+        raise ValueError("ResourceLocation must have one location type set")
+
+
+@dataclass
+class EnterprisePolicy:
+    """Browser enterprise policy.
+
+    Attributes:
+        location: Location of the enterprise policy file
+        type: "MANAGED" for CreateBrowser or "RECOMMENDED" for StartBrowserSession
+    """
+
+    location: ResourceLocation
+    type: str
+
+    def __post_init__(self):
+        """Validate enterprise policy type."""
+        if self.type not in ["MANAGED", "RECOMMENDED"]:
+            raise ValueError(f"type must be 'MANAGED' or 'RECOMMENDED', got '{self.type}'")
+
+    def to_dict(self) -> Dict:
+        """Convert to API-compatible dictionary."""
+        return {
+            "location": self.location.to_dict(),
+            "type": self.type,
+        }
+
+
+@dataclass
+class SecretsManagerLocation:
+    """Secrets Manager location for a certificate.
+
+    Attributes:
+        secret_arn: ARN of the Secrets Manager secret containing the certificate
+    """
+
+    secret_arn: str
+
+    def to_dict(self) -> Dict:
+        """Convert to API-compatible dictionary."""
+        return {"secretArn": self.secret_arn}
+
+
+@dataclass
+class CertificateLocation:
+    """Location from which to retrieve a certificate.
+
+    Attributes:
+        secrets_manager: Secrets Manager location containing the certificate
+    """
+
+    secrets_manager: SecretsManagerLocation
+
+    def to_dict(self) -> Dict:
+        """Convert to API-compatible dictionary."""
+        return {"secretsManager": self.secrets_manager.to_dict()}
+
+
+@dataclass
+class Certificate:
+    """Root CA certificate for browser or code interpreter.
+
+    Attributes:
+        location: Location of the certificate
+    """
+
+    location: CertificateLocation
+
+    def to_dict(self) -> Dict:
+        """Convert to API-compatible dictionary."""
+        return {"location": self.location.to_dict()}
+
+    @classmethod
+    def from_secret_arn(cls, secret_arn: str) -> "Certificate":
+        """Create a Certificate from a Secrets Manager ARN.
+
+        Args:
+            secret_arn: ARN of the secret containing the certificate
+        """
+        return cls(location=CertificateLocation(secrets_manager=SecretsManagerLocation(secret_arn=secret_arn)))
+
+
 def create_browser_config(
     name: str,
     execution_role_arn: str,