Dev

aws_manage_secrets.py

@@ -0,0 +1,67 @@
+import boto3
+import botocore
+import argparse
+
+def manage_secrets(region):
+    """Lists all secrets in AWS Secrets Manager and provides an option to delete them interactively."""
+    try:
+        # It's recommended to configure your AWS region.
+        # You can pass it to the client, e.g., boto3.client('secretsmanager', region_name='us-east-1')
+        # Or configure it via environment variables (AWS_REGION) or ~/.aws/config
+        client = boto3.client('secretsmanager', region_name=region)
+
+        region_name = region if region else client.meta.region_name
+        print(f"Fetching all secrets from AWS Secrets Manager in region: {region_name}...")
+
+        secrets_to_delete = []
+        paginator = client.get_paginator('list_secrets')
+        page_iterator = paginator.paginate()
+
+        for page in page_iterator:
+            for secret in page['SecretList']:
+                secrets_to_delete.append(secret['ARN'])
+                print(f"  - Found secret: {secret['Name']} (ARN: {secret['ARN']})")
+
+        if not secrets_to_delete:
+            print("No secrets found.")
+            return
+
+        print(f"\nFound a total of {len(secrets_to_delete)} secrets.")
+
+        # Ask for confirmation before proceeding with deletion
+        proceed = input("Do you want to proceed with deleting secrets? (yes/no): ").lower()
+        if proceed != 'yes':
+            print("Aborting. No secrets will be deleted.")
+            return
+
+        for secret_arn in secrets_to_delete:
+            try:
+                # Confirm deletion for each secret
+
+                print(f"Deleting secret: {secret_arn}")
+                # By default, secrets are not deleted immediately. 
+                # They are scheduled for deletion after a recovery window of 7 to 30 days.
+                # To force immediate deletion without recovery, use ForceDeleteWithoutRecovery=True
+                client.delete_secret(
+                    SecretId=secret_arn,
+                    RecoveryWindowInDays=7  # Set a recovery window (minimum 7 days)
+                    # ForceDeleteWithoutRecovery=True # Uncomment for immediate, irreversible deletion
+                )
+                print(f"Successfully scheduled secret {secret_arn} for deletion.")
+            except botocore.exceptions.ClientError as e:
+                print(f"Error deleting secret {secret_arn}: {e}")
+
+        print("\nSecret management process finished.")
+
+    except botocore.exceptions.NoCredentialsError:
+        print("AWS credentials not found. Please configure your credentials.")
+    except botocore.exceptions.ClientError as e:
+        print(f"An AWS client error occurred: {e}")
+    except Exception as e:
+        print(f"An unexpected error occurred: {e}")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='List and delete secrets from AWS Secrets Manager.')
+    parser.add_argument('--region', type=str, help='The AWS region to use. If not specified, the default region will be used.')
+    args = parser.parse_args()
+    manage_secrets(region=args.region)

src/agentstr/agents/agentstr.py

@@ -106,7 +106,24 @@ def checkpointer(self):
             return self._checkpointer
         checkpointer = None
         if self.database.conn_str.startswith("postgres"):
-            checkpointer = AsyncPostgresSaver.from_conn_string(self.database.conn_str)
+            key_manager = os.getenv("AGENT_VAULT_KEY_MANAGER")
+            key_manager_prefix = os.getenv("AGENT_VAULT_KEY_MANAGER_PREFIX")
+            if key_manager:
+                try:
+                    from agent_vault.langgraph import async_insecure_postgres_saver, async_secure_postgres_saver
+                    from agent_vault.utils.key_manager import AWSSecretsManagerKeyManager, AzureKeyVaultKeyManager
+                except ImportError:
+                    raise ValueError("agent_vault is not installed")
+                if key_manager == "none":
+                    checkpointer = async_insecure_postgres_saver(self.database.conn_str)
+                elif key_manager == "aws":
+                    checkpointer = async_secure_postgres_saver(self.database.conn_str, AWSSecretsManagerKeyManager(prefix=key_manager_prefix))
+                elif key_manager == "azure":
+                    checkpointer = async_secure_postgres_saver(self.database.conn_str, AzureKeyVaultKeyManager(prefix=key_manager_prefix))
+                else:
+                    raise ValueError(f"Unsupported key manager: {key_manager}")
+            else:
+                checkpointer = AsyncPostgresSaver.from_conn_string(self.database.conn_str)
         elif self.database.conn_str.startswith("sqlite"):
             conn_str = self.database.conn_str.replace("sqlite://", "", 1)
             checkpointer = AsyncSqliteSaver.from_conn_string(conn_str)

src/agentstr/cli/__init__.py

@@ -327,6 +327,28 @@ def _parse_kv(entries: tuple[str, ...], label: str, target: dict[str, str]):
             secret_ref = provider.put_secret(f'AGENTSTR-{deployment_name_safe}-{key.strip()}', val.strip())
             secrets_dict[key.strip()] = secret_ref
 
+    # Try to resolve agent vault
+    if config.get("agent_vault_key_manager"):
+        click.echo("Handling agent vault ...")
+        key_manager = config.get("agent_vault_key_manager")
+        if key_manager.lower() == 'aws':
+            val = 'aws'
+        elif key_manager.lower() == 'azure':
+            val = 'azure'
+        elif key_manager.lower() == 'none':
+            val = 'none'
+        else:
+            raise click.ClickException(f"Invalid agent_vault_key_manager: {key_manager}. Must be 'aws', 'azure', or 'none'.")
+        secret_ref = provider.put_secret(f'AGENTSTR-{deployment_name_safe}-AGENT_VAULT_KEY_MANAGER', val.strip())
+        secrets_dict['AGENT_VAULT_KEY_MANAGER'] = secret_ref
+        if val in {'aws', 'azure'}:
+            # Need to grant permission to the agent vault
+            secret_prefix = f'AGENTSTR-{deployment_name_safe}-AGENT_VAULT_KEYS-' 
+            secret_ref = provider.put_secret(f'AGENTSTR-{deployment_name_safe}-AGENT_VAULT_KEY_MANAGER_PREFIX', secret_prefix.strip())
+            secrets_dict['AGENT_VAULT_KEY_MANAGER_PREFIX'] = secret_ref          
+            # TODO: Grant permission to the agent vault
+            #provider.grant_secret_rw_access(secret_prefix)
+
     # 2. Load from config 'secrets', overwriting env_file
     config_secrets = cfg.get("secrets", {})
     if config_secrets: