Enable use case with RBAC to improve security

Makefile

@@ -1,6 +1,25 @@
-install:
+BUILD_DIR = ./dist
+SRC_DIR = ./src
+NODE_MODULES = ./node_modules
+SRC_FILES := $(shell find $(SRC_DIR) -type f)
+
+.PHONY: all install build run
+
+all: build
+
+$(NODE_MODULES):
 	pnpm install
 
+install: $(NODE_MODULES)
+
+$(BUILD_DIR): $(SRC_FILES) $(NODE_MODULES)
+	pnpm build
+
+build: $(BUILD_DIR)
+
+run: build
+	node ./bin/index.js --profile "${AWS_PROFILE}" --access-key "${ACCESS_KEY}" --secret-key "${SECRET_KEY}" --role-arn "${ROLE_ARN}" --region "${AWS_REGION}" -S $(SLACK_TOKEN_AWS_COST_CLI) -C $(SLACK_CHANNEL_ID)
+
 test:
 	pnpm test
 

jest.config.cjs

@@ -9,6 +9,7 @@ var esmModules = [
   'ansi-regex',
   'is-interactive',
   'stdin-discarder',
+  'aws-sdk-client-mock',
 ];
 
 module.exports = {

package.json

@@ -52,6 +52,7 @@
   },
   "homepage": "https://github.com/kamranahmedse/aws-cost-cli#readme",
   "dependencies": {
+    "@aws-sdk/client-sts": "^3.521.0",
     "@aws-sdk/shared-ini-file-loader": "^3.254.0",
     "aws-sdk": "^2.1299.0",
     "chalk": "^5.2.0",
@@ -66,11 +67,14 @@
     "@babel/preset-env": "^7.24.6",
     "@babel/preset-typescript": "^7.24.6",
     "@eslint/js": "^9.4.0",
+    "@smithy/types": "^3.0.0",
     "@types/eslint__js": "^8.42.3",
     "@types/jest": "^29.5.12",
     "@types/node": "^18.11.18",
     "@typescript-eslint/eslint-plugin": "^7.11.0",
     "@typescript-eslint/parser": "^7.11.0",
+    "aws-sdk-client-mock": "^2.2.0",
+    "aws-sdk-client-mock-jest": "^4.0.0",
     "babel-jest": "^29.7.0",
     "eslint": "^9.4.0",
     "jest": "^29.7.0",

readme.md

@@ -35,6 +35,7 @@ $ aws-cost --help
     -k, --access-key [key]         AWS access key
     -s, --secret-key [key]         AWS secret key
     -r, --region [region]          AWS region (default: us-east-1)
+    -a, --role-arn [arn]           AWS role ARN to assume
 
     -p, --profile [profile]        AWS profile to use (default: "default")
 
@@ -63,6 +64,12 @@ aws-cost
 
 To configure the credentials using aws-cli, have a look at the [aws-cli docs](https://github.com/aws/aws-cli#configuration) for more information.
 
+If you need to assume a role, you can pass the `role-arn` option:
+
+```bash
+aws-cost -a arn:aws:iam::123456789012:role/your-role-arn
+```
+
 ## Detailed Breakdown
 
 > The default usage is to get the cost breakdown by service
@@ -197,7 +204,53 @@ Regarding the credentials, you need to have the following permissions in order t
 }
 ```
 
-Also, please note that this tool uses AWS Cost Explorer under the hood which [costs $0.01 per request](https://aws.amazon.com/aws-cost-management/aws-cost-explorer/pricing/).
+If you need to use Role-Based Access Control (RBAC), you will need to configure two IAM roles: the provider role and the consumer role.
+
+1. **Provider Role**
+
+   This role provides the necessary permissions for `aws-cost-cli`. It requires the above permissions policy and the following trust policy.
+
+   **Trust Policy**
+
+   Replace `arn:aws:iam::YOUR_ACCOUNT_ID:role/YourConsumerRole` with the ARN of the consumer role.
+
+   ```json
+   {
+     "Version": "2012-10-17",
+     "Statement": [
+       {
+         "Effect": "Allow",
+         "Principal": {
+           "AWS": "arn:aws:iam::YOUR_ACCOUNT_ID:role/YourConsumerRole"
+         },
+         "Action": "sts:AssumeRole"
+       }
+     ]
+   }
+   ```
+
+2. **Consumer Role**
+
+   This role is used by the user or service (such as GitHub Actions) that needs to assume the provider role to access cost information. It requires the following permissions policy.
+
+   **Permissions Policy**
+
+   Replace `arn:aws:iam::YOUR_ACCOUNT_ID:role/YourProviderRole` with the ARN of the provider role.
+
+   ```json
+   {
+     "Version": "2012-10-17",
+     "Statement": [
+       {
+         "Effect": "Allow",
+         "Action": "sts:AssumeRole",
+         "Resource": "arn:aws:iam::YOUR_ACCOUNT_ID:role/YourProviderRole"
+       }
+     ]
+   }
+   ```
+
+Please also note that this tool uses AWS Cost Explorer under the hood which [costs $0.01 per request](https://aws.amazon.com/aws-cost-management/aws-cost-explorer/pricing/).
 
 ## License
 

src/config.test.ts

@@ -0,0 +1,78 @@
+import {
+  STSClient,
+  AssumeRoleCommand,
+  AssumeRoleCommandOutput,
+  Credentials,
+} from '@aws-sdk/client-sts';
+import 'aws-sdk-client-mock-jest';
+import { AwsClientStub, mockClient } from 'aws-sdk-client-mock';
+import { getAwsConfigFromOptionsOrFile } from './config';
+
+let stsMock: AwsClientStub<STSClient>;
+
+beforeEach(() => {
+  stsMock = mockClient(STSClient) as AwsClientStub<STSClient>;
+});
+
+afterEach(() => {
+  stsMock.restore();
+});
+
+describe('should assume role if roleArn is provided', (): void => {
+  const mockCredentials: Credentials = {
+    AccessKeyId: 'mockAccessKeyId',
+    SecretAccessKey: 'mockSecretAccessKey',
+    SessionToken: 'mockSessionToken',
+    Expiration: new Date(),
+  };
+
+  it('should assume role if `roleArn` is provided', async () => {
+    const roleArn = 'arn:aws:iam::123456789012:role/test-role';
+    stsMock.on(AssumeRoleCommand).resolves({
+      Credentials: mockCredentials,
+      $metadata: {},
+    } as AssumeRoleCommandOutput);
+
+    const awsConfig = await getAwsConfigFromOptionsOrFile({
+      profile: 'default',
+      accessKey: '',
+      secretKey: '',
+      sessionToken: '',
+      region: 'us-east-1',
+      roleArn,
+    });
+
+    expect(stsMock).toHaveReceivedCommandWith(AssumeRoleCommand, {
+      RoleArn: roleArn,
+    });
+
+    expect(awsConfig.credentials).toEqual({
+      accessKeyId: mockCredentials.AccessKeyId,
+      secretAccessKey: mockCredentials.SecretAccessKey,
+      sessionToken: mockCredentials.SessionToken,
+    });
+  });
+
+  it('should allow ABAC if `{accessKey, secretKey, sessionToken}` provided', async () => {
+    const accessKey = 'testAccessKey';
+    const secretKey = 'testSecretKey';
+    const sessionToken = 'testSessionToken';
+
+    const awsConfig = await getAwsConfigFromOptionsOrFile({
+      profile: 'default',
+      accessKey: accessKey,
+      secretKey: secretKey,
+      sessionToken: sessionToken,
+      region: 'us-east-1',
+      roleArn: '',
+    });
+
+    expect(stsMock).toHaveReceivedCommandTimes(AssumeRoleCommand, 0);
+
+    expect(awsConfig.credentials).toEqual({
+      accessKeyId: accessKey,
+      secretAccessKey: secretKey,
+      sessionToken: sessionToken,
+    });
+  });
+});

src/config.ts

@@ -1,4 +1,5 @@
 import { loadSharedConfigFiles } from '@aws-sdk/shared-ini-file-loader';
+import { STSClient, AssumeRoleCommand } from '@aws-sdk/client-sts';
 import chalk from 'chalk';
 import { printFatalError } from './logger';
 
@@ -23,8 +24,10 @@ export async function getAwsConfigFromOptionsOrFile(options: {
   secretKey;
   sessionToken;
   region: string;
+  roleArn?: string;
 }): Promise<AWSConfig> {
-  const { profile, accessKey, secretKey, sessionToken, region } = options;
+  const { profile, accessKey, secretKey, sessionToken, region, roleArn } =
+    options;
 
   if (accessKey || secretKey) {
     if (!accessKey || !secretKey) {
@@ -46,7 +49,7 @@ export async function getAwsConfigFromOptionsOrFile(options: {
   }
 
   return {
-    credentials: await loadAwsCredentials(profile),
+    credentials: await loadAwsCredentials(profile, region, roleArn),
     region: region,
   };
 }
@@ -57,6 +60,8 @@ export async function getAwsConfigFromOptionsOrFile(options: {
  */
 async function loadAwsCredentials(
   profile: string = 'default',
+  region: string,
+  roleArn = '',
 ): Promise<AWSConfig['credentials'] | undefined> {
   const configFiles = await loadSharedConfigFiles();
 
@@ -71,34 +76,54 @@ async function loadAwsCredentials(
   // https://github.com/kamranahmedse/aws-cost-cli/issues/1
   // const configFile = configFiles.configFile;
   // const region: string = configFile?.[profile]?.region;
+  if (accessKey && secretKey) {
+    return {
+      accessKeyId: accessKey,
+      secretAccessKey: secretKey,
+      sessionToken: sessionToken,
+    };
+  } else {
+    try {
+      const stsClient = new STSClient({ region: region });
+      const assumeRoleCommand = new AssumeRoleCommand({
+        RoleArn: roleArn,
+        RoleSessionName: 'aws-cost-cli',
+      });
+
+      const { Credentials } = await stsClient.send(assumeRoleCommand);
+
+      if (Credentials) {
+        return {
+          accessKeyId: Credentials.AccessKeyId,
+          secretAccessKey: Credentials.SecretAccessKey,
+          sessionToken: Credentials.SessionToken,
+        };
+      }
+    } catch (error) {
+      console.error('Error fetching temporary credentials:', error);
+    }
 
-  if (!accessKey || !secretKey) {
     const sharedCredentialsFile =
       process.env.AWS_SHARED_CREDENTIALS_FILE || '~/.aws/credentials';
     const sharedConfigFile = process.env.AWS_CONFIG_FILE || '~/.aws/config';
 
     printFatalError(`
     Could not find the AWS credentials in the following files for the profile "${profile}":
-      ${chalk.bold(sharedCredentialsFile)}
-      ${chalk.bold(sharedConfigFile)}
+    ${chalk.bold(sharedCredentialsFile)}
+    ${chalk.bold(sharedConfigFile)}
 
     If the config files exist at different locations, set the following environment variables:
-      ${chalk.bold(`AWS_SHARED_CREDENTIALS_FILE`)}
-      ${chalk.bold(`AWS_CONFIG_FILE`)}
+    ${chalk.bold(`AWS_SHARED_CREDENTIALS_FILE`)}
+    ${chalk.bold(`AWS_CONFIG_FILE`)}
 
     You can also configure the credentials via the following command:
-      ${chalk.bold(`aws configure --profile ${profile}`)}
+    ${chalk.bold(`aws configure --profile ${profile}`)}
 
     You can also provide the credentials via the following options:
-      ${chalk.bold(`--access-key`)}
-      ${chalk.bold(`--secret-key`)}
-      ${chalk.bold(`--region`)}
+    ${chalk.bold(`--access-key`)}
+    ${chalk.bold(`--secret-key`)}
+    ${chalk.bold(`--region`)}
+    ${chalk.bold(`--role-arn`)}
     `);
   }
-
-  return {
-    accessKeyId: accessKey,
-    secretAccessKey: secretKey,
-    sessionToken: sessionToken,
-  };
 }

src/index.ts

@@ -24,6 +24,7 @@ program
   .option('-s, --secret-key [key]', 'AWS secret key')
   .option('-t, --session-Token [key]', 'AWS session Token')
   .option('-r, --region [region]', 'AWS region', 'us-east-1')
+  .option('--role-arn [arn]', 'ARN of IAM role')
   // Output variants
   .option('-j, --json', 'Get the output as JSON')
   .option('-u, --summary', 'Get only the summary without service breakdown')
@@ -45,6 +46,7 @@ type OptionsType = {
   secretKey: string;
   sessionToken: string;
   region: string;
+  roleArn: string;
   // AWS profile to use
   profile: string;
   // Output variants
@@ -71,6 +73,7 @@ const awsConfig = await getAwsConfigFromOptionsOrFile({
   secretKey: options.secretKey,
   sessionToken: options.sessionToken,
   region: options.region,
+  roleArn: options.roleArn,
 });
 
 const alias = await getAccountAlias(awsConfig);