limit policy to only read one s3 file

Sage-Bionetworks · Jan 23, 2025 · 4edfa73 · 4edfa73
1 parent f960187
commit 4edfa73
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 13 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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());