Fix is_valid_timestamp for strings with timezone

orchestrator/lambdas/layers/heimdall_utils/heimdall_utils/utils.py

@@ -3,7 +3,7 @@
 import logging
 import time
 from collections import namedtuple
-from datetime import datetime, timedelta
+from datetime import datetime, timedelta, timezone
 from json import JSONDecodeError
 from typing import Optional
 
@@ -17,7 +17,7 @@
 logging.getLogger("botocore").setLevel(logging.WARNING)
 
 DYNAMODB_TTL_DAYS = 60
-TIMESTAMP_FORMAT = "%Y-%m-%dT%H:%M:%SZ"
+TIMESTAMP_FORMAT = "%Y-%m-%dT%H:%M:%S%z"
 
 
 # -- This constant represents the maximum age for a HEIMDALL Scan --
@@ -133,6 +133,6 @@ def is_valid_timestamp(timestamp: str) -> bool:
 
 
 def get_datetime_from_past(days: int) -> datetime:
-    current_timestamp = datetime.now()
+    current_timestamp = datetime.now(timezone.utc)
     three_months_ago = current_timestamp - timedelta(days=days)
     return three_months_ago

orchestrator/lambdas/tests/test_utils.py

@@ -26,3 +26,15 @@ def test_get_json_from_response_success(self):
         result = self.json_utils.get_json_from_response(json.dumps(test_dict))
 
         self.assertEqual(test_dict, result)
+
+    def test_is_valid_timestamp_utc(self):
+        # Checks relative to current time minus a max scan age.
+        # Make sure to change this before 9998
+        time = "9999-01-01T00:00:00Z"
+        self.assertTrue(utils.is_valid_timestamp(time))
+
+    def test_is_valid_timestamp_localized(self):
+        # Checks relative to current time minus a max scan age.
+        # Make sure to change this before 9998
+        time = "9999-01-01T00:00:00-08:00"
+        self.assertTrue(utils.is_valid_timestamp(time))