feat: add ami-updater lambda function with tests and config

lambdas/functions/ami-updater/.eslintrc.js

@@ -0,0 +1,20 @@
+module.exports = {
+  parser: '@typescript-eslint/parser',
+  extends: [
+    'eslint:recommended',
+    'plugin:@typescript-eslint/recommended',
+  ],
+  parserOptions: {
+    ecmaVersion: 2020,
+    sourceType: 'module',
+  },
+  env: {
+    node: true,
+    es6: true,
+  },
+  rules: {
+    '@typescript-eslint/explicit-function-return-type': 'off',
+    '@typescript-eslint/no-explicit-any': 'off',
+    '@typescript-eslint/no-unused-vars': ['error', { argsIgnorePattern: '^_' }],
+  },
+};

lambdas/functions/ami-updater/package.json

@@ -0,0 +1,51 @@
+{
+    "name": "@aws-github-runner/ami-updater",
+    "version": "1.0.0",
+    "main": "lambda.ts",
+    "type": "module",
+    "license": "MIT",
+    "scripts": {
+        "start": "ts-node-dev src/local.ts",
+        "test": "vitest run",
+        "test:watch": "vitest",
+        "lint": "eslint src",
+        "watch": "ts-node-dev --respawn --exit-child src/local.ts",
+        "build": "ncc build src/lambda.ts -o dist",
+        "dist": "yarn build && cp package.json dist/ && cd dist && zip ../ami-updater.zip *",
+        "format": "prettier --write \"**/*.ts\"",
+        "format-check": "prettier --check \"**/*.ts\"",
+        "all": "yarn build && yarn format && yarn lint && yarn test"
+    },
+    "dependencies": {
+        "@aws-github-runner/aws-powertools-util": "*",
+        "@aws-github-runner/aws-ssm-util": "*",
+        "@aws-sdk/client-ec2": "^3.767.0",
+        "@aws-sdk/client-ssm": "^3.759.0"
+    },
+    "devDependencies": {
+        "@aws-sdk/types": "^3.734.0",
+        "@types/aws-lambda": "^8.10.147",
+        "@types/node": "^20.10.4",
+        "@typescript-eslint/eslint-plugin": "^6.13.2",
+        "@typescript-eslint/parser": "^6.13.2",
+        "@vercel/ncc": "^0.38.3",
+        "aws-sdk-client-mock": "^4.1.0",
+        "aws-sdk-client-mock-jest": "^4.1.0",
+        "eslint": "^8.55.0",
+        "prettier": "^3.0.0",
+        "typescript": "^5.3.3",
+        "vitest": "^3.0.9"
+    },
+    "nx": {
+        "includedScripts": [
+            "build",
+            "dist",
+            "format",
+            "format-check",
+            "lint",
+            "start",
+            "watch",
+            "all"
+        ]
+    }
+}

lambdas/functions/ami-updater/src/ami.test.ts

