fix: Fix bug where default automatic remediation configuration was in…

…valid
vibe-io · Jan 11, 2024 · 96c763b · 96c763b
1 parent 3d4c814
commit 96c763b
Show file tree

Hide file tree

Showing 5 changed files with 162 additions and 5 deletions.
diff --git a/API.md b/API.md
diff --git a/src/ec2-required-role-config-rule.ts b/src/ec2-required-role-config-rule.ts
@@ -1,4 +1,4 @@
-import { Resource, ResourceProps } from 'aws-cdk-lib';
+import { Duration, Resource, ResourceProps } from 'aws-cdk-lib';
 import { CustomRule, ResourceType, RuleScope } from 'aws-cdk-lib/aws-config';
 import { IInstanceProfile, IRole, InstanceProfile, Role, ServicePrincipal } from 'aws-cdk-lib/aws-iam';
 import { IFunction } from 'aws-cdk-lib/aws-lambda';
@@ -10,6 +10,8 @@ import { EvaluationFunction } from './evaluation-function';
 export interface Ec2RequiredRoleRemediationOptions {
   readonly automatic?: boolean;
   readonly enabled?: boolean;
+  readonly maxAutomaticAttempts?: number;
+  readonly retryPeriod?: Duration;
   readonly role?: IRole;
 }
 
@@ -61,6 +63,8 @@ export class Ec2RequiredRoleConfigRule extends Resource {
       new Ec2RequiredRoleRemediationConfiguration(this, 'remediation', {
         automatic: props.remediation?.automatic,
         instanceProfile: this.defaultInstanceProfile,
+        maxAutomaticAttempts: props.remediation?.maxAutomaticAttempts,
+        retryPeriod: props.remediation?.retryPeriod,
         rule: this.rule,
       });
     }

diff --git a/src/ec2-required-role-remediation-configuration.ts b/src/ec2-required-role-remediation-configuration.ts
@@ -1,38 +1,59 @@
-import { Resource, ResourceProps } from 'aws-cdk-lib';
+import { Duration, Resource, ResourceProps } from 'aws-cdk-lib';
 import { CfnRemediationConfiguration, IRule } from 'aws-cdk-lib/aws-config';
