Merge pull request #140 from code42/INTEG-2844/timestamp-bugfix

when querying sessions, properly use timestamp in milliseconds

Author: ceciliastevens

src/_incydr_sdk/queries/utils.py

@@ -41,6 +41,13 @@ def parse_ts_to_posix_ts(timestamp: Union[str, datetime]):
     return dt.timestamp()
 
 
+def parse_ts_to_ms_ts(timestamp: Union[str, datetime]):
+    """
+    Parse epoch ms timestamp from DATE/DATETIME str or datetime obj.
+    """
+    return parse_ts_to_posix_ts(timestamp) * 1000
+
+
 def parse_str_to_dt(timestamp: str):
     try:
         dt = parser.parse(timestamp)

src/_incydr_sdk/sessions/client.py

@@ -8,7 +8,7 @@
 from _incydr_sdk.enums.sessions import ContentInspectionStatuses
 from _incydr_sdk.enums.sessions import SessionStates
 from _incydr_sdk.enums.sessions import SortKeys
-from _incydr_sdk.queries.utils import parse_ts_to_posix_ts
+from _incydr_sdk.queries.utils import parse_ts_to_ms_ts
 from _incydr_sdk.sessions.models.models import SessionsChangeStateRequest
 from _incydr_sdk.sessions.models.models import SessionsCriteriaRequest
 from _incydr_sdk.sessions.models.models import SessionsQueryRequest
@@ -76,9 +76,9 @@ def get_page(
 
         # Parse timestamps
         if start_time and not isinstance(start_time, (int, float)):
-            start_time = parse_ts_to_posix_ts(start_time)
+            start_time = parse_ts_to_ms_ts(start_time)
         if end_time and not isinstance(end_time, (int, float)):
-            end_time = parse_ts_to_posix_ts(end_time)
+            end_time = parse_ts_to_ms_ts(end_time)
 
         if states and not isinstance(states, List):
             states = [states]
@@ -253,9 +253,9 @@ def update_state_by_criteria(
 
         # Parse timestamps
         if start_time and not isinstance(start_time, (int, float)):
-            start_time = parse_ts_to_posix_ts(start_time)
+            start_time = parse_ts_to_ms_ts(start_time)
         if end_time and not isinstance(end_time, (int, float)):
-            end_time = parse_ts_to_posix_ts(end_time)
+            end_time = parse_ts_to_ms_ts(end_time)
 
         if states and not isinstance(states, List):
             states = [states]

tests/test_sessions.py

@@ -22,7 +22,11 @@
 
 TEST_SESSION_ID = "123-session-1"
 DATETIME_INSTANT = datetime(2024, 1, 1, tzinfo=timezone.utc)
-POSIX_TS = int(DATETIME_INSTANT.timestamp())
+POSIX_TS = int(DATETIME_INSTANT.timestamp()) * 1000
+START_DATE = "2024-12-19"
+START_TIMESTAMP = 1734566400000  # in ms
+END_DATE = "2024-12-20"
+END_TIMESTAMP = 1734652800000  # in ms
 
 TEST_SESSION = {
     "actorId": TEST_SESSION_ID,
@@ -158,6 +162,50 @@ def test_get_page_when_custom_params_returns_expected_data(httpserver_auth: HTTP
     assert len(page.items) == 1 == page.total_count
 
 
+def test_get_page_when_given_date_uses_correct_timestamp(httpserver_auth: HTTPServer):
+    query = {
+        "actor_id": "actor-id",
+        "on_or_after": START_TIMESTAMP,
+        "before": END_TIMESTAMP,
+        "has_alerts": "false",
+        "risk_indicators": ["risk-indicator"],
+        "state": ["OPEN"],
+        "severity": [3],
+        "rule_id": ["rule-id"],
+        "watchlist_id": ["watchlist-id"],
+        "content_inspection_status": "PENDING",
+        "order_by": "score",
+        "sort_direction": "desc",
+        "page_number": 2,
+        "page_size": 10,
+    }
+    sessions_page = {"items": [TEST_SESSION], "totalCount": 1}
+    httpserver_auth.expect_request(
+        "/v1/sessions", method="GET", query_string=urlencode(query, doseq=True)
+    ).respond_with_json(sessions_page)
+
+    client = Client()
+    page = client.sessions.v1.get_page(
+        actor_id="actor-id",
+        start_time=START_DATE,
+        end_time=END_DATE,
+        has_alerts=False,
+        sort_key=SortKeys.SCORE,
+        risk_indicators=["risk-indicator"],
+        sort_dir=SortDirection.DESC,
+        states=SessionStates.OPEN,
+        severities=3,
+        rule_ids="rule-id",
+        watchlist_ids="watchlist-id",
+        page_num=2,
+        page_size=10,
+        content_inspection_status=ContentInspectionStatuses.PENDING,
+    )
+    assert isinstance(page, SessionsPage)
+    assert page.items[0].json() == json.dumps(TEST_SESSION)
+    assert len(page.items) == 1 == page.total_count
+
+
 def test_iter_all_when_default_params_returns_expected_data(
     httpserver_auth: HTTPServer,
 ):
@@ -357,6 +405,56 @@ def test_update_state_by_criteria_makes_expected_calls(httpserver_auth: HTTPServ
         assert response.status_code == 200
 
 
+def test_update_state_by_criteria_when_given_date_uses_correct_timestamp(
+    httpserver_auth: HTTPServer,
+):
+    query = {
+        "actor_id": "actor-id",
+        "on_or_after": START_TIMESTAMP,
+        "before": END_TIMESTAMP,
+        "has_alerts": "false",
+        "risk_indicators": ["risk-indicator"],
+        "state": ["OPEN"],
+        "severity": [3],
+        "rule_id": ["rule-id"],
+        "watchlist_id": ["watchlist-id"],
+        "content_inspection_status": "PENDING",
+    }
+
+    token = "123-token"
+    httpserver_auth.expect_request(
+        "/v1/sessions/change-states",
+        query_string=urlencode(query, doseq=True),
+        method="POST",
+        json={"continuationToken": None, "newState": "CLOSED"},
+    ).respond_with_json({"continuationToken": token})
+    httpserver_auth.expect_request(
+        "/v1/sessions/change-states",
+        query_string=urlencode(query, doseq=True),
+        method="POST",
+        json={"continuationToken": token, "newState": "CLOSED"},
+    ).respond_with_json({"continuationToken": None})
+
+    client = Client()
+    responses = client.sessions.v1.update_state_by_criteria(
+        new_state=SessionStates.CLOSED,
+        actor_id="actor-id",
+        start_time=START_DATE,
+        end_time=END_DATE,
+        has_alerts=False,
+        risk_indicators=["risk-indicator"],
+        states=SessionStates.OPEN,
+        severities=3,
+        rule_ids="rule-id",
+        watchlist_ids="watchlist-id",
+        content_inspection_status=ContentInspectionStatuses.PENDING,
+    )
+    assert responses[0].json()["continuationToken"] == token
+    assert responses[1].json()["continuationToken"] is None
+    for response in responses:
+        assert response.status_code == 200
+
+
 # ************************************************ CLI ************************************************