AWS Machine Learning Blog

Use Amazon Bedrock Agents for code scanning, optimization, and remediation

Amazon Bedrock is a fully managed service that makes foundation models (FMs) from leading AI startups and Amazon available through an API, so you can choose from a wide range of FMs to find the model that best suits your use case. With the Amazon Bedrock serverless experience, you can get started quickly, privately customize FMs with your own data, integrate and deploy them into your application using Amazon Web Services (AWS) tools without having to manage any infrastructure.

For enterprises in the realm of cloud computing and software development, providing secure code repositories is essential. As sophisticated cybersecurity threats become more prevalent, organizations must adopt proactive measures to protect their assets. Amazon Bedrock offers a powerful solution by automating the process of scanning repositories for vulnerabilities and remediating them. This post explores how you can use Amazon Bedrock to enhance the security of your repositories and maintain compliance with organizational and regulatory standards.

This solution demonstrates how Amazon Bedrock Agents can be configured to scan a specific code repository, remediate vulnerabilities, and push the changes to a new branch. This approach can accelerate development, reduce errors, and adhere to security guidelines.

Solution overview

There are three high-level steps to deploy the solution:

  1. Configure the Amazon Bedrock Agent
  2. Configure the AWS Lambda function for the action group
  3. Add the action group to the Amazon Bedrock agent

There are two key steps in the architecture, as illustrated in the following diagram:

  1. The user provides the necessary information through the Amazon Bedrock agent chat console. They supply the code repository URL, such as https://github.com/abc/test, and specify the branch name to scan, for instance, main. Then they list the folders to exclude from the scan, such as test, and specify file extensions to exclude, such as .md and .txt. Then they provide a new branch name where the remediated code will be uploaded.
  2. The Amazon Bedrock agent forwards the details to an action group that invokes a Lambda function. This function retrieves the code, scans it for vulnerabilities using a preselected large language model (LLM), applies remediation, and pushes the remediated code to a new branch for user validation. The excluded folders and file extensions aren’t scanned. Upon completion, the action group (Lambda function) sends the information back to the Amazon Bedrock agent, which then displays the status to the user.

Figure 1. Architecture Diagram

Prerequisites

To implement the solution, you need the following:

Configure the Amazon Bedrock agent

To configure the Amazon Bedrock agent, complete the following steps:

  1. On the Amazon Bedrock console, choose Agents in the navigation pane, then choose Create Agent.
  2. (Optional) Provide agent details, including agent name and description.
  3. Grant the agent permissions to AWS services through the IAM service role. This gives your agent access to required services, such as Lambda.
  4. Select an FM in Amazon Bedrock (such as Anthropic’s Claude 3 Sonnet).
  5. To scan a code repository and remediate vulnerabilities through Amazon Bedrock Agents, attach the following instruction to the agent:

You are a code scanning and remediating AI assistant. Greet the user and ask user for repository_url and branch_name that needs to be scanned. Ask user for list of folders that needs to be excluded from scanning and also ask user for list of specific file extensions that needs to be excluded from scanning. Ask user new branch name to push the remediated code. Pass those inputs to trigger code-scan-remediation action group.

Configure the Lambda for the action group

After initial agent configuration and adding the preceding instruction to the agent, you create one Lambda function that will be used for the action group.

Create a Lambda function designed to scan a code repository for vulnerabilities, remediate the vulnerabilities, and push the changes to a new user-specified branch. This function will be used by the action group, which will be invoked by the Amazon Bedrock agent following the user’s input of the code repository URL, branch name, and the list of folders and file extensions to exclude from the scan. Reference to the Lambda code. Confirm that the Lambda function has the required IAM permissions and set up a Resource-based policy on the Lambda function to allow Amazon Bedrock Agent to invoke the Lambda using the lambda:InvokeFunction action. Refer to the policy here.

Add the action group to the Amazon Bedrock agent

Complete the following steps to add the action groups to the Amazon Bedrock agent:

  • Add an action group to the Amazon Bedrock agent.
  • Assign a descriptive name to the action group and detail the function in the description field. This helps clarify the purpose of the action group within the workflow.
  • For Action group type, select Define with function details.
  • For Action group invocation, select the Lambda function that you have created previously.

This function runs the business logic required when an action is invoked. Make sure to choose the correct version of the Lambda function and that the GitHub token is set as an environment variable. For more on how to configure Lambda functions for action groups, refer to Configure Lambda functions to send information an Amazon Bedrock agent elicits from the user.

  • For the Action group function 1, select JSON Editor and add the required parameters. Reference to the JSON file.

The following screenshot shows an example of the user interaction with Amazon Bedrock Agents.

Amazon Bedrock Agent sample interaction

Figure 2. User Interaction with Amazon Bedrock Agent

The following screenshot shows an example of remediated code.

Example output

Figure 3. Sample difference of Actual and Remediated Code 

Best practices

Follow these best practices:

  • Add automation tests to validate the code before committing it to the repository and review the remediated code before merging it into the default branch
  • Use descriptive branch names when creating new branches during remediation to maintain clear version control
  • Configure IAM roles and permissions with the principle of least privilege to secure the Amazon Bedrock agent and Lambda functions
  • Update prompts to target and remediate use-case specific vulnerabilities

Clean up

The services used in this demo can incur costs. Complete the following steps to clean up your resources:

  1. Delete the Lambda function if it’s no longer required
  2. Delete the action group and agents you created
  3. Remove the generated branch from the GitHub repository

Conclusion

Amazon Bedrock Agents uses generative AI to transform code repositories by scanning for vulnerabilities and automatically applying fixes. This capability is essential for engineers because it speeds up the process of securing code and maintaining compliance with established best practices from the outset.

The interactive features of Amazon Bedrock Agents automate the vulnerability scanning and remediation process, not only streamlining the initial setup but also significantly enhancing ongoing code maintenance. Although this post focuses on code scanning and remediation, the interactive capabilities of Amazon Bedrock Agents can be applied across various AWS services, offering a dynamic and comprehensive solution for managing and optimizing cloud infrastructure.

Are you ready to streamline your cloud deployment process with the generative AI of Amazon Bedrock? Start by exploring the Amazon Bedrock User Guide to learn how it can facilitate your organization’s transition to the cloud. For specialized assistance, consider engaging with AWS Professional Services to maximize the efficiency and benefits of using Amazon Bedrock.

Embrace the potential for a swift, secure, and efficient cloud transformation with Amazon Bedrock. Take the first step today and discover how using generative AI can revolutionize your approach to cloud infrastructure.

About the authors

Rama Krishna Yalla is an Associate DevOps Consultant at AWS, adept at designing scalable, reliable, and secure cloud environments. He leverages automation and CI/CD best practices to streamline software delivery, reduce downtime, and enhance operational efficiency. Rama is experienced in managing infrastructure as code (IaC) ensuring consistent and repeatable deployments. He also focuses on implementing robust monitoring and logging solutions, enabling proactive issue resolution and optimized performance. Outside of work, Rama enjoys playing badminton and often participates in local tournaments.

Akhil Raj Yallamelli is a Cloud Infrastructure Architect at AWS, specializing in architecting cloud infrastructure solutions for enhanced data security and cost efficiency. He is experienced in integrating technical solutions with business strategies to create scalable, reliable, and secure cloud environments. Akhil enjoys developing solutions focusing on customer business outcomes, incorporating generative AI (Gen AI) technologies to drive innovation and cloud enablement. He holds an MS degree in Computer Science. Outside of his professional work, Akhil enjoys watching and playing sports.