Merge pull request #252 from aquavis12/feat/ec2-security-keypair-imdsv2

kuettai · web-flow · commit 42d8faf43d90 · 2026-07-03T14:20:47.000+08:00
Feat/ec2 security keypair imdsv2
diff --git a/services/ec2/Ec2.py b/services/ec2/Ec2.py
@@ -10,6 +10,7 @@
 import time
 
 from utils.Tools import _pi
+from utils.Tools import _warn
 
 from utils.Config import Config
 from services.Service import Service
@@ -26,6 +27,8 @@
 from services.ec2.drivers.Ec2Vpc import Ec2Vpc
 from services.ec2.drivers.Ec2NACL import Ec2NACL
 from services.ec2.drivers.Ec2Regional import Ec2Regional
+from services.ec2.drivers.Ec2KeyPair import Ec2KeyPair
+from services.ec2.drivers.Ec2SSM import Ec2SSM
 
 class Ec2(Service):
     CHARTSTYPE = {
@@ -385,9 +388,60 @@ def getNetworkACLs(self):
             result = self.ec2Client.describe_network_acls(
                 NextToken = result.get('NextToken')
             )
-            networkACLs = networkACLs + result.get('NetworkAcls')
+        networkACLs = networkACLs + result.get('NetworkAcls')
         return networkACLs
 
+    def getKeyPairs(self):
+        filters = []
+        if self.tags:
+            filters = self.tags
+
+        result = self.ec2Client.describe_key_pairs(
+            Filters=filters
+        )
+        return result.get('KeyPairs', [])
+
+    def getInstanceKeyNames(self, instances):
+        """Collect all key names actively used by running/stopped instances"""
+        keyNames = set()
+        for instanceArr in instances:
+            for instanceData in instanceArr['Instances']:
+                keyName = instanceData.get('KeyName')
+                if keyName:
+                    keyNames.add(keyName)
+        return keyNames
+
+    def getSSMManagedInstances(self):
+        """Get set of instance IDs managed by SSM.
+
+        Returns a set of managed instance IDs on success, or None if the
+        SSM API call fails (e.g. missing ssm:DescribeInstanceInformation
+        permission, or SSM unavailable in the region). Returning None lets
+        the caller skip the SSM check entirely instead of flagging every
+        instance as unmanaged (false positives).
+        """
+        managedSet = set()
+        try:
+            # MaxResults=50 is the API maximum — request full pages to
+            # minimize round-trips for large fleets.
+            results = self.ssmClient.describe_instance_information(MaxResults=50)
+            for info in results.get('InstanceInformationList', []):
+                managedSet.add(info['InstanceId'])
+            
+            while results.get('NextToken'):
+                results = self.ssmClient.describe_instance_information(
+                    MaxResults=50,
+                    NextToken=results['NextToken']
+                )
+                for info in results.get('InstanceInformationList', []):
+                    managedSet.add(info['InstanceId'])
+        except botocore.exceptions.ClientError as e:
+            _warn("Unable to retrieve SSM managed instances ({}); skipping SSM check to avoid false positives.".format(e.response['Error']['Code']))
+            return None
+        except Exception as e:
+            _warn("Unable to retrieve SSM managed instances ({}); skipping SSM check to avoid false positives.".format(e))
+            return None
+        return managedSet
 
     def getChartGenCost(self):
         '''
@@ -665,6 +719,29 @@ def advise(self):
             objs[f"NACL::{nacl['NetworkAclId']}"] = obj.getInfo()
         
         
+        # Key Pair Checks
+        keyPairs = self.getKeyPairs()
+        instanceKeyNames = self.getInstanceKeyNames(instances)
+        for kp in keyPairs:
+            _pi('EC2::Key Pair', kp['KeyName'])
+            obj = Ec2KeyPair(kp, instanceKeyNames)
+            obj.run(self.__class__)
+            objs[f"KeyPair::{kp['KeyName']}"] = obj.getInfo()
+        
+        # SSM Managed Instance Checks
+        ssmManagedSet = self.getSSMManagedInstances()
+        # Skip the SSM check entirely if the managed-instance lookup failed
+        # (getSSMManagedInstances returns None) to avoid false positives.
+        if ssmManagedSet is not None:
+            for instanceArr in instances:
+                for instanceData in instanceArr['Instances']:
+                    if instanceData['State']['Name'] not in ('running', 'stopped'):
+                        continue
+                    _pi('EC2::SSM Check', instanceData['InstanceId'])
+                    obj = Ec2SSM(instanceData['InstanceId'], ssmManagedSet)
+                    obj.run(self.__class__)
+                    objs[f"SSM::{instanceData['InstanceId']}"] = obj.getInfo()
+        
         if self.getChartGenCost():
             self.setChartData({"EC2 Instance Family Pricing": self.getChartGenCost()})
 
diff --git a/services/ec2/drivers/Ec2Instance.py b/services/ec2/drivers/Ec2Instance.py
@@ -1,6 +1,7 @@
 import boto3
 import botocore
 import datetime
+import re
 import time
 from utils.Config import Config
 
@@ -423,7 +424,73 @@ def _checkEC2HasTag(self):
         if self.ec2InstanceData.get('Tags') is None:
             self.results['EC2HasTag'] = [-1, '']
         return
-    
+
+    def _checkIMDSv2(self):
+        """Check if instance enforces IMDSv2 (HttpTokens = required)"""
+        instance = self.ec2InstanceData
+        metadataOptions = instance.get('MetadataOptions', {})
+        httpTokens = metadataOptions.get('HttpTokens', 'optional')
+        
+        if httpTokens != 'required':
+            self.results['EC2IMDSv2'] = [-1, 'Not enforced']
+        return
+
+    def _checkStoppedTooLong(self):
+        """Flag instances stopped for more than 30 days (still incurring EBS cost)"""
+        instance = self.ec2InstanceData
+        if instance['State']['Name'] != 'stopped':
+            return
+        
+        reason = instance.get('StateTransitionReason', '')
+        # Format: "User initiated (2024-01-15 10:30:00 GMT)"
+        if '(' in reason and ')' in reason:
+            match = re.search(r'\((\d{4}-\d{2}-\d{2})', reason)
+            if match:
+                stopped_date = datetime.datetime.strptime(match.group(1), '%Y-%m-%d').date()
+                days_stopped = (datetime.date.today() - stopped_date).days
+                if days_stopped > 30:
+                    self.results['EC2StoppedTooLong'] = [-1, f'{days_stopped} days']
+        return
+
+    def _checkTerminationProtection(self):
+        """Check if instance has termination protection enabled.
+
+        Stopped instances are intentionally included: they can be critical
+        (e.g. reserved for DR failover) and still benefit from deletion
+        protection. Only terminated/shutting-down instances are skipped.
+
+        This uses describe_instance_attribute, which has no batch form and
+        is therefore called once per instance. The shared boto client is
+        configured with standard retry mode (max_attempts=5), so transient
+        throttling is retried automatically by botocore. If throttling still
+        surfaces after retries, the check is skipped (with a warning) rather
+        than emitting a misleading result.
+        """
+        instance = self.ec2InstanceData
+        
+        # Skip terminated/shutting-down instances (nothing left to protect)
+        if instance['State']['Name'] in ('terminated', 'shutting-down'):
+            return
+        
+        try:
+            resp = self.ec2Client.describe_instance_attribute(
+                InstanceId=instance['InstanceId'],
+                Attribute='disableApiTermination'
+            )
+            protected = resp.get('DisableApiTermination', {}).get('Value', False)
+            if not protected:
+                self.results['EC2NoTerminationProtection'] = [-1, 'Disabled']
+        except botocore.exceptions.ClientError as e:
+            code = e.response['Error']['Code']
+            if code in ('RequestLimitExceeded', 'Throttling', 'ThrottlingException'):
+                _warn("Throttled while checking termination protection for {}; skipping.".format(instance['InstanceId']), forcePrint=False)
+            # Other client errors (e.g. permission) are skipped silently to
+            # avoid emitting a false 'Disabled' result.
+            return
+        except Exception:
+            return
+        return
+
     def checkInstanceTypeAvailable(self, instanceType):
         resp = self.ec2Client.describe_instance_type_offerings(
             LocationType='region',
diff --git a/services/ec2/drivers/Ec2KeyPair.py b/services/ec2/drivers/Ec2KeyPair.py
@@ -0,0 +1,21 @@
+import botocore
+
+from services.Evaluator import Evaluator
+
+class Ec2KeyPair(Evaluator):
+    """Check for EC2 key pairs that are not associated with any running instance."""
+    
+    def __init__(self, keyPair, instanceKeyNames):
+        super().__init__()
+        self.keyPair = keyPair
+        self.instanceKeyNames = instanceKeyNames
+
+        self._resourceName = keyPair['KeyName']
+
+        self.init()
+        
+    def _checkKeyPairInUse(self):
+        """Flag key pairs not attached to any running instance"""
+        keyName = self.keyPair['KeyName']
+        if keyName not in self.instanceKeyNames:
+            self.results['EC2KeyPairNotInUse'] = [-1, keyName]
diff --git a/services/ec2/drivers/Ec2SSM.py b/services/ec2/drivers/Ec2SSM.py
@@ -0,0 +1,20 @@
+import botocore
+
+from services.Evaluator import Evaluator
+
+class Ec2SSM(Evaluator):
+    """Check if an EC2 instance is managed by AWS Systems Manager."""
+    
+    def __init__(self, instanceId, ssmManagedSet):
+        super().__init__()
+        self.instanceId = instanceId
+        self.ssmManagedSet = ssmManagedSet
+
+        self._resourceName = instanceId
+
+        self.init()
+        
+    def _checkSSMManaged(self):
+        """Flag instances not registered with SSM (no patching/compliance visibility)"""
+        if self.instanceId not in self.ssmManagedSet:
+            self.results['EC2SSMNotManaged'] = [-1, 'Not managed']
diff --git a/services/ec2/ec2.reporter.json b/services/ec2/ec2.reporter.json
@@ -51,6 +51,75 @@
             "[Elastic IP Charges]<https://aws.amazon.com/premiumsupport/knowledge-center/elastic-ip-charges/>"       
         ]
     },
+    "EC2KeyPairNotInUse": {
+        "category": "S",
+        "^description": "Unused Key Pairs: {$COUNT} of your EC2 key pairs are not associated with any running instance. Remove unused key pairs to reduce the attack surface and improve security hygiene.",
+        "downtime": 0,
+        "slowness": 0,
+        "additionalCost": 0,
+        "criticality": "L",
+        "needFullTest": 0,
+        "shortDesc": "Remove Unused Key Pairs",
+        "ref": [
+            "[Amazon EC2 key pairs]<https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/ec2-key-pairs.html>",
+            "[Security best practices]<https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/ec2-security.html>"
+        ]
+    },
+    "EC2IMDSv2": {
+        "category": "S",
+        "^description": "IMDSv2 Not Enforced: {$COUNT} of your instances have not enforced IMDSv2 (Instance Metadata Service v2). IMDSv2 adds session-based authentication and mitigates SSRF attacks. Set HttpTokens to 'required' to enforce IMDSv2.",
+        "downtime": 0,
+        "slowness": 0,
+        "additionalCost": 0,
+        "criticality": "H",
+        "needFullTest": 1,
+        "shortDesc": "Enforce IMDSv2",
+        "ref": [
+            "[Configure IMDSv2]<https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/configuring-instance-metadata-service.html>",
+            "[IMDSv2 Best Practices]<https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instance-metadata-transition-to-version-2.html>"
+        ]
+    },
+    "EC2StoppedTooLong": {
+        "category": "C",
+        "^description": "Stopped Instances (>30 days): {$COUNT} of your instances have been in a stopped state for more than 30 days. Stopped instances still incur EBS storage costs. Consider creating an AMI and terminating the instance, or terminate if no longer needed.",
+        "downtime": 0,
+        "slowness": 0,
+        "additionalCost": 0,
+        "criticality": "M",
+        "needFullTest": 0,
+        "shortDesc": "Terminate Long-Stopped Instances",
+        "ref": [
+            "[Instance lifecycle]<https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/ec2-instance-lifecycle.html>",
+            "[Stop and start instances]<https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/Stop_Start.html>"
+        ]
+    },
+    "EC2NoTerminationProtection": {
+        "category": "R",
+        "^description": "No Termination Protection: {$COUNT} of your instances do not have termination protection enabled. Enable termination protection to prevent accidental deletion of critical instances.",
+        "downtime": 0,
+        "slowness": 0,
+        "additionalCost": 0,
+        "criticality": "M",
+        "needFullTest": 0,
+        "shortDesc": "Enable Termination Protection",
+        "ref": [
+            "[Enable termination protection]<https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/Using_ChangingDisableAPITermination.html>"
+        ]
+    },
+    "EC2SSMNotManaged": {
+        "category": "O",
+        "^description": "SSM Not Managed: {$COUNT} of your instances are not managed by AWS Systems Manager. SSM provides patching, compliance visibility, remote access, and inventory management. Install the SSM agent and attach the required IAM role.",
+        "downtime": 0,
+        "slowness": 0,
+        "additionalCost": 0,
+        "criticality": "M",
+        "needFullTest": 0,
+        "shortDesc": "Register with SSM",
+        "ref": [
+            "[Setting up Systems Manager]<https://docs.aws.amazon.com/systems-manager/latest/userguide/systems-manager-setting-up.html>",
+            "[SSM Agent]<https://docs.aws.amazon.com/systems-manager/latest/userguide/ssm-agent.html>"
+        ]
+    },
     "EC2MemoryMonitor": {
         "category": "P",
         "^description": "EC2 Memory Monitoring: Memory monitoring has not been enabled for {$COUNT} of your instances. Install CloudWatch agent and enable the monitoring",
