netskope_log_fetcher.py

"""Main file. Run to pull down Netskope logs."""


from datetime import datetime
import logging
import json
import os
import re

from dotenv import load_dotenv

from netskope_fetcher.bootstrap import NetskopeAsyncBootstrap
from netskope_fetcher.token import Token
from netskope_fetcher.events import EventClient
from netskope_fetcher.alerts import AlertClient
from netskope_fetcher.logger import setup_logger


class TinyTimeWriter:
    """Manages writing / reading of the timestamp file."""

    def __init__(self, file_path=None):
        self.time_file_path = file_path or os.path.join(
            os.path.dirname(__file__), "time.log"
        )

    def save_last_log_time(self, _time):
        """
        Saves the end-time defined for the current run of the program.
        Parameters
        ----------
        format_: str
            Format you wish the string output to be in
        Returns
        ----------
        datetime string
            String format of the datetime object.
        """
        with open(self.time_file_path, "w") as _file:
            _file.write(str(_time))

    def get_last_log_time(self):
        """ Reads the last log time from file

        Returns
        ----------
        Epoch time stamp for last time the program ran.
        """
        try:
            with open(self.time_file_path, "r") as _file:
                time_stamp = int(_file.readline())
        except FileNotFoundError:
            # File doesn't exist yet
            return None
        except ValueError:
            # If non-int or empty file.
            return None
        else:
            if time_stamp <= 0:
                time_stamp = None
            return time_stamp


def write_logs(netskope_object):
    """ Writes logs to the type-specific log file.

        Pull the log files from netskope_object.log_dictionary, and
        write them to file.  If log_dictionary looks like this:
        {
            'application': [list of logs],
            'page': [list of logs],
        }
        Then this is an 'EventClient' object and we will write to two
        files like this:

        /file/path/to/logs/event/application.log
        and
        /file/path/to/logs/event/page.log

    Parameters
    ----------
    netskope_object: netskope_fetcher.events.EventClient
                     OR netskope_fetcher.alerts.AlertClient
        Object contains the log files in log_dictionary.
    """

    _current_directory = os.path.dirname(__file__)
    for type_, log_list in netskope_object.log_dictionary.items():
        # Some types have spaces, replace them with underscores
        file_ = replace_spaces(type_)
        # logs/alert or logs/event
        log_path = os.path.join("logs", netskope_object.endpoint_type)
        make_dir_if_needed(_current_directory, log_path)
        # Ex: base/file/path/logs/alert/type.log
        log_file = os.path.join(_current_directory, log_path, "{}.log".format(file_))

        with open(log_file, "a+") as _f:
            logging.debug("Writing to %s log file.", log_file)
            try:
                for log in log_list:
                    _f.write("{}\n".format(json.dumps(log)))
            except TypeError as _t:
                # Most likely that log_list is not an iterable
                logging.warning("Couldn't write logs for %s: %s", type_, {_t})


def make_dir_if_needed(current_dir, log_dir):
    """ Helper function to create approriate log directories if they
        don't already exist

    Parameters
    ----------
    current_dir: str
        /path/to/current/directory
    log_dir: str
        logs/event or logs/alert. Will be appended to current_dir
    """

    required_dir = os.path.join(current_dir, log_dir)
    if not os.path.isdir(required_dir):
        os.mkdir(required_dir)


def replace_spaces(some_string):
    """ Substitute spaces with underscores"""

    return re.sub(" ", "_", some_string)


if __name__ == "__main__":
    try:
        setup_logger()

        CURRENT_DIRECTORY = os.path.dirname(__file__)
        load_dotenv(dotenv_path=os.path.join(CURRENT_DIRECTORY, ".env"))

        TINY_TIME = TinyTimeWriter()
        TOKEN = Token()

        # End time will always be 'right now'
        # start_time will be the end of the last successful run, or in the
        #    case that the last timestamp isn't available, set the start
        #    time to ten minutes ago.
        END_TIME = int(datetime.now().timestamp())
        START_TIME = TINY_TIME.get_last_log_time() or (END_TIME - 600)

        logging.info(
            "Running from %s to %s",
            datetime.strftime(datetime.fromtimestamp(START_TIME), "%c"),
            datetime.strftime(datetime.fromtimestamp(END_TIME), "%c"),
        )

        CLIENTS = [
            EventClient(token=TOKEN, start=START_TIME, end=END_TIME),
            AlertClient(token=TOKEN, start=START_TIME, end=END_TIME),
        ]

        BOOTSTRAP = NetskopeAsyncBootstrap(client_list=CLIENTS)
        BOOTSTRAP.run()

        # Write to the log files
        for client in CLIENTS:
            write_logs(client)

        # Save the end time so that it can be used in the next run.
        # This is purposely left at the end of the program so that the
        # subsequent run of the program will gather logs that may have been
        # missed if the script were to fail mid-stream.
        TINY_TIME.save_last_log_time(END_TIME)
    except Exception as _e:
        logging.exception("Exception Occurred: %s.", _e)
        raise

    exit()