# Copyright (C) 2005-2025 Splunk Inc. All Rights Reserved.


from . import statestore
from time import sleep

from splunk.auth import getCurrentUser

from ITOA.setup_logging import logger

LOG_CHANGE_TRACKING = "[change_tracking]"


class ITOAStorage(object):
    """
    Defines a storage interface for each of the handlers.
    The different storage options will eventually be registered through a lookup file
    or state store, or some other persistent thingy
    """
    def __init__(self, **kwargs):
        """
        All simple initialization, for adding the storage options, read a file in the storage directory
        Registers the different options, and sets the default backend
        """
        self.backend = None
        self.app = self.get_app_name()
        self.itoa_storage_options = {}
        self.register_storage_option('statestore', statestore.StateStore(**kwargs))
        # Set the default primary option
        self.set_storage_backend('statestore')
        self.should_init = True

    def wait_for_storage_init(self, session_key, host_base_uri=''):
        """
        KV store can take long to get initialized on splunkd restart,
        use this method to wait until it is inited

        @param session_key: Session key to splunkd to use to check KV store init
        @type session_key: basestring

        @param host_base_uri: The base uri <scheme>://<host>:<port> of the target host.  '' targets local host.
        @type host_bast_uri: basestring

        @return: True if KV store is inited, else False
        @rtype: boolean
        """
        timeout_seconds = 300  # 5 Mins
        retry_seconds = 5
        while not self.is_available(session_key, host_base_uri=host_base_uri):
            logger.debug('KV store is not initialized, retrying after 5 seconds')
            sleep(retry_seconds)
            timeout_seconds -= retry_seconds
            if timeout_seconds == 0:
                logger.error("KV store does not seem to have been initialized"
                             " for the last 5 minutes. Stopping retry.")
                return False
        logger.debug('KV store has been initialized.')
        return True

    def get_app_name(self):
        """
        Gets the name of the app
        We used to be clever, and get it from the filesystem
        ...But clever was wrong
        """
        return "SA-ITOA"

    def register_storage_option(self, option, optioninstance):
        """
        Method for allowing custom storage options to be registered,
        should a customer or external party decide that this is desirable
        @param option: The new or existing option that they want
        @type option: string

        @param optioninstance: The interface to the option they want, duck typing used
        @type option: Anything as long as it supports the interface
        """
        self.itoa_storage_options[option] = optioninstance

    def get_storage_options(self):
        """
        Returns an array of the different options available for storing data
        @param self: The reference to self
        @type self: itoa_storage instance

        @return: An array of keys for supported backend storage options
        @rtype: list
        """
        return list(self.itoa_storage_options.keys())

    def get_backend(self, session_key, option=None, host_base_uri=''):
        """
        Gets the backend store and throws an exception if it doesnt exist

        @param option: An optional parameter that specifies which backend interface
        to use.  If blank, uses the current defined backend
        @type option: string

        @param host_base_uri: The base uri <scheme>://<host>:<port> of the target host.  '' targets local host.
        @type host_bast_uri: basestring

        @return: The backend interface
        @rtype: duck typed.
        """
        if option is None:
            option = self.backend
        elif option != self.backend:
            self.should_init = True
        backend = self.itoa_storage_options.get(option, None)
        if backend is None:
            raise Exception("Backend not registered or undefined")
        if self.should_init is True:
            backend.lazy_init(session_key, host_base_uri=host_base_uri)  # The backend should provide an initialize method
            self.backend = option
            self.should_init = False
        return backend

    def set_storage_backend(self, option):
        """
        Sets the option that should be used to store all of this crap

        @param option: The string index of the currently existing option to use
        @type string: string
        """
        self.backend = option

    def is_available(self, session_key, host_base_uri=''):
        """
        Tries to hit a non-existent collection endpoint.
        If the KV store is ready to serve requests, it will return a 500; if it's not up yet, it will
        return a 503.

        Will throw exception if the current backend is not KV store

        @param session_key: splunkd session key
        @type session_key: basestring

        @param host_base_uri: The base uri <scheme>://<host>:<port> of the target host.  '' targets local host.
        @type host_base_uri: basestring

        @returns true if response is not 503 else false
        @rtype bool
        """
        if self.backend != 'statestore':
            raise Exception("`is_available` method can only be run on the KV store backend.")
        backend = self.get_backend(session_key, host_base_uri=host_base_uri)
        retval = backend.is_available(session_key, host_base_uri=host_base_uri)
        logger.debug("Querying if KV store is available: %s", retval)
        return retval

    def check_payload_size(self, session_key, data_list, throw_on_violation=True, host_base_uri=''):
        """
        Method to verify payload size isnt larger than limit of size in backend

        @param session_key: splunkd session key
        @type: basestring

        @param data_list: JSON list payload to verify
        @type: list

        @param throw_on_violation: True if violation should trigger exception, else returns bool indicating
        @type: boolean
            if violation detected

        @param host_base_uri: The base uri <scheme>://<host>:<port> of the target host.  '' targets local host.
        @type host_base_uri: basestring

        @rtype: boolean
        @return: True if no violation detected, False if violation detected
        """
        return self.get_backend(session_key, host_base_uri=host_base_uri).check_payload_size(
            data_list, session_key=session_key, throw_on_violation=throw_on_violation)