@@ -0,0 +1,124 @@
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+import {
+  EC2Client,
+  DescribeImagesCommand,
+  DescribeLaunchTemplatesCommand,
+  DescribeLaunchTemplateVersionsCommand,
+  CreateLaunchTemplateVersionCommand,
+  ModifyLaunchTemplateCommand,
+} from '@aws-sdk/client-ec2';
+import { AMIManager } from './ami';
+
+vi.mock('@aws-sdk/client-ec2');
+vi.mock('../../shared/aws-powertools-util', () => ({
+  logger: {
+    info: vi.fn(),
+    warn: vi.fn(),
+    error: vi.fn(),
+  },
+}));
+
+describe('AMIManager', () => {
+  let ec2Client: EC2Client;
+  let amiManager: AMIManager;
+
+  beforeEach(() => {
+    ec2Client = new EC2Client({});
+    amiManager = new AMIManager(ec2Client);
+    vi.clearAllMocks();
+  });
+
+  describe('getLatestAmi', () => {
+    it('should return the latest AMI ID', async () => {
+      const mockResponse = {
+        Images: [
+          { ImageId: 'ami-2', CreationDate: '2023-12-02' },
+          { ImageId: 'ami-1', CreationDate: '2023-12-01' },
+        ],
+      };
+
+      vi.mocked(ec2Client.send).mockResolvedValueOnce(mockResponse);
+
+      const config = {
+        owners: ['self'],
+        filters: [{ Name: 'tag:Environment', Values: ['prod'] }],
+      };
+
+      const result = await amiManager.getLatestAmi(config);
+      expect(result).toBe('ami-2');
+      expect(ec2Client.send).toHaveBeenCalledWith(expect.any(DescribeImagesCommand));
+    });
+
+    it('should throw error when no AMIs found', async () => {
+      vi.mocked(ec2Client.send).mockResolvedValueOnce({ Images: [] });
+
+      const config = {
+        owners: ['self'],
+        filters: [{ Name: 'tag:Environment', Values: ['prod'] }],
+      };
+
+      await expect(amiManager.getLatestAmi(config)).rejects.toThrow('No matching AMIs found');
+    });
+  });
+
+  describe('updateLaunchTemplate', () => {
+    it('should update launch template with new AMI ID', async () => {
+      vi.mocked(ec2Client.send)
+        .mockResolvedValueOnce({
+          // getCurrentAmiId - DescribeLaunchTemplatesCommand
+          LaunchTemplates: [{ LatestVersionNumber: 1 }],
+        })
+        .mockResolvedValueOnce({
+          // getCurrentAmiId - DescribeLaunchTemplateVersionsCommand
+          LaunchTemplateVersions: [{ LaunchTemplateData: { ImageId: 'ami-old' } }],
+        })
+        .mockResolvedValueOnce({
+          // updateLaunchTemplate - DescribeLaunchTemplatesCommand
+          LaunchTemplates: [{ LatestVersionNumber: 1 }],
+        });
+
+      const result = await amiManager.updateLaunchTemplate('test-template', 'ami-new', false);
+
+      expect(result.success).toBe(true);
+      expect(result.message).toBe('Updated successfully');
+      expect(ec2Client.send).toHaveBeenCalledWith(expect.any(CreateLaunchTemplateVersionCommand));
+      expect(ec2Client.send).toHaveBeenCalledWith(expect.any(ModifyLaunchTemplateCommand));
+    });
+
+    it('should not update if AMI ID is the same', async () => {
+      vi.mocked(ec2Client.send)
+        .mockResolvedValueOnce({
+          // getCurrentAmiId - DescribeLaunchTemplatesCommand
+          LaunchTemplates: [{ LatestVersionNumber: 1 }],
+        })
+        .mockResolvedValueOnce({
+          // getCurrentAmiId - DescribeLaunchTemplateVersionsCommand
+          LaunchTemplateVersions: [{ LaunchTemplateData: { ImageId: 'ami-1' } }],
+        });
+
+      const result = await amiManager.updateLaunchTemplate('test-template', 'ami-1', false);
+
+      expect(result.success).toBe(true);
+      expect(result.message).toBe('Already using latest AMI');
+      expect(ec2Client.send).not.toHaveBeenCalledWith(expect.any(CreateLaunchTemplateVersionCommand));
+    });
+
+    it('should handle dry run mode', async () => {
+      vi.mocked(ec2Client.send)
+        .mockResolvedValueOnce({
+          // getCurrentAmiId - DescribeLaunchTemplatesCommand
+          LaunchTemplates: [{ LatestVersionNumber: 1 }],
+        })
+        .mockResolvedValueOnce({
+          // getCurrentAmiId - DescribeLaunchTemplateVersionsCommand
+          LaunchTemplateVersions: [{ LaunchTemplateData: { ImageId: 'ami-old' } }],
+        });
+
+      const result = await amiManager.updateLaunchTemplate('test-template', 'ami-new', true);
+
+      expect(result.success).toBe(true);
+      expect(result.message).toBe('Would update AMI (Dry Run)');
+      expect(ec2Client.send).not.toHaveBeenCalledWith(expect.any(CreateLaunchTemplateVersionCommand));
+    });
+  });
+});

lambdas/functions/ami-updater/src/ami.ts