-import { IInstanceProfile } from 'aws-cdk-lib/aws-iam';
+import { IInstanceProfile, Role, ServicePrincipal } from 'aws-cdk-lib/aws-iam';
 import { IConstruct } from 'constructs';
 import { Ec2RequiredRoleRemediationDocument } from './ec2-required-role-remediation-document';
 
 
 export interface Ec2RequiredRoleRemediationConfigurationProps extends ResourceProps {
   readonly automatic?: boolean;
   readonly instanceProfile: IInstanceProfile;
+  readonly maxAutomaticAttempts?: number;
+  readonly retryPeriod?: Duration;
   readonly rule: IRule;
 }
 
 export class Ec2RequiredRoleRemediationConfiguration extends Resource {
+  public static readonly DEFAULT_MAX_AUTOMATIC_ATTEMPTS: number = 5;
+  public static readonly DEFAULT_RETRY_PERIOD: Duration = Duration.seconds(60);
+
   public readonly automatic: boolean;
   public readonly instanceProfile: IInstanceProfile;
+  public readonly maxAutomaticAttempts: number;
+  public readonly retryPeriod: Duration;
   public readonly rule: IRule;
 
+  public readonly automationRole?: Role;
+
 
   public constructor(scope: IConstruct, id: string, props: Ec2RequiredRoleRemediationConfigurationProps) {
     super(scope, id, props);
 
     this.automatic = props.automatic ?? false;
     this.instanceProfile = props.instanceProfile;
+    this.maxAutomaticAttempts = props.maxAutomaticAttempts ?? Ec2RequiredRoleRemediationConfiguration.DEFAULT_MAX_AUTOMATIC_ATTEMPTS;
+    this.retryPeriod = props.retryPeriod ?? Ec2RequiredRoleRemediationConfiguration.DEFAULT_RETRY_PERIOD;
     this.rule = props.rule;
 
     const automation = new Ec2RequiredRoleRemediationDocument(this, 'automation');
 
+    if (this.automatic) {
+      this.automationRole = new Role(this, 'automation-role', {
+        assumedBy: new ServicePrincipal('ssm.amazonaws.com'),
+        path: '/service-role/',
+      });
+
+      automation.grantExecute(this.automationRole);
+    }
+
     new CfnRemediationConfiguration(this, 'Resource', {
       automatic: this.automatic,
       configRuleName: this.rule.configRuleName,
+      maximumAutomaticAttempts: this.automatic ? this.maxAutomaticAttempts : undefined,
       parameters: {
         AutomationAssumeRole: {
           StaticValue: {
-            Values: [],
+            Values: this.automationRole ? [this.automationRole.roleArn] : [],
           },
         },
         InstanceId: {
@@ -49,6 +70,7 @@ export class Ec2RequiredRoleRemediationConfiguration extends Resource {
         },
       },
       resourceType: 'AWS::EC2::Instance',
+      retryAttemptSeconds: this.automatic ? this.retryPeriod.toSeconds() : undefined,
       targetId: automation.documentName,
       targetType: 'SSM_DOCUMENT',
     });

diff --git a/src/ec2-required-role-remediation-document.ts b/src/ec2-required-role-remediation-document.ts
@@ -1,4 +1,5 @@
 import { ArnFormat, Resource, ResourceProps } from 'aws-cdk-lib';
+import { IGrantable, PolicyStatement } from 'aws-cdk-lib/aws-iam';
 import { CfnDocument } from 'aws-cdk-lib/aws-ssm';
 import { IConstruct } from 'constructs';
 
@@ -84,4 +85,20 @@ export class Ec2RequiredRoleRemediationDocument extends Resource {
   public automationDefinitionArnForVersion(version: string): string {
     return `${this.automationDefinitionArn}:${version}`;
   }
+
+  public grantExecute(principal: IGrantable): void {
+    principal.grantPrincipal.addToPrincipalPolicy(new PolicyStatement({
+      actions: [
+        'ec2:AssociateIamInstanceProfile',
+      ],
+      resources: [
+        this.stack.formatArn({
+          arnFormat: ArnFormat.SLASH_RESOURCE_NAME,
+          resource: 'instance',
+          resourceName: '*',
+          service: 'ec2',
+        }),
+      ],
+    }));
+  }
 }
diff --git a/test/index.test.ts b/test/index.test.ts
@@ -1,4 +1,4 @@
-import { Stack } from 'aws-cdk-lib';
+import { ArnFormat, CfnResource, Stack } from 'aws-cdk-lib';
 import { Template } from 'aws-cdk-lib/assertions';
 import { Ec2RequiredRoleConfigRule } from '../src';
 
@@ -12,4 +12,96 @@ test('template should contain default set of resources', () => {
   template.resourceCountIs('AWS::Config::RemediationConfiguration', 1);
   template.resourceCountIs('AWS::Lambda::Function', 1);
   template.resourceCountIs('AWS::SSM::Document', 1);
+});
+
+test('automatic remediation should be reflected in configuration', () => {
+  const stack = new Stack();
+  new Ec2RequiredRoleConfigRule(stack, 'rule', {
+    remediation: {
+      automatic: true,
+    },
+  });
+
+  const template = Template.fromStack(stack);
+
+  template.hasResourceProperties('AWS::Config::RemediationConfiguration', {
+    Automatic: true,
+  });
+});
+
+test('automatic remediation should have required properties', () => {
+  const stack = new Stack();
+  new Ec2RequiredRoleConfigRule(stack, 'rule', {
+    remediation: {
+      automatic: true,
+    },
+  });
+
+  const template = Template.fromStack(stack);
+
+  template.hasResourceProperties('AWS::Config::RemediationConfiguration', {
+    MaximumAutomaticAttempts: 5,
+    RetryAttemptSeconds: 60,
+  });
+});
+
+test('automatic remediation should properly configure permissions', () => {
+  const stack = new Stack();
+  const rule = new Ec2RequiredRoleConfigRule(stack, 'rule', {
+    remediation: {
+      automatic: true,
+    },
+  });
+
+  const template = Template.fromStack(stack);
+
+  const remediation = rule.node.tryFindChild('remediation');
+  const automationRoleL2 = remediation?.node.tryFindChild('automation-role');
+  const automationRoleL1 = automationRoleL2?.node.defaultChild;
+
+  expect(CfnResource.isCfnResource(automationRoleL1)).toBe(true);
+  const automationRoleResource = automationRoleL1 as CfnResource;
+
+  template.hasResourceProperties('AWS::IAM::Role', {
+    AssumeRolePolicyDocument: {
+      Statement: [{
+        Action: 'sts:AssumeRole',
+        Effect: 'Allow',
+        Principal: {
+          Service: 'ssm.amazonaws.com',
+        },
+      }],
+    },
+    Path: '/service-role/',
+  });
+
+  template.hasResourceProperties('AWS::IAM::Policy', {
+    PolicyDocument: {
+      Statement: [{
+        Action: 'ec2:AssociateIamInstanceProfile',
+        Effect: 'Allow',
+        Resource: stack.resolve(stack.formatArn({
+          arnFormat: ArnFormat.SLASH_RESOURCE_NAME,
+          resource: 'instance',
+          resourceName: '*',
+          service: 'ec2',
+        })),
+      }],
+    },
+    Roles: [
+      stack.resolve(automationRoleResource.ref),
+    ],
+  });
+
+  template.hasResourceProperties('AWS::Config::RemediationConfiguration', {
+    Parameters: {
+      AutomationAssumeRole: {
+        StaticValue: {
+          Values: [
+            stack.resolve(automationRoleResource.getAtt('Arn')),
+          ],
+        },
+      },
+    },
+  });
 });