fix(upgrade): disable EventBridge rule NewLogGroup_rule during asea prep

Author: crissupb

reference-artifacts/Custom-Scripts/lza-upgrade/.prettierrc

@@ -0,0 +1,7 @@
+{
+  "tabWidth": 2,
+  "printWidth": 120,
+  "singleQuote": true,
+  "trailingComma": "all",
+  "arrowParens": "avoid"
+}

reference-artifacts/Custom-Scripts/lza-upgrade/src/common/utils/accounts.ts

@@ -13,6 +13,14 @@
 
 import { DynamoDB } from '../aws/dynamodb';
 import { Account } from '../outputs/accounts';
+import { STS } from '../aws/sts';
+
+
+export interface Environment {
+    accountId: string;
+    accountKey: string;
+    region: string;
+}
 
 export async function loadAccounts(tableName: string, client: DynamoDB): Promise<Account[]> {
   let index = 0;
@@ -31,3 +39,22 @@ export async function loadAccounts(tableName: string, client: DynamoDB): Promise
   }
   return accounts;
 }
+
+export function getEnvironments(accounts: Account[], regions: string[]): Environment[] {
+  const environments: Environment[] = [];
+  for (const account of accounts) {
+    for (const region of regions) {
+      environments.push({
+        accountId: account.id,
+        accountKey: account.key,
+        region,
+      });
+    }
+  }
+  return environments;
+}
+
+export async function assumeRole(accountId: string, roleName: string) {
+  const sts = new STS();
+  return sts.getCredentialsForAccountAndRole(accountId, roleName);
+}

reference-artifacts/Custom-Scripts/lza-upgrade/src/disable-rules.ts

@@ -0,0 +1,105 @@
+/**
+ *  Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ *  with the License. A copy of the License is located at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  or in the 'license' file accompanying this file. This file is distributed on an 'AS IS' BASIS, WITHOUT WARRANTIES
+ *  OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions
+ *  and limitations under the License.
+ */
+
+import { Account } from './common/outputs/accounts';
+import { DynamoDB } from './common/aws/dynamodb';
+import { Environment, getEnvironments, loadAccounts } from './common/utils/accounts';
+import { loadAseaConfig } from './asea-config/load';
+import { Config } from './config';
+import { AcceleratorConfig } from './asea-config';
+import { STS } from './common/aws/sts';
+import { throttlingBackOff } from './common/aws/backoff';
+import { EventBridgeClient, DisableRuleCommand, ResourceNotFoundException } from '@aws-sdk/client-eventbridge';
+
+export class DisableRules {
+    homeRegion: string;
+    assumeRoleName: string;
+    configRepositoryName: string;
+    accountList: Account[];
+    enabledRegions: string[];
+    sts: STS;
+    constructor(config: Config, 
+        disableRuleConfig: { accountList: Account[], enabledRegions: string[], acceleratorConfig: AcceleratorConfig },
+    ) {
+      this.homeRegion = config.homeRegion;
+      this.sts = new STS();
+      this.assumeRoleName = config.assumeRoleName ?? 'OrganizationAccountAccessRole';
+      this.configRepositoryName = config.repositoryName;
+      this.accountList = disableRuleConfig.accountList;
+      this.enabledRegions = disableRuleConfig.enabledRegions;
+    }
+
+    static async init(config: Config) {
+        const accountList = await loadAccounts(config.parametersTableName, new DynamoDB(undefined, config.homeRegion));
+        const acceleratorConfig = await loadAseaConfig({
+            filePath: 'raw/config.json',
+            repositoryName: config.repositoryName,
+            defaultRegion: config.homeRegion,
+        });
+        const enabledRegions = acceleratorConfig['global-options']['supported-regions'];
+        const disableRules = new DisableRules(config, {accountList, enabledRegions, acceleratorConfig});
+        return disableRules;
+    }
+    async disableAllAccountRules(prefix: string) {
+      const environments = getEnvironments(this.accountList, this.enabledRegions);
+      const eventBridgeClientMap = await this.getEventBridgeClientMap(environments);
+      const deleteRulesPromises = environments.map(environment => this.disableEventBridgeRules(eventBridgeClientMap, environment.accountId, environment.region, prefix.replaceAll('-', '')));
+      await Promise.all(deleteRulesPromises);
+    }
+  
+    async disableEventBridgeRules(eventBridgeClientMap: Map<string, EventBridgeClient>, accountId: string, region: string, prefix: string,){
+        const client = eventBridgeClientMap.get(`${accountId}-${region}`);
+        if (!client) {
+            throw new Error(`No client found for account ${accountId} and region ${region}`);
+        };
+        const suffixes = [
+            'NewLogGroup_rule'
+        ];
+        for (const suffix of suffixes) {
+            try {
+                const command = new DisableRuleCommand({
+                Name: `${prefix}-${suffix}`,
+                });
+                await throttlingBackOff (() => client.send(command));
+            } catch (e: any) {
+                if (e instanceof ResourceNotFoundException) {
+                    continue;
+                }
+            }
+            console.log(`Disabling rule ${prefix}-${suffix} in ${accountId}-${region}`);
+        }
+    }
+
+  async getEventBridgeClientMap(environments: Environment[]) {
+    const eventBridgeClientMap = new Map<string, EventBridgeClient>();
+    const promises = [];
+    for (const environment of environments) {
+        promises.push(this.createEventBridgeClients(this.assumeRoleName, environment.accountId, environment.region));
+    }
+    const eventBridgeClients = await Promise.all(promises);
+    eventBridgeClients.forEach((client) => {
+        eventBridgeClientMap.set(`${client.accountId}-${client.region}`, client.client);
+    });
+    return eventBridgeClientMap;
+  }
+
+  async createEventBridgeClients(assumeRoleName: string, accountId: string, region: string) {
+    const credentials = await this.sts.getCredentialsForAccountAndRole(accountId, assumeRoleName);
+    const client = new EventBridgeClient({ credentials, region });
+    return {
+        accountId,
+        region,
+        client,
+    };
+  }
+}

