From 4edfa7398cde6b38716d2e17a99f6c2e775801c5 Mon Sep 17 00:00:00 2001 From: john-hill Date: Thu, 23 Jan 2025 12:13:48 -0800 Subject: [PATCH] limit policy to only read one s3 file --- .../template/TemplateGuiceModule.java | 1 - .../agent/BedrockAgentContextProvider.java | 27 ++++++++++++------- .../repo/agent/bedrock_agent_template.json | 5 +++- .../RepositoryTemplateBuilderImplTest.java | 6 ++++- 4 files changed, 26 insertions(+), 13 deletions(-) diff --git a/src/main/java/org/sagebionetworks/template/TemplateGuiceModule.java b/src/main/java/org/sagebionetworks/template/TemplateGuiceModule.java index 0a74d30a..d081361a 100644 --- a/src/main/java/org/sagebionetworks/template/TemplateGuiceModule.java +++ b/src/main/java/org/sagebionetworks/template/TemplateGuiceModule.java @@ -11,7 +11,6 @@ import static org.sagebionetworks.template.TemplateUtils.loadFromJsonFile; import java.io.IOException; -import java.time.Duration; import org.apache.http.client.HttpClient; import org.apache.http.impl.client.HttpClientBuilder; diff --git a/src/main/java/org/sagebionetworks/template/repo/agent/BedrockAgentContextProvider.java b/src/main/java/org/sagebionetworks/template/repo/agent/BedrockAgentContextProvider.java index b0766fd7..e653e507 100644 --- a/src/main/java/org/sagebionetworks/template/repo/agent/BedrockAgentContextProvider.java +++ b/src/main/java/org/sagebionetworks/template/repo/agent/BedrockAgentContextProvider.java @@ -4,7 +4,6 @@ import static org.sagebionetworks.template.Constants.PROPERTY_KEY_STACK; import java.util.StringJoiner; -import java.util.UUID; import org.apache.velocity.VelocityContext; import org.json.JSONArray; @@ -35,24 +34,27 @@ public void addToContext(VelocityContext context) { String agentName = new StringJoiner("-").add(stack).add(instance).add("agent").toString(); String openApiSchemaBucket = String.format("%s-configuration.sagebase.org", stack); - String openApiSchemakey = String.format("chat/openapi/%s/%s.json", instance, UUID.randomUUID().toString()); + String openApiSchemakey = String.format("chat/openapi/%s.json", instance); String openApiSchemJsonString = TemplateUtils.loadContentFromFile("templates/repo/agent/agent_open_api.json"); s3Cient.putObject(openApiSchemaBucket, openApiSchemakey, openApiSchemJsonString); - + + String openApiSchemaS3Arn = String.format("arn:aws:s3:::%s/%s", openApiSchemaBucket, openApiSchemakey); JSONObject baseTemplate = new JSONObject(TemplateUtils.loadContentFromFile("templates/repo/agent/bedrock_agent_template.json")); JSONObject resources = baseTemplate.getJSONObject("Resources"); + JSONArray roleStatements = resources + .getJSONObject("bedrockAgentRole") + .getJSONObject("Properties") + .getJSONArray("Policies") + .getJSONObject(0) + .getJSONObject("PolicyDocument") + .getJSONArray("Statement"); + // Since the agent template is shared to external people, we need to hack it to replace parameters that do not exist in our template - JSONArray bedrockAgentRoleKbResource = resources - .getJSONObject("bedrockAgentRole") - .getJSONObject("Properties") - .getJSONArray("Policies") - .getJSONObject(0) - .getJSONObject("PolicyDocument") - .getJSONArray("Statement") + JSONArray bedrockAgentRoleKbResource = roleStatements .getJSONObject(1) .getJSONArray("Fn::If") .getJSONObject(1) @@ -71,6 +73,11 @@ public void addToContext(VelocityContext context) { kbProperty.getJSONObject("KnowledgeBaseId").put("Ref", "SynapseHelpKnowledgeBase"); kbProperty.put("Description", baseTemplate.getJSONObject("Parameters").getJSONObject("knowledgeBaseDescription").getString("Default")); + + // set bucket and key + JSONObject statementTwo = roleStatements.getJSONObject(2); + statementTwo.put("Resource", openApiSchemaS3Arn); + JSONObject s3 = bedrockAgentProps.getJSONArray("ActionGroups").getJSONObject(1).getJSONObject("ApiSchema") .getJSONObject("S3"); s3.put("S3BucketName", openApiSchemaBucket); diff --git a/src/main/resources/templates/repo/agent/bedrock_agent_template.json b/src/main/resources/templates/repo/agent/bedrock_agent_template.json index 420e69e8..2cbd6b09 100644 --- a/src/main/resources/templates/repo/agent/bedrock_agent_template.json +++ b/src/main/resources/templates/repo/agent/bedrock_agent_template.json @@ -112,7 +112,10 @@ }, { "Effect": "Allow", - "Action": "s3:*", + "Action": [ + "s3:ListObject", + "s3:GetObject" + ], "Resource": "*" } ] diff --git a/src/test/java/org/sagebionetworks/template/repo/RepositoryTemplateBuilderImplTest.java b/src/test/java/org/sagebionetworks/template/repo/RepositoryTemplateBuilderImplTest.java index 8a97bdad..b67c1817 100644 --- a/src/test/java/org/sagebionetworks/template/repo/RepositoryTemplateBuilderImplTest.java +++ b/src/test/java/org/sagebionetworks/template/repo/RepositoryTemplateBuilderImplTest.java @@ -394,6 +394,9 @@ public void testBuildAndDeployProd() throws InterruptedException { assertTrue(resources.has("SynapseHelpKnowledgeBase")); assertTrue(resources.has("bedrockAgentRole")); assertTrue(resources.has("bedrockAgent")); + + assertTrue(resources.getJSONObject("bedrockAgentRole").toString().contains("arn:aws:s3:::prod-configuration.sagebase.org/chat/openapi/101.json")); + JSONObject bedrockAgentProps = resources.getJSONObject("bedrockAgent").getJSONObject("Properties"); assertEquals("prod-101-agent", bedrockAgentProps.get("AgentName")); @@ -403,12 +406,13 @@ public void testBuildAndDeployProd() throws InterruptedException { } void validateOpenApiSchema(JSONObject bedrockAgentProps) { + JSONObject s3 = bedrockAgentProps.getJSONArray("ActionGroups").getJSONObject(1).getJSONObject("ApiSchema") .getJSONObject("S3"); String openApiBucket = s3.getString("S3BucketName"); assertEquals("prod-configuration.sagebase.org", openApiBucket); String openApiKey = s3.getString("S3ObjectKey"); - assertTrue(openApiKey.startsWith("chat/openapi/101/")); + assertEquals("chat/openapi/101.json",s3.getString("S3ObjectKey")); verify(mockS3Client).putObject(eq(openApiBucket), eq(openApiKey), jsonStringCaptor.capture()); JSONObject openApiSchema = new JSONObject(jsonStringCaptor.getValue());