feat: Add AWS Role Assumption

Author: moretalk

package-lock.json

@@ -61,6 +61,7 @@
     },
     "dependencies": {
         "@aws-crypto/sha256-js": "^4.0.0",
+        "@aws-sdk/client-sts": "^3.312.0",
         "@aws-sdk/signature-v4": "^3.310.0",
         "abort-controller": "^3.0.0",
         "axios": "^0.27.2",

package.json

@@ -9,6 +9,7 @@ export type OpenAPIConfig = {
     token?: string | Resolver<string>;
     accountSid?: string | Resolver<string>;
     authToken?: string | Resolver<string>;
+    awsRoleArn?: string | Resolver<string>;
 }
 
 export type OpenAPIConfig2 = {
@@ -18,4 +19,5 @@ export type OpenAPIConfig2 = {
     token?: string | Resolver<string>;
     accountSid?: string | Resolver<string>;
     authToken?: string | Resolver<string>;
+    awsRoleArn?: string | Resolver<string>;
 }

src/templates/core/OpenAPI.hbs

@@ -15,6 +15,7 @@ import { BaseHttpRequest } from './BaseHttpRequest';
 {{#if @root.awsSign}}
 import { SignatureV4 } from '@aws-sdk/signature-v4';
 import { Sha256 } from '@aws-crypto/sha256-js';
+import { AssumeRoleCommand, STSClient, AssumeRoleCommandOutput, Credentials } from "@aws-sdk/client-sts";
 
 {{>functions/signAWS}}
 {{/if}}
@@ -91,7 +92,37 @@ export class AxiosHttpRequest extends BaseHttpRequest {
                 const headers = await getHeaders(options, this.openApiConfig, url);
 
                 {{#if @root.awsSign}}
-                const awsSigned = await signAWS({ url, options, headers, body: body ?? formData });
+                let awsSigned = { headers: {} };
+                if(this.openApiConfig.awsRoleArn) {
+                    const awsCredentails: AssumeRoleCommandOutput = await assumeAWSRole({ awsRoleArn: this.openApiConfig.awsRoleArn }, options);
+
+                    if(!awsCredentails.Credentials) {
+                        throw new Error("Invalid AWS Role Arn");
+                    }
+                    
+                    awsSigned = await signAWS({ 
+                        credentials: awsCredentails.Credentials, 
+                        url, 
+                        options, 
+                        headers, 
+                        body: body ?? formData 
+                    });
+                }
+                else {
+                    awsSigned = await signAWS({ 
+                        credentials: {
+                            AccessKeyId: process.env.AWS_ACCESS_KEY_ID,
+                            SecretAccessKey: process.env.AWS_SECRET_ACCESS_KEY,
+                            SessionToken: process.env.AWS_SESSION_TOKEN,
+                            Expiration: new Date()
+                        }, 
+                        url, 
+                        options, 
+                        headers, 
+                        body: body ?? formData 
+                    });
+                }
+                    
                 {{/if}}
 
                 if (!onCancel.isCancelled) {

src/templates/core/axios/request.hbs

@@ -1,20 +1,28 @@
-async function signAWS({ url, options, headers, body }: { url: URL, options: ApiRequestOptions, headers: { [key: string]: any }, body?: any }): Promise<any> {
-    if(!process.env.AWS_REGION || !process.env.AWS_ACCESS_KEY_ID || !process.env.AWS_SECRET_ACCESS_KEY || !process.env.AWS_SESSION_TOKEN) {
+interface SignAWSPayload {
+    credentials: Credentials
+    url: URL, 
+    options: ApiRequestOptions, 
+    headers: { [key: string]: any }, 
+    body?: any 
+}
+
+async function signAWS({ credentials, url, options, headers, body }: SignAWSPayload): Promise<any> {
+    if(!process.env.AWS_REGION || !credentials.AccessKeyId || !credentials.SecretAccessKey || !credentials.SessionToken) {
         throw new Error("Missing AWS Credentials");
     }
 
     const sigv4 = new SignatureV4({
         service: 'execute-api',
         region: process.env.AWS_REGION,
         credentials: {
-            accessKeyId: process.env.AWS_ACCESS_KEY_ID,
-            secretAccessKey: process.env.AWS_SECRET_ACCESS_KEY,
-            sessionToken: process.env.AWS_SESSION_TOKEN,
+            accessKeyId: credentials.AccessKeyId,
+            secretAccessKey: credentials.SecretAccessKey,
+            sessionToken: credentials.SessionToken,
         },
         sha256: Sha256,
     });
 
-    const signed = await sigv4.sign({
+    return sigv4.sign({
         method: options.method,
         hostname: url.host,
         path: url.pathname,
@@ -23,6 +31,25 @@ async function signAWS({ url, options, headers, body }: { url: URL, options: Api
         headers: headers,
         body: body
   });
+}
+
+interface AssumeAWSRolePayload {
+    awsRoleArn: string | Resolver<string>;
+}
+
+async function assumeAWSRole({ awsRoleArn }: AssumeAWSRolePayload, options: ApiRequestOptions): Promise<AssumeRoleCommandOutput> {
+    const arn = await resolve(options, awsRoleArn);
+    
+    if(!process.env.AWS_REGION || !arn) {
+        throw new Error("Missing AWS Credentials");
+    }
+
+    const stsClient = new STSClient({ region: process.env.AWS_REGION });
+    const command = new AssumeRoleCommand({
+        RoleArn: arn,
+        RoleSessionName: "session1",
+        DurationSeconds: 900
+    });
 
-  return signed;
+    return stsClient.send(command);
 }

src/templates/core/functions/signAWS.hbs

@@ -25,7 +25,8 @@ export class {{{clientName}}} {{#if service}}extends {{{service.name}}} {{/if}}{
             version: '{{{version}}}',
             token: openApiConfig?.token,
             accountSid: openApiConfig?.accountSid,
-            authToken: openApiConfig?.authToken
+            authToken: openApiConfig?.authToken,
+            awsRoleArn: openApiConfig?.awsRoleArn
         });
         {{#if service}}
         super(request);