Release shutdown 1.1.1

Author: aeghbali-mw

aws/ec2-shutdown-lambda/v1/.version

@@ -1 +1 @@
-v1.1.0
+v1.1.1

aws/ec2-shutdown-lambda/v1/ec2-shutdown-lambda.yml

@@ -169,8 +169,6 @@ Resources:
                   instance: Dict[str, Any] = response['Reservations'][0]['Instances'][0]
                   return instance
 
-
-
               # Handle case where an invalid or non-existent instance ID is provided
               except ClientError as e:
                   print(f"Invalid or non-existent EC2 Instance ID: {instance_id}. Error: {e}")
@@ -219,19 +217,35 @@ Resources:
                   # If an instance is not found, an error string will be printed by get_instance().
                   return
               
-              if not is_instance_running(instance):
-                  # If an instance is not running, an error string will be printed by is_instance_running() containing the current state.
-                  return
-              
+              instance_last_start_time = instance.get('LaunchTime')
               tags: Dict[str, str] = {tag['Key']: tag['Value'] for tag in instance.get('Tags', [])}
               
+              if not is_instance_running(instance):
+                  # Delete the shutdown tag if it exists
+                  if tag_to_monitor in tags:
+                      ec2.delete_tags(Resources=[instance_id],Tags=[{'Key': tag_to_monitor}])
+                      print(f"Removed {tag_to_monitor} tag from EC2 instance {instance_id} (instance not running).")
+                  return
+
               if tag_to_monitor in tags:
                   shutdown_time: datetime = datetime.strptime(tags[tag_to_monitor], '%a, %d %b %Y %H:%M:%S GMT')
                   
-                  if datetime.now(UTC).replace(tzinfo=None) >= shutdown_time:
+                  if instance_last_start_time is None:
+                      print("LaunchTime of EC2 instance is None, skipping comparison with shutdown_time.")
+                      return
+
+                  instance_last_start_time_naive = instance_last_start_time.replace(tzinfo=None)
+
+                  # Stale tag: instance was restarted after the scheduled shutdown time, skip action
+                  if instance_last_start_time_naive >= shutdown_time:
+                      print(f"Shutdown tag is stale (instance started at {instance_last_start_time_naive}, shutdown was scheduled for {shutdown_time}). Skipping.")
+                      return
+                  
+                  now = datetime.now(UTC).replace(tzinfo=None)
+                  if now >= shutdown_time:
                       shutdown_instance(instance_id, tag_to_monitor)
                   else:
