#!/usr/bin/env python
# coding=utf-8

__name__ = "trackme_rest_handler_licensing.py"
__author__ = "TrackMe Limited"
__copyright__ = "Copyright 2022-2026, TrackMe Limited, U.K."
__credits__ = "TrackMe Limited, U.K."
__license__ = "TrackMe Limited, all rights reserved"
__version__ = "0.1.0"
__maintainer__ = "TrackMe Limited, U.K."
__email__ = "support@trackme-solutions.com"
__status__ = "PRODUCTION"

# Built-in libraries
import json
import os
import sys
import time
import uuid
from collections import OrderedDict

# splunk home
splunkhome = os.environ["SPLUNK_HOME"]

# append current directory
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

# import libs
import import_declare_test

# set logging
from trackme_libs_logging import setup_logger

logger = setup_logger(
    "trackme.rest.licensing_admin", "trackme_rest_api_licensing_admin.log"
)
# Redirect global logging to use the same handler
import logging
logging.getLogger().handlers = logger.handlers
logging.getLogger().setLevel(logger.level)


# import rest handler
import trackme_rest_handler

# import trackme libs
from trackme_libs import trackme_getloglevel

# import trackme licensing libs
from trackme_libs_licensing import (
    trackme_return_license_status,
    trackme_start_trial,
    trackme_return_license_status_offline,
)

# import Splunk libs
import splunklib.client as client

# import cryptolense
from licensing.models import *
from licensing.methods import Key, Helpers


