Bugfix/audit-logs checkpoints (#169)

* added audit-logs command

* remove format option

* Added send-to command

* added changelog

* fix test

* changed argument naming convention to singular

* refactor search intervals

* added format options

* refactor

* make begin date required

* fix tests

* added message when no results found and default header for table output

* fix send-to when no results found

* toggle py42 version

* Fix integration test, removed output validation as made input dates dynamic

* Add integration test for auditlogs

* Fix style

* added docs string and renamed variable

* refactor: _send_to function

* Rectify doc

* Changed output fields for table format

* add checkpointing to audit-logs commands

* sort audit-log events

* fix style failures

* fix tests

* fix moar tests

* re-implement checkpointing with deduping by event hashes (which are now also stored in cursor)

* use parent class methods since they're the same

* update tests, fix style

* fix missing mock logger

* add docstring to dedupe function

* style

* add auditlog cursor to profile delete test

* make response fixtures generators to match actual get_all behavior

* clarify use_checkpoint value is checkpoint name not a boolean

* add TestAuditLogCursorStore class

* no need to return when writing

* fix clear_checkpoint helptext

* style

* fix "timestamp w/ missing ms" bug and add test

* style

Co-authored-by: Kiran Chaudhary <kiran.chaudhary@code42.com>

Author: timabrmsn

src/code42cli/cmds/auditlogs.py

@@ -1,9 +1,12 @@
-import json
-from _collections import OrderedDict
+from collections import OrderedDict
+from datetime import datetime
+from datetime import timezone
 
 import click
 
 from code42cli.click_ext.groups import OrderedGroup
+from code42cli.cmds.search.cursor_store import AuditLogCursorStore
+from code42cli.cmds.search.options import BeginOption
 from code42cli.date_helper import parse_max_timestamp
 from code42cli.date_helper import parse_min_timestamp
 from code42cli.logger import get_logger_for_server
@@ -14,10 +17,12 @@
 from code42cli.options import send_to_format_options
 from code42cli.options import server_options
 from code42cli.output_formats import OutputFormatter
+from code42cli.util import hash_event
 from code42cli.util import warn_interrupt
 
 EVENT_KEY = "events"
 AUDIT_LOGS_KEYWORD = "audit-logs"
+AUDIT_LOG_TIMESTAMP_FORMAT = "%Y-%m-%dT%H:%M:%S.%f"
 
 AUDIT_LOGS_DEFAULT_HEADER = OrderedDict()
 AUDIT_LOGS_DEFAULT_HEADER["timestamp"] = "Timestamp"
@@ -26,8 +31,7 @@
 AUDIT_LOGS_DEFAULT_HEADER["actorIpAddress"] = "ActorIpAddress"
 AUDIT_LOGS_DEFAULT_HEADER["userName"] = "AffectedUser"
 AUDIT_LOGS_DEFAULT_HEADER["userId"] = "AffectedUserUID"
-# AUDIT_LOGS_DEFAULT_HEADER["success"] = "Success"
-# AUDIT_LOGS_DEFAULT_HEADER["resultCount"] = "ResultCount"
+
 
 filter_option_usernames = click.option(
     "--username", required=False, help="Filter results by usernames.", multiple=True,
@@ -67,7 +71,7 @@ def filter_options(f):
         f,
         AUDIT_LOGS_KEYWORD,
         callback=lambda ctx, param, arg: parse_min_timestamp(arg),
-        required=True,
+        cls=BeginOption,
     )
     f = end_option(
         f, AUDIT_LOGS_KEYWORD, callback=lambda ctx, param, arg: parse_max_timestamp(arg)
@@ -81,16 +85,34 @@ def filter_options(f):
     return f
 
 
+checkpoint_option = click.option(
+    "-c",
+    "--use-checkpoint",
+    metavar="checkpoint",
+    help="Only get audit-log events that were not previously retrieved.",
+)
+
+
 @click.group(cls=OrderedGroup)
 @sdk_options(hidden=True)
 def audit_logs(state):
     """Retrieve audit logs."""
-    pass
+    # store cursor getter on the group state so shared --begin option can use it in validation
+    state.cursor_getter = _get_audit_log_cursor_store
+
+
+@audit_logs.command()
+@click.argument("checkpoint-name")
+@sdk_options()
+def clear_checkpoint(state, checkpoint_name):
+    """Remove the saved audit log checkpoint from `--use-checkpoint/-c` mode."""
+    _get_audit_log_cursor_store(state.profile.name).delete(checkpoint_name)
 
 
 @audit_logs.command()
 @filter_options
 @format_option
+@checkpoint_option
 @sdk_options()
 def search(
     state,
@@ -103,11 +125,19 @@ def search(
     affected_user_id,
     affected_username,
     format,
+    use_checkpoint,
 ):
     """Search audit logs."""
-    _search(
+    formatter = OutputFormatter(format, AUDIT_LOGS_DEFAULT_HEADER)
+    cursor = _get_audit_log_cursor_store(state.profile.name)
+    if use_checkpoint:
+        checkpoint_name = use_checkpoint
+        checkpoint = cursor.get(checkpoint_name)
+        if checkpoint is not None:
+            begin = checkpoint
+
+    events = _get_all_audit_log_events(
         state.sdk,
-        format,
         begin_time=begin,
         end_time=end,
         event_types=event_type,
@@ -117,10 +147,25 @@ def search(
         affected_user_ids=affected_user_id,
         affected_usernames=affected_username,
     )
+    if use_checkpoint:
+        checkpoint_name = use_checkpoint
+        events = list(
+            _dedupe_checkpointed_events_and_store_updated_checkpoint(
+                cursor, checkpoint_name, events
+            )
+        )
+    if not events:
+        click.echo("No results found.")
+        return
+    elif len(events) > 10:
+        click.echo_via_pager(formatter.get_formatted_output(events))
+    else:
+        formatter.echo_formatted_list(events)
 
 
 @audit_logs.command()
 @filter_options
+@checkpoint_option
 @server_options
 @send_to_format_options
 @sdk_options()
@@ -137,13 +182,19 @@ def send_to(
     user_ip,
     affected_user_id,
     affected_username,
+    use_checkpoint,
 ):
     """Send audit logs to the given server address."""
-    _send_to(
+    logger = get_logger_for_server(hostname, protocol, format)
+    cursor = _get_audit_log_cursor_store(state.profile.name)
+    if use_checkpoint:
+        checkpoint_name = use_checkpoint
+        checkpoint = cursor.get(checkpoint_name)
+        if checkpoint is not None:
+            begin = checkpoint
+
+    events = _get_all_audit_log_events(
         state.sdk,
-        hostname,
-        protocol,
-        format,
         begin_time=begin,
         end_time=end,
         event_types=event_type,
@@ -153,43 +204,81 @@ def send_to(
         affected_user_ids=affected_user_id,
         affected_usernames=affected_username,
     )
+    if use_checkpoint:
+        checkpoint_name = use_checkpoint
+        events = list(
+            _dedupe_checkpointed_events_and_store_updated_checkpoint(
+                cursor, checkpoint_name, events
+            )
+        )
+    with warn_interrupt():
+        event = None
+        for event in events:
+            logger.info(event)
+        if event is None:  # generator was empty
+            click.echo("No results found.")
 
 
-def _search(sdk, format, **filter_args):
-
-    formatter = OutputFormatter(format, AUDIT_LOGS_DEFAULT_HEADER)
+def _get_all_audit_log_events(sdk, **filter_args):
     response_gen = sdk.auditlogs.get_all(**filter_args)
-
     events = []
     try:
-        for response in response_gen:
-            response_dict = json.loads(response.text)
-            if EVENT_KEY in response_dict:
-                events.extend(response_dict.get(EVENT_KEY))
+        responses = list(response_gen)
     except KeyError:
         # API endpoint (get_page) returns a response without events key when no records are found
         # e.g {"paginationRangeStartIndex": 10000, "paginationRangeEndIndex": 10000, "totalResultCount": 1593}
-        pass
+        # we can remove this check once PL-93211 is resolved and deployed.
+        return events
 
-    event_count = len(events)
-    if not event_count:
-        click.echo("No results found.")
-    elif event_count > 10:
-        click.echo_via_pager(formatter.get_formatted_output(events))
-    else:
-        formatter.echo_formatted_list(events)
+    for response in responses:
+        if EVENT_KEY in response.data:
+            response_events = response.data.get(EVENT_KEY)
+            events.extend(response_events)
 
+    return sorted(events, key=lambda x: x.get("timestamp"))
 
-def _send_to(sdk, hostname, protocol, format, **filter_args):
-    logger = get_logger_for_server(hostname, protocol, format)
-    with warn_interrupt():
-        response_gen = sdk.auditlogs.get_all(**filter_args)
-        try:
-            for response in response_gen:
-                if EVENT_KEY in response:
-                    for event in response[EVENT_KEY]:
-                        logger.info(event)
-                else:
-                    logger.info("No results found.")
-        except KeyError:
-            pass
+
+def _dedupe_checkpointed_events_and_store_updated_checkpoint(
+    cursor, checkpoint_name, events
+):
+    """De-duplicates events across checkpointed runs. Since using the timestamp of the last event
+    processed as the `--begin` time of the next run causes the last event to show up again in the
+    next results, we hash the last event(s) of each run and store those hashes in the cursor to
+    filter out on the next run. It's also possible that two events have the exact same timestamp, so
+    `checkpoint_events` needs to be a list of hashes so we can filter out everything that's actually
+    been processed.
+    """
+
+    checkpoint_events = cursor.get_events(checkpoint_name)
+    new_timestamp = None
+    new_events = []
+    for event in events:
+        event_hash = hash_event(event)
+        if event_hash not in checkpoint_events:
+            if event["timestamp"] != new_timestamp:
+                new_timestamp = event["timestamp"]
+                new_events.clear()
+            new_events.append(event_hash)
+            yield event
+            ts = _parse_audit_log_timestamp_string_to_timestamp(new_timestamp)
+            cursor.replace(checkpoint_name, ts)
+            cursor.replace_events(checkpoint_name, new_events)
+
+
+def _get_audit_log_cursor_store(profile_name):
+    return AuditLogCursorStore(profile_name)
+
+
+def _parse_audit_log_timestamp_string_to_timestamp(ts):
+    # example: {"property": "bar", "timestamp": "2020-11-23T17:13:26.239647Z"}
+    ts = ts[:-1]
+    try:
+        dt = datetime.strptime(ts, AUDIT_LOG_TIMESTAMP_FORMAT).replace(
+            tzinfo=timezone.utc
+        )
+    except ValueError:
+        ts = ts + ".0"  # handle timestamps that are missing ms
+        dt = datetime.strptime(ts, AUDIT_LOG_TIMESTAMP_FORMAT).replace(
+            tzinfo=timezone.utc
+        )
+    return dt.timestamp()

src/code42cli/cmds/search/cursor_store.py

@@ -1,3 +1,4 @@
+import json
 import os
 from os import path
 
@@ -37,7 +38,7 @@ def replace(self, cursor_name, new_timestamp):
         """Replaces the last stored date observed timestamp with the given one."""
         location = path.join(self._dir_path, cursor_name)
         with open(location, "w") as checkpoint:
-            return checkpoint.write(str(new_timestamp))
+            checkpoint.write(str(new_timestamp))
 
     def delete(self, cursor_name):
         """Removes a single cursor from the store."""
@@ -75,13 +76,31 @@ def __init__(self, profile_name):
         super().__init__(dir_path)
 
 
-def get_file_event_cursor_store(profile_name):
-    return FileEventCursorStore(profile_name)
+class AuditLogCursorStore(BaseCursorStore):
+    def __init__(self, profile_name):
+        dir_path = get_user_project_path("audit_log_checkpoints", profile_name)
+        super().__init__(dir_path)
 
+    def get_events(self, cursor_name):
+        try:
+            location = path.join(self._dir_path, cursor_name) + "_events"
+            with open(location) as checkpoint:
+                try:
+                    return json.loads(checkpoint.read())
+                except json.JSONDecodeError:
+                    return []
+        except FileNotFoundError:
+            return []
 
-def get_alert_cursor_store(profile_name):
-    return AlertCursorStore(profile_name)
+    def replace_events(self, cursor_name, new_events):
+        location = path.join(self._dir_path, cursor_name) + "_events"
+        with open(location, "w") as checkpoint:
+            checkpoint.write(json.dumps(new_events))
 
 
 def get_all_cursor_stores_for_profile(profile_name):
-    return [FileEventCursorStore(profile_name), AlertCursorStore(profile_name)]
+    return [
+        FileEventCursorStore(profile_name),
+        AlertCursorStore(profile_name),
+        AuditLogCursorStore(profile_name),
+    ]

src/code42cli/util.py

@@ -1,7 +1,9 @@
+import json
 import os
 import shutil
 from collections import OrderedDict
 from functools import wraps
+from hashlib import md5
 from os import path
 from signal import getsignal
 from signal import SIGINT
@@ -167,3 +169,9 @@ def _get_default_header(header_items):
             if key not in header and isinstance(key, str):
                 header[key] = key
     return header
+
+
+def hash_event(event):
+    if isinstance(event, dict):
+        event = json.dumps(event, sort_keys=True)
+    return md5(event.encode()).hexdigest()