-                      remaining_minutes: int = int((shutdown_time - datetime.now(UTC).replace(tzinfo=None)).total_seconds() // 60)
+                      remaining_minutes: int = int((shutdown_time - now).total_seconds() // 60)
                       print(f"No action needed. EC2 instance {instance_id} will shut down at {shutdown_time}. Remaining minutes: {remaining_minutes}")
               else:
                   if shutdown_behaviour == 'Never':

aws/ec2-shutdown-lambda/v1/tests/test-ec2-shutdown-lambda.py

@@ -33,7 +33,7 @@ def load_lambda_module(self, lambda_function_files):
         """Helper to load the Lambda module. Must be called within mock_aws context."""
         _, temp_path, _ = lambda_function_files
         spec = importlib.util.spec_from_file_location("lambda_module", temp_path)
-        lambda_module = importlib.util.module_from_spec(spec)        
+        lambda_module = importlib.util.module_from_spec(spec)
         spec.loader.exec_module(lambda_module)
         return lambda_module
 
@@ -52,7 +52,32 @@ def create_instance(self, ec2_client, image_id: str = 'ami-12345678', instance_t
             InstanceType=instance_type
         )
         return response['Instances'][0]['InstanceId']
-    
+
+    def _get_launch_time_naive(self, ec2_client, instance_id: str) -> datetime:
+        response = ec2_client.describe_instances(InstanceIds=[instance_id])
+        instance = response['Reservations'][0]['Instances'][0]
+        launch_time = instance.get('LaunchTime')
+        if launch_time is None:
+            raise AssertionError("LaunchTime is None in test setup; cannot construct non-stale shutdown tag.")
+        if launch_time.tzinfo is not None:
+            return launch_time.replace(tzinfo=None)
+        return launch_time
+
+    def _freeze_lambda_now(self, monkeypatch, lambda_module, frozen_now_naive: datetime):
+        """
+        Freeze lambda_module.datetime.now(...) to return frozen_now_naive.
+        The Lambda imports `datetime` directly (`from datetime import datetime, ...`),
+        so we patch the module attribute `datetime`.
+        """
+        class FrozenDatetime(datetime):
+            @classmethod
+            def now(cls, tz=None):
+                if tz is None:
+                    return frozen_now_naive
+                return frozen_now_naive.replace(tzinfo=tz)
+
+        monkeypatch.setattr(lambda_module, "datetime", FrozenDatetime)
+
     @mock_aws
     def test_handler_with_non_existent_instance(self, capsys, lambda_function_files):
         """Test handler behavior with non-existent instance ID"""
@@ -97,6 +122,40 @@ def test_handler_with_stopped_instance(self, capsys, lambda_function_files):
         assert "skipping any action" in captured.out
         assert "(current state:" in captured.out
 
+    @mock_aws
+    def test_handler_removes_tag_from_stopped_instance(self, capsys, lambda_function_files):
+        """Test handler removes shutdown tag from a stopped instance that has one"""
+        lambda_module = self.load_lambda_module(lambda_function_files)
+
+        # Create and stop an instance
+        ec2 = boto3.client('ec2', region_name='us-east-1')
+        instance_id = self.create_instance(ec2)
+
+        # Add a shutdown tag before stopping
+        future_time = datetime.now(UTC).replace(tzinfo=None) + timedelta(hours=2)
+        tag_value = future_time.strftime('%a, %d %b %Y %H:%M:%S GMT')
+        ec2.create_tags(
+            Resources=[instance_id],
+            Tags=[{'Key': 'mw-autoshutdown', 'Value': tag_value}]
+        )
+
+        ec2.stop_instances(InstanceIds=[instance_id])
+
+        # Set environment variables
+        self.set_env_vars(instance_id, "After 2 hours")
+
+        # Execute handler
+        lambda_module.handler(event=None, context=None)
+
+        # Verify tag was removed
+        captured = capsys.readouterr()
+        response = ec2.describe_instances(InstanceIds=[instance_id])
+        instance = response['Reservations'][0]['Instances'][0]
+        tags = {tag['Key']: tag['Value'] for tag in instance.get('Tags', [])}
+
+        assert os.environ['TAG_TO_MONITOR'] not in tags
+        assert f"Removed {os.environ['TAG_TO_MONITOR']} tag from EC2 instance {instance_id}" in captured.out
+
     @mock_aws
     def test_handler_with_terminated_instance(self, capsys, lambda_function_files):
         """Test handler behavior with terminated instance"""
@@ -193,6 +252,9 @@ def test_handler_respects_future_shutdown_tag(self, capsys, lambda_function_file
         # Create instance
         ec2 = boto3.client('ec2', region_name='us-east-1')
         instance_id = self.create_instance(ec2)
+
+        # Set environment variables (ensures TAG_TO_MONITOR exists deterministically)
+        self.set_env_vars(instance_id, "After 1 hour")
 
         # Add future shutdown tag
         future_time = datetime.now(UTC).replace(tzinfo=None) + timedelta(hours=3)
@@ -202,9 +264,6 @@ def test_handler_respects_future_shutdown_tag(self, capsys, lambda_function_file
             Tags=[{'Key': os.environ['TAG_TO_MONITOR'], 'Value': tag_value}]
         )
 
-        # Set environment variables
-        self.set_env_vars(instance_id, "After 1 hour")
-        
         # Execute handler
         lambda_module.handler(event=None, context=None)
 
@@ -221,23 +280,30 @@ def test_handler_respects_future_shutdown_tag(self, capsys, lambda_function_file
         assert "Remaining minutes:" in captured.out
 
     @mock_aws
-    def test_handler_stops_instance_when_time_reached(self, capsys, lambda_function_files):
+    def test_handler_stops_instance_when_time_reached(self, capsys, monkeypatch, lambda_function_files):
         """Test handler stops instance when shutdown time is reached"""
         lambda_module = self.load_lambda_module(lambda_function_files)
 
         # Create instance
         ec2 = boto3.client('ec2', region_name='us-east-1')
         instance_id = self.create_instance(ec2)
 
-        # Add past shutdown tag (should trigger shutdown)
-        past_time = "Thu, 01 Jan 1970 00:00:00 GMT"
+        # Set environment variables FIRST (deterministic TAG_TO_MONITOR)
+        self.set_env_vars(instance_id, "After 2 hours")
+
+        # Construct a non-stale shutdown time: AFTER LaunchTime but BEFORE "now" (frozen)
+        launch_time = self._get_launch_time_naive(ec2, instance_id)
+        shutdown_time = launch_time + timedelta(seconds=5)
+        tag_value = shutdown_time.strftime('%a, %d %b %Y %H:%M:%S GMT')
+
         ec2.create_tags(
             Resources=[instance_id],
-            Tags=[{'Key': os.environ['TAG_TO_MONITOR'], 'Value': past_time}]
+            Tags=[{'Key': os.environ['TAG_TO_MONITOR'], 'Value': tag_value}]
         )
 
-        # Set environment variables
-        self.set_env_vars(instance_id, "After 2 hours")
+        # Freeze lambda "now" to after shutdown_time to trigger shutdown
+        frozen_now = shutdown_time + timedelta(minutes=1)
+        self._freeze_lambda_now(monkeypatch, lambda_module, frozen_now)
 
         # Execute handler
         lambda_module.handler(event=None, context=None)
@@ -257,7 +323,7 @@ def test_handler_stops_instance_when_time_reached(self, capsys, lambda_function_
         assert f"The tag {os.environ['TAG_TO_MONITOR']} has been removed from EC2 instance {instance_id}" in captured.out
 
     @mock_aws
-    def test_handler_with_custom_tag_name(self, capsys, lambda_function_files):
+    def test_handler_with_custom_tag_name(self, capsys, monkeypatch, lambda_function_files):
         """Test handler with custom tag name (not the default mw-autoshutdown)"""
         lambda_module = self.load_lambda_module(lambda_function_files)
 
@@ -274,20 +340,27 @@ def test_handler_with_custom_tag_name(self, capsys, lambda_function_files):
 
         # Verify custom tag was added
         response = ec2.describe_instances(InstanceIds=[instance_id])
-        tags = {tag['Key']: tag['Value'] 
+        tags = {tag['Key']: tag['Value']
                 for tag in response['Reservations'][0]['Instances'][0].get('Tags', [])}
 
         assert custom_tag in tags
 
-        # Now add past time to trigger shutdown
+        # Now force a due shutdown tag that is NOT stale:
+        # shutdown_time must be after LaunchTime, and we freeze now after it.
+        launch_time = self._get_launch_time_naive(ec2, instance_id)
+        shutdown_time = launch_time + timedelta(seconds=5)
         ec2.create_tags(
             Resources=[instance_id],
-            Tags=[{'Key': custom_tag, 'Value': 'Thu, 01 Jan 1970 00:00:00 GMT'}]
+            Tags=[{'Key': custom_tag, 'Value': shutdown_time.strftime('%a, %d %b %Y %H:%M:%S GMT')}]
         )
 
         # Clear output
         capsys.readouterr()
 
+        # Freeze lambda "now" so shutdown triggers
+        frozen_now = shutdown_time + timedelta(minutes=1)
+        self._freeze_lambda_now(monkeypatch, lambda_module, frozen_now)
+        
         # Execute handler again - should stop and remove custom tag
         lambda_module.handler(event=None, context=None)
 
@@ -327,7 +400,7 @@ def test_handler_with_different_hour_values(self, capsys, config_string, expecte
 
         # Verify tag was set with correct time
         response = ec2.describe_instances(InstanceIds=[instance_id])
-        tags = {tag['Key']: tag['Value'] 
+        tags = {tag['Key']: tag['Value']
                 for tag in response['Reservations'][0]['Instances'][0].get('Tags', [])}
 
         assert os.environ['TAG_TO_MONITOR'] in tags, f"Tag should be added for {config_string}"
@@ -347,27 +420,36 @@ def test_handler_with_different_hour_values(self, capsys, config_string, expecte
         assert "minutes" in captured.out
 
     @mock_aws
-    def test_handler_preserves_other_tags(self, capsys, lambda_function_files):
+    def test_handler_preserves_other_tags(self, capsys, monkeypatch, lambda_function_files):
         """Test handler preserves other tags when stopping instance"""
         lambda_module = self.load_lambda_module(lambda_function_files)
 
         # Create instance with multiple tags
         ec2 = boto3.client('ec2', region_name='us-east-1')
         instance_id = self.create_instance(ec2)
 
-        # Add multiple tags including shutdown tag with past time
+        # Set environment variables FIRST (deterministic TAG_TO_MONITOR)
+        self.set_env_vars(instance_id, "After 2 hours")
+
+        # Construct a due, non-stale shutdown time
+        launch_time = self._get_launch_time_naive(ec2, instance_id)
+        shutdown_time = launch_time + timedelta(seconds=5)
+        tag_value = shutdown_time.strftime('%a, %d %b %Y %H:%M:%S GMT')
+
+        # Add multiple tags including shutdown tag
         ec2.create_tags(
             Resources=[instance_id],
             Tags=[
-                {'Key': os.environ['TAG_TO_MONITOR'], 'Value': 'Thu, 01 Jan 1970 00:00:00 GMT'},
+                {'Key': os.environ['TAG_TO_MONITOR'], 'Value': tag_value},
                 {'Key': 'Name', 'Value': 'TestInstance'},
                 {'Key': 'Environment', 'Value': 'Testing'},
                 {'Key': 'Owner', 'Value': 'TestUser'}
             ]
         )
 
-        # Set environment variables
-        self.set_env_vars(instance_id, "After 2 hours")
+        # Freeze lambda "now" to after shutdown_time
+        frozen_now = shutdown_time + timedelta(minutes=1)
+        self._freeze_lambda_now(monkeypatch, lambda_module, frozen_now)
 
         # Execute handler - should stop and remove only shutdown tag
         lambda_module.handler(event=None, context=None)
@@ -384,23 +466,28 @@ def test_handler_preserves_other_tags(self, capsys, lambda_function_files):
         assert tags.get('Owner') == 'TestUser'
 
     @mock_aws
-    def test_handler_respects_never_with_existing_tag(self, capsys, lambda_function_files):
+    def test_handler_respects_never_with_existing_tag(self, capsys, monkeypatch, lambda_function_files):
         """Test that 'Never' shutdown behavior still checks existing tags"""
         lambda_module = self.load_lambda_module(lambda_function_files)
 
         # Create instance
         ec2 = boto3.client('ec2', region_name='us-east-1')
         instance_id = self.create_instance(ec2)
 
-        # Add a past shutdown tag
-        past_time = "Thu, 01 Jan 1970 00:00:00 GMT"
+        # Set environment variables FIRST (deterministic TAG_TO_MONITOR)
+        self.set_env_vars(instance_id, "Never")
+        
+        # Add a due, non-stale shutdown tag
+        launch_time = self._get_launch_time_naive(ec2, instance_id)
+        shutdown_time = launch_time + timedelta(seconds=5)
         ec2.create_tags(
             Resources=[instance_id],
-            Tags=[{'Key': os.environ['TAG_TO_MONITOR'], 'Value': past_time}]
+            Tags=[{'Key': os.environ['TAG_TO_MONITOR'], 'Value': shutdown_time.strftime('%a, %d %b %Y %H:%M:%S GMT')}]
         )
 
-        # Set environment variables with Never - but existing tag should still be processed
-        self.set_env_vars(instance_id, "Never")
+        # Freeze lambda "now" so shutdown triggers
+        frozen_now = shutdown_time + timedelta(minutes=1)
+        self._freeze_lambda_now(monkeypatch, lambda_module, frozen_now)
 
         # Execute handler - should still stop because tag exists
         lambda_module.handler(event=None, context=None)
@@ -416,5 +503,45 @@ def test_handler_respects_never_with_existing_tag(self, capsys, lambda_function_
         assert os.environ['TAG_TO_MONITOR'] not in tags
         assert f"Stopping EC2 instance {instance_id}" in captured.out
 
+    @mock_aws
+    def test_handler_ignores_stale_shutdown_tag(self, capsys, lambda_function_files):
+        """
+        Test handler ignores a stale shutdown tag.
+
+        A stale tag is one where the shutdown_time is <= instance LaunchTime,
+        which can happen if the instance was restarted after a previously scheduled shutdown time.
+        In this case, the Lambda should take no action and must not stop the instance.
+        """
+        lambda_module = self.load_lambda_module(lambda_function_files)
+
+        # Create instance
+        ec2 = boto3.client('ec2', region_name='us-east-1')
+        instance_id = self.create_instance(ec2)
+
+        # Set environment variables FIRST (deterministic TAG_TO_MONITOR)
+        self.set_env_vars(instance_id, "After 2 hours")
+
+        # Create a stale shutdown tag: shutdown_time BEFORE LaunchTime
+        launch_time = self._get_launch_time_naive(ec2, instance_id)
+        stale_shutdown_time = launch_time - timedelta(minutes=1)
+        ec2.create_tags(
+            Resources=[instance_id],
+            Tags=[{
+                'Key': os.environ['TAG_TO_MONITOR'],
+                'Value': stale_shutdown_time.strftime('%a, %d %b %Y %H:%M:%S GMT')
+            }]
+        )
+
+        # Execute handler - should return early (no action)
+        lambda_module.handler(event=None, context=None)
+
+        response = ec2.describe_instances(InstanceIds=[instance_id])
+        instance = response['Reservations'][0]['Instances'][0]
+        tags = {tag['Key']: tag['Value'] for tag in instance.get('Tags', [])}
+
+        # Instance should still be running and stale tag should not be deleted
+        assert instance['State']['Name'] == 'running'
+        assert os.environ['TAG_TO_MONITOR'] in tags
+
 if __name__ == "__main__":
     pytest.main([__file__, '-sv'])