###############################################################################
# Generic Crud methods
###############################################################################

    def create(self, session_key, owner, objecttype, data, current_user_name=None, host_base_uri=''):
        """
        A generic creation method, used by all of the other components
        to create a particular entity or service based on the json passed in
        @param session_key: The splunkd session key
        @type session_key: string

        @param owner: user who is performing this operation
        @type owner: basestring

        @param objecttype: name of objecttype being created
        @type objecttype: basestring

        @param data: data
        @type data: object

        @param current_user_name: user name of who is performing this operation
        @type current_user_name: basestring

        @param host_base_uri: The base uri <scheme>://<host>:<port> of the target host.  '' targets local host.
        @type host_base_uri: basestring

        @return a json structure containing an id field and an id
        @retval dict
        """
        backend = self.get_backend(session_key, host_base_uri=host_base_uri)
        # The backend will create an id for us to use
        # The we cannot assign it here, it must be on the lower levels
        user_name = current_user_name
        if user_name is None:
            user_name = getCurrentUser()['name']
        logger.info("%s user=%s method=create objecttype=%s attempt",
                    LOG_CHANGE_TRACKING, user_name, objecttype)
        self.entity_clarification(objecttype, data)
        retval = backend.create(
            session_key,
            owner,
            objecttype,
            data,
            host_base_uri=host_base_uri
        )
        logger.info("%s user=%s method=create objecttype=%s key=%s",
                    LOG_CHANGE_TRACKING, user_name, objecttype, retval.get("_key"))
        return retval

    def edit(self, session_key, owner, objecttype, identifier, data, current_user_name=None, host_base_uri=''):
        """
        A generic edit method, used by all of the other components to
        create a particular service or entity based on the json passed in
        @param session_key: The splunkd session key
        @type session_key: string

        @param owner: user who is performing this operation
        @type owner: basestring

        @param objecttype: name of objecttype being updated
        @type objecttype: basestring

        @param identifier: key of object being updated
        @type identifer: basestring

        @param data: data
        @type data: object

        @param current_user_name: user name of who is performing this operation
        @type current_user_name: basestring

        @param host_base_uri: The base uri <scheme>://<host>:<port> of the target host.  '' targets local host.
        @type host_base_uri: basestring

        @return a json structure containing an id field and an id
        @retval dict
        """
        backend = self.get_backend(session_key, host_base_uri=host_base_uri)
        user_name = current_user_name
        if user_name is None:
            user_name = getCurrentUser()['name']
        logger.info("%s user=%s method=edit objecttype=%s key=%s",
                    LOG_CHANGE_TRACKING, user_name, objecttype, identifier)
        self.entity_clarification(objecttype, data)
        retval = backend.edit(
            session_key,
            owner,
            objecttype,
            identifier,
            data,
            host_base_uri=host_base_uri
        )
        return retval

    def get(self, session_key, owner, objecttype, identifier, current_user_name=None, host_base_uri=''):
        """
        A generic get method to retrieve the item specified by
        its identifier.  This is only used to retrieve by identifier

        @param session_key: The splunkd session key
        @type session_key: string

        @param owner: user who is performing this operation
        @type owner: basestring

        @param objecttype: name of objecttype being fetched
        @type objecttype: basestring

        @param identifier: key of object being fetched
        @type identifer: basestring

        @param current_user_name: user name of who is performing this operation
        @type current_user_name: basestring

        @param host_base_uri: The base uri <scheme>://<host>:<port> of the target host.  '' targets local host.
        @type host_base_uri: basestring

        @return a json-like dict structure containing the fields of the requested item
        @retval dict
        """
        backend = self.get_backend(session_key, host_base_uri=host_base_uri)
        user_name = current_user_name
        if user_name is None:
            user_name = getCurrentUser()['name']
        logger.info("%s user=%s method=get objecttype=%s key=%s",
                    LOG_CHANGE_TRACKING, user_name, objecttype, identifier)
        return backend.get(
            session_key,
            owner,
            objecttype,
            identifier,
            host_base_uri=host_base_uri
        )

    def delete(self, session_key, owner, objecttype, identifier, current_user_name=None, host_base_uri=''):
        """
        A generic delete method used to delete
        @param session_key: The splunkd session key
        @type session_key: string

        @param owner: user who is performing this operation
        @type owner: basestring

        @param objecttype: name of objecttype being deleted
        @type objecttype: basestring

        @param identifier: key of object being deleted
        @type identifer: basestring

        @param current_user_name: user name of who is performing this operation
        @type current_user_name: basestring

        @param host_base_uri: The base uri <scheme>://<host>:<port> of the target host.  '' targets local host.
        @type host_base_uri: basestring
        """
        backend = self.get_backend(session_key, host_base_uri=host_base_uri)
        user_name = current_user_name
        if user_name is None:
            user_name = getCurrentUser()['name']
        logger.info("%s user=%s method=delete objecttype=%s key=%s",
                    LOG_CHANGE_TRACKING, user_name, objecttype, identifier)
        backend.delete(
            session_key,
            owner,
            objecttype,
            identifier,
            host_base_uri=host_base_uri
        )

    def get_all(self, session_key, owner, objecttype, sort_key=None, sort_dir=None,
                filter_data={}, fields=None, skip=None, limit=None, current_user_name=None,
                host_base_uri=''):
        """
        Get all of a particular thing, returned as a list of json structures
        @param session_key: The splunkd session key
        @type session_key: string

        @param owner: user who is performing this operation
        @type owner: basestring

        @param objecttype: name of objecttype being fetched
        @type objecttype: basestring

        @param sort_key: field to sort on
        @type sort_key: basestring

        @param sort_dir: direction of sort
        @type sort_dir: basestring

        @param filter_data: filter to apply to objects bein fetched
        @type filter_data: dict

        @param skip: skip
        @type skip: object

        @param limit: max count of objects being fetched
        @type limit: basestring

        @param current_user_name: user name of who is performing this operation
        @type current_user_name: basestring

        @param host_base_uri: The base uri <scheme>://<host>:<port> of the target host.  '' targets local host.
        @type host_base_uri: basestring

        @return: a dict suitable for json conversion of the object types
        @rtype: list
        """
        backend = self.get_backend(session_key, host_base_uri=host_base_uri)
        user_name = current_user_name
        if user_name is None:
            user_name = getCurrentUser()['name']
        logger.info("%s user=%s method=get_all objecttype=%s filter=%s sort_key=%s sort_dir=%s limit=%s skip=%s",
                    LOG_CHANGE_TRACKING, user_name, objecttype, filter_data, sort_key, sort_dir, limit, skip)
        return backend.get_all(
            session_key,
            owner,
            objecttype,
            sort_key=sort_key,
            sort_dir=sort_dir,
            filter_data=filter_data,
            fields=fields,
            skip=skip,
            limit=limit,
            host_base_uri=host_base_uri
        )

    def delete_all(self, session_key, owner, objecttype, filter_data, current_user_name=None, host_base_uri=''):
        """
        Delete all of a particular thing, no return value
        @param session_key: The splunkd session key
        @type session_key: basestring

        @param owner: The owner of the particular thing
        @type session_key: basestring

        @param objecttype:    The type of object to delete
        @type objecttype: basestring

        @param filterdata: Particular filtering parameters - very much required because
            I don't want to allow people to delete the entire database just yet
        @type filterdata: dict

        @param current_user_name: user name of who is performing this operation
        @type current_user_name: basestring

        @param host_base_uri: The base uri <scheme>://<host>:<port> of the target host.  '' targets local host.
        @type host_base_uri: basestring
        """
        backend = self.get_backend(session_key, host_base_uri=host_base_uri)
        user_name = current_user_name
        if user_name is None:
            user_name = getCurrentUser()['name']
        if filter_data is None or len(filter_data) == 0:
            filter_data = {"object_type": objecttype}
        logger.info("%s user=%s method=delete_all objecttype=%s filter=%s",
                    LOG_CHANGE_TRACKING, user_name, objecttype, filter_data)
        backend.delete_all(
            session_key,
            owner,
            objecttype,
            filter_data,
            host_base_uri=host_base_uri
        )
        return

    def batch_save(self, session_key, owner, data_list, objecttype='UNKNOWN', current_user_name=None, host_base_uri=''):
        """
        WARNING: object_type must be set for all the elements of data_list before calling this function

        Save multiple objects in single save

        @param session_key: session_key
        @type session_key: basestring

        @param owner: user who is performing this operation
        @type owner: string

        @param data_list: objects to upsert
        @type data_list: list of dictionary

        @param host_base_uri: The base uri <scheme>://<host>:<port> of the target host.  '' targets local host.
        @type host_base_uri: basestring

        @param objecttype: The type of object to save. Required for audit logging purposes (FedRamp compliance)
        @type objecttype: basestring

        @param current_user_name: user name of who is performing this operation
        @type current_user_name: basestring

        @return: ids of objects upserted on success, throws exceptions on errors
        @rtype: list of strings
        """
        backend = self.get_backend(session_key, host_base_uri=host_base_uri)
        user_name = current_user_name
        if user_name is None:
            user_name = getCurrentUser()['name']

        if objecttype == 'UNKNOWN' and data_list is not None and isinstance(data_list, list) and len(data_list) > 0:
            for data in data_list:
                if 'object_type' in data and data['object_type'] is not None:
                    objecttype = data['object_type']
                    break

        logger.info("%s user=%s method=batch_save objecttype=%s",
                    LOG_CHANGE_TRACKING, user_name, objecttype)

        results = backend.batch_save(session_key, owner, data_list, host_base_uri=host_base_uri)
        return results

