From 1f0ac7de4e56416add751e0b557de8841a38e890 Mon Sep 17 00:00:00 2001
From: Garrett Marconet <41752849+g-marconet@users.noreply.github.com>
Date: Thu, 6 Feb 2025 09:24:11 -0800
Subject: [PATCH] Fix `is_valid_timestamp` for strings with timezone (#417)

* fix timestamp_format timezone

* add tests for is_valid_timestamp

* ruff format

* make comments consistent

* update generated timezone comment

* delete utc blurb
---
 .../layers/heimdall_utils/heimdall_utils/utils.py    |  9 ++++-----
 orchestrator/lambdas/tests/test_utils.py             | 12 ++++++++++++
 2 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/orchestrator/lambdas/layers/heimdall_utils/heimdall_utils/utils.py b/orchestrator/lambdas/layers/heimdall_utils/heimdall_utils/utils.py
index 69c8b660..522d86cf 100644
--- a/orchestrator/lambdas/layers/heimdall_utils/heimdall_utils/utils.py
+++ b/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 --
@@ -102,14 +102,13 @@ def parse_timestamp(timestamp: Optional[str] = None) -> str:
     Examples:
         >>> parse_timestamp()
         {"level":"DEBUG","location":"parse_timestamp","message":"Generating Default timestamp"}
-        "2024-04-24T22:35:36Z"  # Output will vary based on the current date and time
+        "2024-01-01T6:14:57-0800"  # Output will vary based on the current date, time, and timezone
 
         >>> parse_timestamp("2024-06-24T22:50:00Z")
         "2024-06-24T22:50:00Z"
 
     Notes:
         - The function uses the current system time to calculate the 3-month offset.
-        - All returned timestamps are in UTC (denoted by the 'Z' suffix).
     """
     if timestamp and is_valid_timestamp(timestamp):
         return timestamp
@@ -133,6 +132,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
diff --git a/orchestrator/lambdas/tests/test_utils.py b/orchestrator/lambdas/tests/test_utils.py
index 5458c144..c8efe0cc 100644
--- a/orchestrator/lambdas/tests/test_utils.py
+++ b/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))