From 8aa25a5cb6c49fb073aab5190da3a13b10101e3f Mon Sep 17 00:00:00 2001
From: peterwoodworth <woodwoop@amazon.com>
Date: Wed, 21 Jun 2023 14:50:04 -0700
Subject: [PATCH] feat: optional policy inputs when assuming role

---
 action.yml        | 6 ++++++
 src/assumeRole.ts | 6 ++++++
 src/index.ts      | 8 ++++++++
 3 files changed, 20 insertions(+)

diff --git a/action.yml b/action.yml
index 8b99025..4555942 100644
--- a/action.yml
+++ b/action.yml
@@ -52,6 +52,12 @@ inputs:
   web-identity-token-file:
     description: Use the web identity token file from the provided file system path in order to assume an IAM role using a web identity, e.g. from within an Amazon EKS worker node.
     required: false
+  inline-session-policy:
+    description: 'Inline session policy'
+    required: false
+  managed-session-policies:
+    description: 'List of managed session policies'
+    required: false
 outputs:
   aws-account-id:
     description: The AWS account ID for the provided credentials
diff --git a/src/assumeRole.ts b/src/assumeRole.ts
index 279fbf6..ea44764 100644
--- a/src/assumeRole.ts
+++ b/src/assumeRole.ts
@@ -71,6 +71,8 @@ export interface assumeRoleParams {
   roleExternalId?: string;
   webIdentityTokenFile?: string;
   webIdentityToken?: string;
+  inlineSessionPolicy?: string;
+  managedSessionPolicies?: any[];
 }
 
 export async function assumeRole(params: assumeRoleParams) {
@@ -84,6 +86,8 @@ export async function assumeRole(params: assumeRoleParams) {
     roleSkipSessionTagging,
     webIdentityTokenFile,
     webIdentityToken,
+    inlineSessionPolicy,
+    managedSessionPolicies
   } = { ...params };
 
   // Load GitHub environment variables
@@ -128,6 +132,8 @@ export async function assumeRole(params: assumeRoleParams) {
     DurationSeconds: roleDuration,
     Tags: tags ? tags : undefined,
     ExternalId: roleExternalId ? roleExternalId : undefined,
+    Policy: inlineSessionPolicy ? inlineSessionPolicy : undefined,
+    PolicyArns: managedSessionPolicies ? managedSessionPolicies : undefined,
   };
   const keys = Object.keys(commonAssumeRoleParams) as Array<keyof typeof commonAssumeRoleParams>;
   keys.forEach((k) => commonAssumeRoleParams[k] === undefined && delete commonAssumeRoleParams[k]);
diff --git a/src/index.ts b/src/index.ts
index 77f82b8..f639a39 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -29,6 +29,12 @@ export async function run() {
     const roleSkipSessionTagging = roleSkipSessionTaggingInput.toLowerCase() === 'true';
     const proxyServer = core.getInput('http-proxy', { required: false });
     const disableOIDC = core.getInput('disable-oidc', { required: false });
+    const inlineSessionPolicy = core.getInput('inline-session-policy', { required: false });
+    const managedSessionPoliciesInput = core.getMultilineInput('managed-session-policies', { required: false })
+    const managedSessionPolicies: any[] = [];
+    for (const managedSessionPolicy of managedSessionPoliciesInput) {
+      managedSessionPolicies.push({arn: managedSessionPolicy})
+    }
 
     // Logic to decide whether to attempt to use OIDC or not
     const useGitHubOIDCProvider = () => {
@@ -110,6 +116,8 @@ export async function run() {
           roleSkipSessionTagging,
           webIdentityTokenFile,
           webIdentityToken,
+          inlineSessionPolicy,
+          managedSessionPolicies,
         });
       }, true);
       core.info(`Authenticated as assumedRoleId ${roleCredentials.AssumedRoleUser!.AssumedRoleId!}`);