From d4485fa78b8511c434727555223d2b62d7847066 Mon Sep 17 00:00:00 2001
From: Cat Zimmermann <cat.zimmermann@lacework.net>
Date: Mon, 29 Jul 2024 21:29:02 +0000
Subject: [PATCH] Add MFA token support

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

diff --git a/action.yml b/action.yml
index 7e98591a9..2132bc411 100644
--- a/action.yml
+++ b/action.yml
@@ -24,6 +24,12 @@ inputs:
   aws-session-token:
     description: AWS Session Token.
     required: false
+  mfa-token:
+    description: Token when the user/role requires MFA.
+    required: false
+  mfa-serial:
+    description: Serial/ARN when the user/role requires MFA.
+    required: false
   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
diff --git a/src/assumeRole.ts b/src/assumeRole.ts
index 7806dec0f..f5f1bd008 100644
--- a/src/assumeRole.ts
+++ b/src/assumeRole.ts
@@ -74,6 +74,8 @@ export interface assumeRoleParams {
   roleExternalId?: string;
   webIdentityTokenFile?: string;
   webIdentityToken?: string;
+  mfaSerial?: string;
+  mfaToken?: string;
   inlineSessionPolicy?: string;
   managedSessionPolicies?: any[];
 }
@@ -89,6 +91,8 @@ export async function assumeRole(params: assumeRoleParams) {
     roleSkipSessionTagging,
     webIdentityTokenFile,
     webIdentityToken,
+    mfaSerial,
+    mfaToken,
     inlineSessionPolicy,
     managedSessionPolicies,
   } = { ...params };
@@ -137,6 +141,8 @@ export async function assumeRole(params: assumeRoleParams) {
     ExternalId: roleExternalId ? roleExternalId : undefined,
     Policy: inlineSessionPolicy ? inlineSessionPolicy : undefined,
     PolicyArns: managedSessionPolicies?.length ? managedSessionPolicies : undefined,
+    SerialNumber: mfaSerial,
+    TokenCode: mfaToken,
   };
   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 da296c70c..1316da91f 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -30,6 +30,8 @@ export async function run() {
     const maskAccountId = maskAccountIdInput.toLowerCase() === 'true';
     const roleExternalId = core.getInput('role-external-id', { required: false });
     const webIdentityTokenFile = core.getInput('web-identity-token-file', { required: false });
+    const mfaSerial = core.getInput('mfa-serial', { required: false });
+    const mfaToken = core.getInput('mfa-token', { required: false });
     const roleDuration = parseInt(core.getInput('role-duration-seconds', { required: false })) || DEFAULT_ROLE_DURATION;
     const roleSessionName = core.getInput('role-session-name', { required: false }) || ROLE_SESSION_NAME;
     const roleSkipSessionTaggingInput = core.getInput('role-skip-session-tagging', { required: false }) || 'false';
@@ -159,6 +161,8 @@ export async function run() {
               roleSkipSessionTagging,
               webIdentityTokenFile,
               webIdentityToken,
+              mfaSerial,
+              mfaToken,
               inlineSessionPolicy,
               managedSessionPolicies,
             });