feat: add ami start and record lambda (#112)

components/base.yml

@@ -0,0 +1,24 @@
+name: InstallBase
+description: This document installs all required packages on top of Amazon Linux 2
+schemaVersion: 1.0
+
+phases:
+    - name: build
+      steps:
+        - name: InstallBase
+          action: ExecuteBash
+          inputs:
+            commands:
+                - sudo rpm --import https://packages.wazuh.com/key/GPG-KEY-WAZUH
+                - sudo echo -e "[wazuh]\ngpgcheck=1\ngpgkey=https://packages.wazuh.com/key/GPG-KEY-WAZUH\nenabled=1\nname=EL-\$releasever - Wazuh\nbaseurl=https://packages.wazuh.com/4.x/yum/\nprotect=1" >> /etc/yum.repos.d/wazuh.repo
+                - WAZUH_MANAGER="10.0.0.2" sudo yum install -y wazuh-agent
+                - sudo systemctl daemon-reload
+                - sudo systemctl enable wazuh-agent
+                - sudo systemctl start wazuh-agent
+                - sudo sed -i "s/^enabled=1/enabled=0/" /etc/yum.repos.d/wazuh.repo
+                - sudo amazon-linux-extras install epel
+                - sudo yum update -y
+                - sudo yum install -y redis
+                - sudo yum install -y postgresql14
+                - sudo yum install -y clamav freshclam clamav-update
+

deploy_imagebuilder.sh

@@ -97,7 +97,7 @@ echo "AWS CDK is installed successfully."
 
 # Check for AWS CDK
 if ! command -v cdk &> /dev/null; then
-    echo "AWS CDK could not be found. Please rerun 'bash install.sh' with Sudo access and ensure the command is available within the \$PATH"
+    echo "AWS CDK could not be found. Please rerun 'bash deploy_imagebuilder.sh' with Sudo access and ensure the command is available within the \$PATH"
     exit 1
 fi
 
@@ -120,7 +120,7 @@ esac
 
 # Check if AWS CLI installation was successful
 if ! command -v aws &> /dev/null; then
-    echo "AWS CLI could not be found. Please rerun 'bash install.sh' with Sudo access and ensure the command is available within the $PATH"
+    echo "AWS CLI could not be found. Please rerun 'bash deploy_imagebuilder.sh' with Sudo access and ensure the command is available within the $PATH"
     exit 1
 fi
 

install.sh

@@ -226,6 +226,26 @@ show_install_options() {
     echo "${bold}${green}2. Production Ready ${reset} - ${bold}${blue}Optimized for scalability and performance, leveraging the power of AWS EKS for robust, enterprise-grade deployments.${reset}"
 }
 
+declare base_ami,envoy_ami,squid_ami
+
+check_image_builder() {
+    ssm=$(aws ssm get-parameters \
+            --names base_image_ami squid_image_ami envoy_image_ami \
+            --query "Parameters[*].{Name:Name,Value:Value}" --output json)
+
+    length=$(echo "$ssm" | jq -r ".|length")
+
+    if [ "$length" -lt 3 ]; then
+        display_error "Unable to find base images for Proxy Servers. Please run the following command: bash deploy_imagebuilder.sh\nIf you have done it already please wait for 15-20 mins until it builds the images"
+        exit 1
+    fi
+
+    base_ami=$(echo "$ssm" | jq '.[]|select(.Name=="base_image_ami")|.Value')
+    envoy_ami=$(echo "$ssm" | jq '.[]|select(.Name=="envoy_image_ami")|.Value')
+    squid_ami=$(echo "$ssm" | jq '.[]|select(.Name=="squid_image_ami")|.Value')
+
+}
+
 # Function to read user input until a valid choice is made
 get_user_choice() {
     while true; do
@@ -247,6 +267,7 @@ get_user_choice() {
 clear
 list_services
 echo
+check_image_builder
 show_install_options
 get_user_choice
 
@@ -481,7 +502,7 @@ if [[ "$INSTALLATION_MODE" == 2 ]]; then
         aws iam delete-role --role-name $ROLE_NAME 2>/dev/null
         cdk bootstrap aws://$AWS_ACCOUNT_ID/$AWS_DEFAULT_REGION -c aws_arn=$AWS_ARN
     fi
-    if cdk deploy --require-approval never -c db_pass=$DB_PASS -c admin_api_key=$ADMIN_API_KEY -c aws_arn=$AWS_ARN -c master_enc_key=$MASTER_ENC_KEY -c vpn_ips=$VPN_IPS $LOCKER; then
+    if cdk deploy --require-approval never -c db_pass=$DB_PASS -c admin_api_key=$ADMIN_API_KEY -c aws_arn=$AWS_ARN -c master_enc_key=$MASTER_ENC_KEY -c vpn_ips=$VPN_IPS -c base_ami=$base_ami -c envoy_ami=$envoy_ami -c squid_ami=$squid_ami $LOCKER; then
         # Wait for the EKS Cluster to be deployed
         echo $(aws eks create-addon --cluster-name hs-eks-cluster --addon-name amazon-cloudwatch-observability)
         aws eks update-kubeconfig --region "$AWS_DEFAULT_REGION" --name hs-eks-cluster

lib/aws/image_builder_stack.ts

@@ -2,10 +2,14 @@ import * as cdk from "aws-cdk-lib";
 import * as image_builder from "aws-cdk-lib/aws-imagebuilder"
 import * as iam from "aws-cdk-lib/aws-iam";
 
+import { Function, Runtime, Code } from "aws-cdk-lib/aws-lambda";
+import { LambdaSubscription } from 'aws-cdk-lib/aws-sns-subscriptions';
 import { Vpc } from './networking';
+import { Topic } from 'aws-cdk-lib/aws-sns';
 import { Construct } from "constructs";
 import { ImageBuilderConfig, VpcConfig } from "./config";
-import { MachineImage } from 'aws-cdk-lib/aws-ec2';
+import { aws_logs as logs } from 'aws-cdk-lib';
+import { MachineImage, SubnetType, SecurityGroup } from 'aws-cdk-lib/aws-ec2';
 
 import { readFileSync } from "fs";
 
@@ -19,13 +23,25 @@ type ImageBuilderProperties = {
     baseimageArn: string;
     description: string;
     compFilePath: string;
+    snsTopicArn: string;
+    subnetId: string;
+    sgId: string;
+
+}
+
+type RecordLambdaProperties = {
+    name: string;
+    id: string;
+    role: iam.Role;
+    ssm_key: string;
+    topic: Topic,
 }
 
 function CreateImagePipeline(
     stack: ImageBuilderStack,
     role: iam.Role,
     props: ImageBuilderProperties,
-) {
+): string {
     const component = new image_builder.CfnComponent(stack, props.comp_id, {
         name: props.comp_name,
         description: props.description,
@@ -48,6 +64,9 @@ function CreateImagePipeline(
         name: props.infra_config_name,
         instanceTypes: ["t3.medium"],
         instanceProfileName: props.profile_name,
+        snsTopicArn: props.snsTopicArn,
+        subnetId: props.subnetId,
+        securityGroupIds: [props.sgId]
     })
     squid_infra_config.addDependency(instance_profile)
 
@@ -58,6 +77,8 @@ function CreateImagePipeline(
     })
 
     pipeline.addDependency(squid_infra_config)
+    return pipeline.attrArn
+
 }
 
 export class ImageBuilderStack extends cdk.Stack {
@@ -71,7 +92,19 @@ export class ImageBuilderStack extends cdk.Stack {
             maxAzs: 2,
         };
 
+        const envoy_channel = new Topic(this, 'ImgBuilderNotificationTopicEnvoy', {});
+        const squid_channel = new Topic(this, 'ImgBuilderNotificationTopicSquid', {});
+        const base_channel = new Topic(this, 'ImgBuilderNotificationTopicBase', {});
+
         let vpc = new Vpc(this, vpcConfig);
+
+        let subnetId = vpc.vpc.selectSubnets({ subnetType: SubnetType.PUBLIC }).subnetIds[0];
+        let ib_SG = new SecurityGroup(this, 'image-server-sg', {
+            vpc: vpc.vpc,
+            allowAllOutbound: true,
+            description: 'security group for a image builder server',
+        });
+
         let role = new iam.Role(this, "StationRole", { roleName: "StationRole", assumedBy: new iam.ServicePrincipal("ec2.amazonaws.com") })
 
         role.addManagedPolicy(iam.ManagedPolicy.fromAwsManagedPolicyName("AmazonSSMManagedInstanceCore"))
@@ -89,14 +122,18 @@ export class ImageBuilderStack extends cdk.Stack {
             baseimageArn: base_image_id,
             description: "Image builder for squid",
             compFilePath: "./components/squid.yml",
+            sgId: ib_SG.securityGroupId,
+            subnetId: subnetId,
+            snsTopicArn: squid_channel.topicArn,
         };
 
-        CreateImagePipeline(
+        let envoy_arn = CreateImagePipeline(
             this,
             role,
             squid_properties,
         )
 
+
         let envoy_properties = {
             pipeline_name: "EnvoyImagePipeline",
             recipe_name: "EnvoyImageRecipe",
@@ -107,12 +144,137 @@ export class ImageBuilderStack extends cdk.Stack {
             baseimageArn: base_image_id,
             description: "Image builder for Envoy",
             compFilePath: "./components/envoy.yml",
+            sgId: ib_SG.securityGroupId,
+            subnetId: subnetId,
+            snsTopicArn: envoy_channel.topicArn,
         };
 
-        CreateImagePipeline(
+        let squid_arn = CreateImagePipeline(
             this,
             role,
             envoy_properties,
         )
+
+
+        let base_properties = {
+            pipeline_name: "BaseImagePipeline",
+            recipe_name: "BaseImageRecipe",
+            profile_name: "BaseStationProfile",
+            comp_name: "HyperswitchBaseImageBuilder",
+            comp_id: "install_base_component",
+            infra_config_name: "BaseInfraConfig",
+            baseimageArn: base_image_id,
+            description: "Image builder for Base Image",
+            compFilePath: "./components/base.yml",
+            sgId: ib_SG.securityGroupId,
+            subnetId: subnetId,
+            snsTopicArn: base_channel.topicArn,
+        };
+
+        let base_arn = CreateImagePipeline(
+            this,
+            role,
+            base_properties,
+        )
+
+        const start_ib_code = readFileSync(
+            "lib/aws/lambda/start_ib.py",
+        ).toString();
+
+
+        const lambda_policy = new iam.PolicyDocument({
+            statements: [
+                new iam.PolicyStatement({
+                    effect: iam.Effect.ALLOW,
+                    actions: ['imagebuilder:StartImagePipelineExecution'],
+                    resources: [
+                        `arn:aws:imagebuilder:${cdk.Stack.of(this).region}:${cdk.Stack.of(this).account
+                        }:image/*`,
+                        `arn:aws:imagebuilder:${cdk.Stack.of(this).region}:${cdk.Stack.of(this).account
+                        }:image-pipeline/*`,
+                    ],
+                }),
+                new iam.PolicyStatement({
+                    effect: iam.Effect.ALLOW,
+                    actions: ['ssm:*'],
+                    resources: ["*"],
+                }),
+            ],
+        });
+
+        const lambda_role = new iam.Role(this, "hyperswitch-ib-lambda-role", {
+            assumedBy: new iam.ServicePrincipal("lambda.amazonaws.com"),
+            inlinePolicies: {
+                "ib-start-role": lambda_policy,
+            },
+        });
+
+        const ib_lambda = new Function(this, "hyperswitch-ib-lambda", {
+            functionName: "HyperswitchIbStartLambda",
+            runtime: Runtime.PYTHON_3_9,
+            handler: "index.lambda_handler",
+            code: Code.fromInline(start_ib_code),
+            timeout: cdk.Duration.minutes(15),
+            role: lambda_role,
+            environment: {
+                envoy_image_pipeline_arn: envoy_arn,
+                squid_image_pipeline_arn: squid_arn,
+                base_image_pipeline_arn: base_arn,
+            },
+        });
+
+
+        const triggerIbStart = new cdk.CustomResource(
+            this,
+            "HyperswitchIbStart",
+            {
+                serviceToken: ib_lambda.functionArn,
+            },
+        );
+
+        addRecordLambda(this, {
+            id: "hyperswitch-record-amid-squid",
+            name: "HyperswitchRecordAmiIdSquid",
+            role: lambda_role,
+            topic: squid_channel,
+            ssm_key: "squid_image_ami",
+        })
+
+        addRecordLambda(this, {
+            id: "hyperswitch-record-amid-envoy",
+            name: "HyperswitchRecordAmiIdEnvoy",
+            role: lambda_role,
+            topic: envoy_channel,
+            ssm_key: "envoy_image_ami",
+        })
+
+        addRecordLambda(this, {
+            id: "hyperswitch-record-amid-base",
+            name: "HyperswitchRecordAmiIdBase",
+            role: lambda_role,
+            topic: base_channel,
+            ssm_key: "base_image_ami",
+        })
     }
 }
+
+
+function addRecordLambda(stack: cdk.Stack, props: RecordLambdaProperties) {
+    const record_amid_code = readFileSync(
+        "lib/aws/lambda/record_ami.py"
+    ).toString();
+
+    let ib_record_function = new Function(stack, props.id, {
+        functionName: props.name,
+        runtime: Runtime.PYTHON_3_9,
+        handler: "index.lambda_handler",
+        code: Code.fromInline(record_amid_code),
+        timeout: cdk.Duration.minutes(15),
+        role: props.role,
+        environment: {
+            IMAGE_SSM_NAME: props.ssm_key,
+        },
+    });
+
+    props.topic.addSubscription(new LambdaSubscription(ib_record_function))
+}

lib/aws/lambda/record_ami.py

@@ -0,0 +1,39 @@
+from typing import Union
+import json
+import os
+import logging
+import boto3
+logging.getLogger().setLevel(logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+def lambda_handler(event, _):
+    logger.info(event)
+
+    message_body = event.get("Records")[0].get("Sns").get("Message")
+    json_body = json.loads(message_body)
+    ami = json_body["outputResources"]["amis"][0]["image"]
+    ssm_key = os.environ["IMAGE_SSM_NAME"]
+    logger.info(f"updating ssm {ssm_key}")
+    ssm_client = boto3.client("ssm")
+    try:
+        result = ssm_client.put_parameter(
+            Name=ssm_key,
+            Value=ami,
+            Type="String",
+            DataType="text",
+            Tier="Advanced",
+            Overwrite=True,
+        )
+        logger.info(result)
+    except Exception as e:
+        logger.error(e)
+
+    return create_response(200, {})
+
+
+def create_response(code: int, body: Union[dict, str]):
+    json_content = {
+        "statusCode": code,
+    }
+    return json_content