@@ -0,0 +1,141 @@
+import {
+  EC2Client,
+  DescribeImagesCommand,
+  DescribeLaunchTemplatesCommand,
+  DescribeLaunchTemplateVersionsCommand,
+  CreateLaunchTemplateVersionCommand,
+  ModifyLaunchTemplateCommand,
+  Image,
+  Filter,
+} from '@aws-sdk/client-ec2';
+import { logger } from '@aws-github-runner/aws-powertools-util';
+
+export interface AMIFilterConfig {
+  owners: string[];
+  filters: Filter[];
+}
+
+export class AMIManager {
+  constructor(private readonly ec2Client: EC2Client) {}
+
+  async getLatestAmi(config: AMIFilterConfig): Promise<string> {
+    try {
+      const response = await this.ec2Client.send(
+        new DescribeImagesCommand({
+          Owners: config.owners,
+          Filters: config.filters,
+        }),
+      );
+
+      if (!response.Images || response.Images.length === 0) {
+        throw new Error('No matching AMIs found');
+      }
+
+      // Sort by creation date to get the latest
+      const sortedImages = response.Images.sort((a: Image, b: Image) => {
+        return (b.CreationDate || '').localeCompare(a.CreationDate || '');
+      });
+
+      if (!sortedImages[0].ImageId) {
+        throw new Error('Latest AMI has no ImageId');
+      }
+
+      return sortedImages[0].ImageId;
+    } catch (error) {
+      logger.error('Error getting latest AMI', { error });
+      throw error;
+    }
+  }
+
+  async getCurrentAmiId(templateName: string): Promise<string | null> {
+    try {
+      const response = await this.ec2Client.send(
+        new DescribeLaunchTemplatesCommand({
+          LaunchTemplateNames: [templateName],
+        }),
+      );
+
+      if (!response.LaunchTemplates || response.LaunchTemplates.length === 0) {
+        logger.warn(`Launch template ${templateName} not found`);
+        return null;
+      }
+
+      const latestVersion = response.LaunchTemplates[0].LatestVersionNumber?.toString();
+      if (!latestVersion) {
+        logger.warn('No latest version found for launch template');
+        return null;
+      }
+
+      const templateData = await this.ec2Client.send(
+        new DescribeLaunchTemplateVersionsCommand({
+          LaunchTemplateName: templateName,
+          Versions: [latestVersion],
+        }),
+      );
+
+      return templateData.LaunchTemplateVersions?.[0]?.LaunchTemplateData?.ImageId || null;
+    } catch (error) {
+      logger.error(`Error getting current AMI ID for ${templateName}`, { error });
+      return null;
+    }
+  }
+
+  async updateLaunchTemplate(
+    templateName: string,
+    amiId: string,
+    dryRun: boolean,
+  ): Promise<{ success: boolean; message: string }> {
+    try {
+      const currentAmi = await this.getCurrentAmiId(templateName);
+      if (!currentAmi) {
+        return { success: false, message: 'Failed to get current AMI ID' };
+      }
+
+      if (currentAmi === amiId) {
+        logger.info(`Template ${templateName} already using latest AMI ${amiId}`);
+        return { success: true, message: 'Already using latest AMI' };
+      }
+
+      if (dryRun) {
+        logger.info(`[DRY RUN] Would update template ${templateName} from AMI ${currentAmi} to ${amiId}`);
+        return { success: true, message: 'Would update AMI (Dry Run)' };
+      }
+
+      // Get the latest version of the launch template
+      const response = await this.ec2Client.send(
+        new DescribeLaunchTemplatesCommand({
+          LaunchTemplateNames: [templateName],
+        }),
+      );
+
+      if (!response.LaunchTemplates || response.LaunchTemplates.length === 0) {
+        logger.warn(`Launch template ${templateName} not found`);
+        return { success: false, message: 'Template not found' };
+      }
+
+      // Create new version with updated AMI ID
+      await this.ec2Client.send(
+        new CreateLaunchTemplateVersionCommand({
+          LaunchTemplateName: templateName,
+          SourceVersion: response.LaunchTemplates[0].LatestVersionNumber?.toString(),
+          LaunchTemplateData: { ImageId: amiId },
+        }),
+      );
+
+      // Set the new version as default
+      await this.ec2Client.send(
+        new ModifyLaunchTemplateCommand({
+          LaunchTemplateName: templateName,
+          DefaultVersion: '$Latest',
+        }),
+      );
+
+      logger.info(`Successfully updated launch template ${templateName} from AMI ${currentAmi} to ${amiId}`);
+      return { success: true, message: 'Updated successfully' };
+    } catch (error) {
+      const errorMessage = error instanceof Error ? error.message : 'Unknown error';
+      logger.error(`Error updating launch template ${templateName}`, { error });
+      return { success: false, message: `Error: ${errorMessage}` };
+    }
+  }
+}