reference-artifacts/Custom-Scripts/lza-upgrade/src/index.ts

@@ -20,6 +20,7 @@ import { PostMigration } from './post-migration';
 import { Preparation } from './preparation';
 import { ResourceMapping } from './resource-mapping';
 import { Snapshot } from './snapshot';
+import { DisableRules } from './disable-rules';
 
 async function main() {
   const args = process.argv.slice(2);
@@ -49,8 +50,8 @@ async function main() {
     case 'asea-prep':
       const preparation = new Preparation(config);
       await preparation.prepareAsea();
-      const disableAccountRules = new Snapshot(config);
-      await disableAccountRules.disableSubscriptionRule();
+      const disableRulesInPrep = await DisableRules.init(config);
+      await disableRulesInPrep.disableAllAccountRules(config.aseaPrefix);
       break;
     case 'lza-prep':
       const lzaPreparation = new Preparation(config);
@@ -74,8 +75,8 @@ async function main() {
       }
       break;
     case 'disable-rules':
-      const disableRules = new Snapshot(config);
-      await disableRules.disableSubscriptionRule();
+      const disableRules = await DisableRules.init(config);
+      await disableRules.disableAllAccountRules(config.aseaPrefix);
       break;
     case 'post-migration':
       await new PostMigration(config, args).process();

reference-artifacts/Custom-Scripts/lza-upgrade/src/preparation/asea.ts

@@ -30,15 +30,11 @@ import {
   ListImagesCommand,
   ListImagesCommandOutput,
 } from '@aws-sdk/client-ecr';
-import { EventBridgeClient, DisableRuleCommand, ResourceNotFoundException } from '@aws-sdk/client-eventbridge';
+import { EventBridgeClient, DisableRuleCommand } from '@aws-sdk/client-eventbridge';
 import { KMSClient, ListAliasesCommand, GetKeyPolicyCommand, PutKeyPolicyCommand } from '@aws-sdk/client-kms';
 
 import { SSMClient, PutParameterCommandInput, PutParameterCommand } from '@aws-sdk/client-ssm';
 
-import { AcceleratorConfig } from '../asea-config';
-import { getCredentials, getAccountList } from '../snapshot/snapshotConfiguration';
-import { AwsCredentialIdentity } from '@aws-sdk/types';
-
 export async function updateSecretsKey(prefix: string, operationsAccountId: string, region: string): Promise<void> {
   const kmsClient = new KMSClient({ region: region });
   const keyAliases = await kmsClient.send(new ListAliasesCommand({}));
@@ -247,43 +243,3 @@ export async function disableASEARules(prefix: string) {
   await Promise.all(disableRulePromises);
   console.log('Disabled Rules');
 }
-
-export async function disableAllAccountAseaRules(aseaConfig: AcceleratorConfig, roleName: string, prefix: string) {
-  const accounts = await getAccountList();
-  const regions = aseaConfig['global-options']['supported-regions'];
-  const accountPromises = []
-  for (const account of accounts) {
-    if (account.Status !== 'SUSPENDED') {
-        const credentials = await getCredentials(account.Id!, roleName);
-        if (credentials === undefined) {
-          continue;
-        }
-        console.log(`Disabling rules in ${account.Id}`);
-        accountPromises.push(
-          disableRules(credentials, regions, prefix),
-        );
-      }
-    }
-  await Promise.all(accountPromises);
-}
-
-async function disableRules(credentials: AwsCredentialIdentity, regions: any, prefix: string){
-  for (const region of regions) {
-    const client = new EventBridgeClient({ region: region, credentials });
-    const suffixes = [
-      'NewLogGroup_rule'
-    ];
-    for (const suffix of suffixes) {
-      try {
-        const command = new DisableRuleCommand({
-          Name: `${prefix}-${suffix}`,
-        });
-        await client.send(command);
-      } catch (e: any) {
-        if (e instanceof ResourceNotFoundException) {
-          continue;
-        }
-      }
-    }
-  }
-}

reference-artifacts/Custom-Scripts/lza-upgrade/src/snapshot.ts

@@ -15,7 +15,6 @@ import { loadAseaConfig } from './asea-config/load';
 import { Config } from './config';
 import { Reset } from './snapshot/common/dynamodb';
 import { getChangedResources } from './snapshot/lib/reporting';
-import { disableAllAccountAseaRules } from './preparation/asea';
 import * as snapshot from './snapshot/snapshotConfiguration';
 
 export class Snapshot {
@@ -79,15 +78,4 @@ export class Snapshot {
     const cleanup = new Reset(this.tableName, this.homeRegion);
     await cleanup.dropTable();
   }
-
-  async disableSubscriptionRule() {
-    const aseaConfig = await loadAseaConfig({
-      filePath: 'raw/config.json',
-      repositoryName: this.aseaConfigRepositoryName,
-      defaultRegion: this.homeRegion,
-      localFilePath: this.localConfigFilePath,
-    });
-
-    await disableAllAccountAseaRules(aseaConfig, this.roleName, this.aseaPrefix.replaceAll('-', ''));
-  }
 }