###############################################################################
# Specific retrieval methods
###############################################################################

    def entity_clarification(self, objecttype, data):
        """
        Define an _type field for entities for use in rules vs non-rules
        Sigh .... We should get rid of this at some point
        """
        if objecttype != 'entity':
            return
        if '_type' in data:
            return
        # At this point we need to make up the key to clarify what type of entity it is
        if 'identifier' in data:
            data['_type'] = 'entity'
        else:
            data['_type'] = 'rule'

#######################################################################################
# CSV Loading Methods
#######################################################################################

    def get_all_aliases(self, session_key, owner, host_base_uri=''):
        """
        Returns all of the entity aliases and all of the stuff they own.
        We needed to add a filter parameter as well.

        @param session_key: session_key
        @type session_key: basestring

        @param owner: user who is performing this operation
        @type owner: string

        @param host_base_uri: The base uri <scheme>://<host>:<port> of the target host.  '' targets local host.
        @type host_base_uri: basestring

        @return: entity alias'
        @rtype: dict
        """
        backend = self.get_backend(session_key, host_base_uri=host_base_uri)
        currentUser = getCurrentUser()
        logger.info("%s user=%s method=get_all_aliases", LOG_CHANGE_TRACKING, currentUser['name'])
        alias_dict = {}
        aliases = backend.paged_get_all(
            session_key,
            owner,
            'entity',
            sort_key='identifying_name',
            sort_dir='asc',
            fields=['identifier.fields', 'informational.fields'],
            host_base_uri=host_base_uri
        )
        for field in ['identifier', 'informational']:
            alias_set = set()
            # Add all of the aliases to the alias set
            for alias in aliases:
                if field not in alias:
                    continue
                if "fields" not in alias[field]:
                    continue
                alias_set.update(alias[field]['fields'])
            alias_dict[field] = list(alias_set)
        return alias_dict

    def get_count(self, session_key, owner, object_type, filter_data=None, host_base_uri=''):
        """
        Gets the count for the specified object type.  Right now its a little hackey

        @param session_key: session_key
        @type session_key: basestring

        @param owner: user who is performing this operation
        @type owner: string

        @param host_base_uri: The base uri <scheme>://<host>:<port> of the target host.  '' targets local host.
        @type host_base_uri: basestring

        @return: entity alias'
        @rtype: dict
        """
        backend = self.get_backend(session_key, host_base_uri=host_base_uri)
        currentUser = getCurrentUser()
        logger.info("%s user=%s method=get_count", LOG_CHANGE_TRACKING, currentUser['name'])

        keys = backend.get_all(
            session_key,
            owner,
            object_type,
            filter_data=filter_data,
            fields=['_key'],
            host_base_uri=host_base_uri
        )
        return {"count": len(keys)}