class TrackMeHandlerLicensingAdmin_v2(trackme_rest_handler.RESTHandler):
    def __init__(self, command_line, command_arg):
        super(TrackMeHandlerLicensingAdmin_v2, self).__init__(
            command_line, command_arg, logger
        )

    def get_resource_group_desc_licensing(self, request_info, **kwargs):
        response = {
            "resource_group_name": "licensing",
            "resource_group_desc": "Endpoints for the purposes of license management (admin operations)",
        }

        return {"payload": response, "status": 200}

    # set the license key
    def post_set_license(self, request_info, **kwargs):
        describe = False

        # Retrieve from data
        try:
            resp_dict = json.loads(str(request_info.raw_args["payload"]))
        except Exception as e:
            resp_dict = None

        if resp_dict is not None:
            try:
                describe = resp_dict["describe"]
                if describe in ("true", "True"):
                    describe = True
            except Exception as e:
                describe = False
        else:
            describe = False

        if not describe:
            license_key = resp_dict["license_key"]

        # if describe is requested, show the usage
        if describe:
            response = {
                "describe": "This endpoint sets the license key for this deployment, it requires a POST call with the no options:",
                "resource_desc": "Get the license status",
                "resource_spl_example": '| trackme url="/services/trackme/v2/licensing/admin/license_key" mode="post"',
                "options": [
                    {
                        "license_key": "The license key to be set",
                    }
                ],
            }
            return {"payload": response, "status": 200}

        # Get splunkd port
        splunkd_port = request_info.server_rest_port

        # Get service
        service = client.connect(
            owner="nobody",
            app="trackme",
            port=splunkd_port,
            token=request_info.system_authtoken,
            timeout=600,
        )

        # set loglevel
        loglevel = trackme_getloglevel(
            request_info.system_authtoken, request_info.server_rest_port
        )
        logger.setLevel(loglevel)

        try:
            # check the license validity first
            response = trackme_return_license_status(license_key)

            if int(response.get("license_is_valid")) == 1:
                collection_name = "kv_trackme_license_key"
                collection = service.kvstore[collection_name]

                # Get the current record
                # Notes: the record is returned as an array, as we search for a specific record, we expect one record only

                try:
                    kvrecords = collection.data.query()

                except Exception as e:
                    kvrecords = None

                if kvrecords:
                    # Remove any record in the KV
                    for kvrecord in kvrecords:
                        key = kvrecord.get("_key")
                        logger.info(
                            f'purging KVrecord="{json.dumps(kvrecord, indent=2)}"'
                        )
                        collection.data.delete(json.dumps({"_key": key}))

                # set
                logger.info("attempting to set the license key KVstore record")
                collection.data.insert(
                    json.dumps(
                        {
                            "_key": license_key,
                            "license_string": response.get("license_string"),
                            "license_type": "subscription",
                        }
                    )
                )

                # log
                logger.info(
                    f'license set key operation terminated, response="{json.dumps(response, indent=2)}"'
                )

                # return
                return {"payload": response, "status": 200}

            else:
                # return
                return {"payload": response, "status": 500}

        except Exception as e:
            response = {
                "action": "failure",
                "message": "An exception was encountered while attempting to setup the license key",
                "exception": str(e),
            }

            return {"payload": response, "status": 500}

    # Upload license file
    def post_upload_license_file(self, request_info, **kwargs):
        describe = False

        # Retrieve from data
        try:
            resp_dict = json.loads(str(request_info.raw_args["payload"]))
        except Exception as e:
            resp_dict = None

        if resp_dict is not None:
            try:
                describe = resp_dict["describe"]
                if describe in ("true", "True"):
                    describe = True
            except Exception as e:
                describe = False
        else:
            describe = False

        if not describe:
            license_file = resp_dict["license_file"]

        # if describe is requested, show the usage
        if describe:
            response = {
                "describe": "This endpoint verifies and upload a license file, it requires a POST call with the following options:",
                "resource_desc": "Upload license file",
                "resource_spl_example": '| trackme url="/services/trackme/v2/licensing/admin/upload_license_file" mode="post"',
                "options": [
                    {
                        "license_file": "The license file content",
                    }
                ],
            }
            return {"payload": response, "status": 200}

        # Get splunkd port
        splunkd_port = request_info.server_rest_port

        # Get service
        service = client.connect(
            owner="nobody",
            app="trackme",
            port=splunkd_port,
            token=request_info.system_authtoken,
            timeout=600,
        )

        # set loglevel
        loglevel = trackme_getloglevel(
            request_info.system_authtoken, request_info.server_rest_port
        )
        logger.setLevel(loglevel)

        try:
            # check the license validity first
            response = trackme_return_license_status_offline(license_file)

            if int(response.get("license_is_valid")) == 1:
                collection_name = "kv_trackme_license_key"
                collection = service.kvstore[collection_name]

                # Get the current record
                # Notes: the record is returned as an array, as we search for a specific record, we expect one record only

                try:
                    kvrecords = collection.data.query()

                except Exception as e:
                    kvrecords = None

                if kvrecords:
                    # Remove any record in the KV
                    for kvrecord in kvrecords:
                        key = kvrecord.get("_key")
                        logger.info(
                            f'purging KVrecord="{json.dumps(kvrecord, indent=2)}"'
                        )
                        collection.data.delete(json.dumps({"_key": key}))

                # set
                logger.info("attempting to set the license key KVstore record")
                collection.data.insert(
                    json.dumps(
                        {
                            "license_string": response.get("license_string"),
                            "license_type": "subscription",
                        }
                    )
                )

                # log
                logger.info(
                    f'license set key operation terminated, response="{json.dumps(response, indent=2)}"'
                )

                # return
                return {"payload": response, "status": 200}

            else:
                # return
                return {"payload": response, "status": 500}

        except Exception as e:
            response = {
                "action": "failure",
                "message": "An exception was encountered while attempting to setup the license key",
                "exception": str(e),
            }

            return {"payload": response, "status": 500}

    # Start trial license
    def post_start_trial(self, request_info, **kwargs):
        describe = False

        # Retrieve from data
        try:
            resp_dict = json.loads(str(request_info.raw_args["payload"]))
        except Exception as e:
            resp_dict = None

        if resp_dict is not None:
            try:
                describe = resp_dict["describe"]
                if describe in ("true", "True"):
                    describe = True
            except Exception as e:
                describe = False
        else:
            describe = False

        # if describe is requested, show the usage
        if describe:
            response = {
                "describe": "This endpoint starts a trial license period for this deployment, it requires a POST call with no options:",
                "resource_desc": "Starts a Trial period",
                "resource_spl_example": '| trackme url="/services/trackme/v2/licensing/admin/start_trial" mode="post"',
            }
            return {"payload": response, "status": 200}

        # Get splunkd port
        splunkd_port = request_info.server_rest_port

        # Get service
        service = client.connect(
            owner="nobody",
            app="trackme",
            port=splunkd_port,
            token=request_info.system_authtoken,
            timeout=600,
        )

        # set loglevel
        loglevel = trackme_getloglevel(
            request_info.system_authtoken, request_info.server_rest_port
        )
        logger.setLevel(loglevel)

        try:
            # check the license validity first
            response = trackme_start_trial(request_info)

            logger.info(f'response="{response}"')

            if int(response.get("license_is_valid")) == 1:
                collection_name = "kv_trackme_license_key"
                collection = service.kvstore[collection_name]

                # Get the current record
                # Notes: the record is returned as an array, as we search for a specific record, we expect one record only

                try:
                    kvrecords = collection.data.query()

                except Exception as e:
                    kvrecords = None

                if kvrecords:
                    # Remove any record in the KV
                    for kvrecord in kvrecords:
                        key = kvrecord.get("_key")
                        logger.info(
                            f'purging KVrecord="{json.dumps(kvrecord, indent=2)}"'
                        )
                        collection.data.delete(json.dumps({"_key": key}))

                # set
                logger.info("attempting to set the Trial license key KVstore record")
                collection.data.insert(
                    json.dumps(
                        {
                            "_key": response.get("trial_key"),
                            "license_string": response.get("license_string"),
                            "license_type": "trial",
                        }
                    )
                )

                # log
                logger.info(
                    f'license Trial operation terminated, response="{json.dumps(response, indent=2)}"'
                )

                # return
                if response.get("action") == "success":
                    return {"payload": response, "status": 200}

                else:
                    return {"payload": response, "status": 500}

            else:
                # return
                return {"payload": response, "status": 500}

        except Exception as e:
            response = {
                "action": "failure",
                "message": "An exception was encountered while attempting to generate a Trial license",
                "exception": str(e),
            }

            return {"payload": response, "status": 500}

    # Enable developer mode
    def post_enable_developer_license(self, request_info, **kwargs):
        describe = False

        # Retrieve from data
        try:
            resp_dict = json.loads(str(request_info.raw_args["payload"]))
        except Exception as e:
            resp_dict = None

        if resp_dict is not None:
            try:
                describe = resp_dict["describe"]
                if describe in ("true", "True"):
                    describe = True
            except Exception as e:
                describe = False
        else:
            describe = False

        # if describe is requested, show the usage
        if describe:
            response = {
                "describe": "This endpoint enables the developer mode for this deployment, it requires a POST call with no options:",
                "resource_desc": "Enable the developer mode license",
                "resource_spl_example": '| trackme url="/services/trackme/v2/licensing/admin/enable_developer_license" mode="post"',
            }
            return {"payload": response, "status": 200}

        # Get splunkd port
        splunkd_port = request_info.server_rest_port

        # Get service
        service = client.connect(
            owner="nobody",
            app="trackme",
            port=splunkd_port,
            token=request_info.system_authtoken,
            timeout=600,
        )

        # set loglevel
        loglevel = trackme_getloglevel(
            request_info.system_authtoken, request_info.server_rest_port
        )
        logger.setLevel(loglevel)

        try:
            collection_name = "kv_trackme_license_key"
            collection = service.kvstore[collection_name]

            # Get the current record
            # Notes: the record is returned as an array, as we search for a specific record, we expect one record only

            try:
                kvrecords = collection.data.query()

            except Exception as e:
                kvrecords = None

            if kvrecords:
                # Remove any record in the KV
                for kvrecord in kvrecords:
                    key = kvrecord.get("_key")
                    logger.info(f'purging KVrecord="{json.dumps(kvrecord, indent=2)}"')
                    collection.data.delete(json.dumps({"_key": key}))

            # set
            logger.info(
                "attempting to set the developer mode license key KVstore record"
            )

            new_record = {
                "license_string": json.dumps(
                    {
                        "uuid": str(uuid.uuid4()),
                        "expires": int(time.time()) + 2592000,
                    },
                    indent=2,
                ),
                "license_type": "developer",
            }

            try:
                collection.data.insert(json.dumps(new_record))
                response = {
                    "action": "success",
                    "license_string": new_record.get("license_string"),
                    "licence_type": new_record.get("license_type"),
                }

            except Exception as e:
                response = {
                    "action": "failure",
                    "license_string": new_record.get("license_string"),
                    "licence_type": new_record.get("license_type"),
                    "exception": str(e),
                }

            # log
            logger.info(
                f'license developer operation terminated, response="{json.dumps(response, indent=2)}"'
            )

            # return
            if response.get("action") == "success":
                return {"payload": response, "status": 200}

            else:
                return {"payload": response, "status": 500}

        except Exception as e:
            response = {
                "action": "failure",
                "message": "An exception was encountered while attempting to enable the developer license",
                "exception": str(e),
            }

            return {"payload": response, "status": 500}

    # Reset licensing
    def post_reset_license(self, request_info, **kwargs):
        describe = False

        # Retrieve from data
        try:
            resp_dict = json.loads(str(request_info.raw_args["payload"]))
        except Exception as e:
            resp_dict = None

        if resp_dict is not None:
            try:
                describe = resp_dict["describe"]
                if describe in ("true", "True"):
                    describe = True
            except Exception as e:
                describe = False
        else:
            describe = False

        # if describe is requested, show the usage
        if describe:
            response = {
                "describe": "This endpoint resets the current license registration, it requires a POST call with no options:",
                "resource_desc": "Enable the developer mode license",
                "resource_spl_example": '| trackme url="/services/trackme/v2/licensing/admin/reset_license" mode="post"',
            }
            return {"payload": response, "status": 200}

        # Get splunkd port
        splunkd_port = request_info.server_rest_port

        # Get service
        service = client.connect(
            owner="nobody",
            app="trackme",
            port=splunkd_port,
            token=request_info.system_authtoken,
            timeout=600,
        )

        # set loglevel
        loglevel = trackme_getloglevel(
            request_info.system_authtoken, request_info.server_rest_port
        )
        logger.setLevel(loglevel)

        try:
            collection_name = "kv_trackme_license_key"
            collection = service.kvstore[collection_name]

            # Get the current record
            # Notes: the record is returned as an array, as we search for a specific record, we expect one record only

            try:
                kvrecords = collection.data.query()

            except Exception as e:
                kvrecords = None

            if kvrecords:
                # Remove any record in the KV
                for kvrecord in kvrecords:
                    key = kvrecord.get("_key")
                    logger.info(f'purging KVrecord="{json.dumps(kvrecord, indent=2)}"')
                    collection.data.delete(json.dumps({"_key": key}))

            # log
            logger.info("reset license operation terminated")

            return {
                "payload": {
                    "action": "success",
                },
                "status": 200,
            }

        except Exception as e:
            response = {
                "action": "failure",
                "message": "An exception was encountered while attempting to reset the current license registration",
                "exception": str(e),
            }

            return {"payload": response, "status": 500}