Splunk_Deploiement/apps/trackme/lib/trackme_libs_licensing.py

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

__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"

import os
import sys
import requests
import re
import json
import time
import datetime
import logging
import uuid
import urllib3

urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

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

# append lib
sys.path.append(os.path.join(splunkhome, "etc", "apps", "trackme", "lib"))

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

# logging:
# To avoid overriding logging destination of callers, the libs will not set on purpose any logging definition
# and rely on callers themselves


def trackme_check_license(server_rest_uri, session_key):
    # build header and target
    header = {
        "Authorization": f"Splunk {session_key}",
        "Content-Type": "application/json",
    }
    target_url = f"{server_rest_uri}/services/trackme/v2/licensing/license_status"

    try:
        response = requests.get(
            target_url,
            headers=header,
            verify=False,
            timeout=600,
        )
        return json.loads(response.text)

    except Exception as e:
        raise Exception(
            f'An exception was encountered while attempting to verify the license status, exception="{str(e)}"'
        )


def trackme_return_license_status(license_key):
    # Get the RSA pub key
    with open(
        os.path.join(
            splunkhome, "etc", "apps", "trackme", "lib", "licensing", "trackme_rsa.pub"
        ),
        "r",
    ) as f:
        for line in f:
            RSAPubKey = line
            break

    # Get the public access token
    with open(
        os.path.join(
            splunkhome,
            "etc",
            "apps",
            "trackme",
            "lib",
            "licensing",
            "trackme_token.pub",
        ),
        "r",
    ) as f:
        for line in f:
            auth = line
            break

    try:
        result = Key.activate(
            token=auth,
            rsa_pub_key=RSAPubKey,
            product_id=18301,
            key=license_key,
            machine_code="",
        )

        response = {}

        action = None
        message = None
        license_is_valid = 0
        license_expiration = None
        license_features = []
        license_string = None

        if result[0] == None:
            action = "failure"
            license_is_valid = 0
            message = str(result[1])

        else:
            license_key = result[0]
            license_expiration = str(license_key.expires)
            license_features.append(
                {
                    "time_limit": str(license_key.f1),
                    "trial": str(license_key.f2),
                    "enterprise": str(license_key.f3),
                    "unlimited": str(license_key.f4),
                    "free_extended": str(license_key.f5),
                }
            )
            license_string = result[0].save_as_string()

            # get the remaining time in seconds
            expiration_dt = datetime.datetime.strptime(
                license_expiration, "%Y-%m-%d %H:%M:%S"
            )
            license_expiration_epoch = round(expiration_dt.timestamp())
            time_before_expiration = round(license_expiration_epoch - time.time())
            logging.debug(f'license_expiration_epoch="{license_expiration_epoch}"')
            logging.debug(f'time_before_expiration="{time_before_expiration}"')

            # the license has expired
            if not time_before_expiration > 0:
                action = "failure"
                license_is_valid = 0
                message = "The license has expired"

            else:
                action = "success"
                license_is_valid = 1
                message = "The license is valid"

        response = {
            "action": action,
            "license_is_valid": license_is_valid,
            "message": message,
            "license_expiration": license_expiration,
            "license_expiration_countdown_sec": time_before_expiration,
            "license_features": license_features,
            "license_string": license_string,
        }

        logging.debug(f'response="{json.dumps(response, indent=2)}"')
        return response

    except Exception as e:
        logging.error(
            f'An exception occurred while attempting to verify the license status, exception="{str(e)}"'
        )
        return response


def trackme_return_license_status_offline(license_string):
    # Get the RSA pub key
    with open(
        os.path.join(
            splunkhome, "etc", "apps", "trackme", "lib", "licensing", "trackme_rsa.pub"
        ),
        "r",
    ) as f:
        for line in f:
            RSAPubKey = line
            break

    try:
        # log
        logging.debug(f'Verifying license from KVstore record="{license_string}"')

        # get license key
        license_key = LicenseKey.load_from_string(RSAPubKey, license_string)

        # init
        response = {}
        action = None
        message = None
        license_is_valid = 0
        license_expiration = None
        license_features = []

        if license_key == None:
            action = "failure"
            license_is_valid = 0
            message = "This license is not valid"
            license_string = None

        else:
            license_expiration = str(license_key.expires)
            license_features.append(
                {
                    "time_limit": str(license_key.f1),
                    "trial": str(license_key.f2),
                    "enterprise": str(license_key.f3),
                    "unlimited": str(license_key.f4),
                    "free_extended": str(license_key.f5),
                }
            )
            license_string = license_key.save_as_string()

            # get the remaining time in seconds
            expiration_dt = datetime.datetime.strptime(
                license_expiration, "%Y-%m-%d %H:%M:%S"
            )
            license_expiration_epoch = round(expiration_dt.timestamp())
            time_before_expiration = round(license_expiration_epoch - time.time())
            logging.debug(f'license_expiration_epoch="{license_expiration_epoch}"')
            logging.debug(f'time_before_expiration="{time_before_expiration}"')

            # the license has expired
            if not time_before_expiration > 0:
                action = "failure"
                license_is_valid = 0
                message = "The license has expired"

            else:
                action = "success"
                license_is_valid = 1
                message = "The license is valid"

        response = {
            "action": action,
            "license_is_valid": license_is_valid,
            "message": message,
            "license_expiration": license_expiration,
            "license_expiration_countdown_sec": time_before_expiration,
            "license_features": license_features,
            "license_string": license_string,
        }

        logging.debug(f'response="{json.dumps(response, indent=2)}"')
        return response

    except Exception as e:
        logging.error(
            f'An exception occurred while attempting to verify the license status, exception="{str(e)}"'
        )
        return response


def trackme_return_license_status_developer(license_string):
    try:
        # load as a dict
        license_string = json.loads(license_string)

        # log
        logging.debug(f'Verifying license from KVstore record="{license_string}"')

        # init
        response = {}
        action = None
        message = None
        license_is_valid = 0
        license_expiration = None
        license_features = []

        # get the remaining time in seconds
        license_expiration_epoch = license_string.get("expires")
        time_before_expiration = round(license_expiration_epoch - time.time())
        logging.debug(f'license_expiration_epoch="{license_expiration_epoch}"')
        logging.debug(f'time_before_expiration="{time_before_expiration}"')

        # convert
        license_expiration = time.strftime(
            "%Y-%m-%d %H:%M:%S", time.localtime(license_expiration_epoch)
        )

        # the license has expired
        if not time_before_expiration > 0:
            action = "failure"
            license_is_valid = 0
            message = "The license has expired"

        else:
            action = "success"
            license_is_valid = 1
            message = "The license is valid"

        response = {
            "action": action,
            "license_is_valid": license_is_valid,
            "message": message,
            "license_expiration": license_expiration,
            "license_expiration_countdown_sec": time_before_expiration,
            "license_features": license_features,
            "license_string": json.dumps(license_string),
        }

        logging.debug(f'response="{json.dumps(response, indent=2)}"')
        return response

    except Exception as e:
        logging.error(
            f'An exception occurred while attempting to verify the license status, exception="{str(e)}"'
        )
        return response


def trackme_start_trial(reqinfo):
    header = {
        "Authorization": f"Splunk {reqinfo.session_key}",
        "Content-Type": "application/json",
    }

    #
    # retrieve the instance guid
    #

    instance_guid = None
    target_url = f"{reqinfo.server_rest_uri}/services/server/info"

    try:
        response = requests.get(target_url, headers=header, verify=False, timeout=600)
        logging.debug(f'success retrieving server info, data="{response.text}"')

        pattern = r'name="guid"\>([^\<]+)\<\/s:key>'
        match = re.search(pattern, response.text)

        if match:
            instance_guid = match.group(1)
            logging.debug(f'instance_guid="{instance_guid}"')

    except Exception as e:
        logging.error(f'failed to retrieve the instance info, exception="{str(e)}"')

    #
    # verify if running in SHC, if so retrieve the shc_label
    #

    target_url = f"{reqinfo.server_rest_uri}/services/server/roles"
    is_shc = False

    try:
        response = requests.get(target_url, headers=header, verify=False, timeout=600)
        logging.debug(f'success retrieving server roles, data="{response.text}"')

        pattern = r"(\<s:item\>shc_member\<\/s:item\>)"
        match = re.search(pattern, response.text)

        if match:
            is_shc = True
            logging.debug("this instance is a member of a SHC cluster")

    except Exception as e:
        logging.error(f'failed to retrieve the instance roles, exception="{str(e)}"')

    # if running in SHC, extract the shc_label
    if is_shc:
        target_url = f"{reqinfo.server_rest_uri}/services/shcluster/config"

        try:
            response = requests.get(
                target_url, headers=header, verify=False, timeout=600
            )
            logging.debug(f'success retrieving shcluster info, data="{response.text}"')

            pattern = r'\<s:key\sname="shcluster_label"\>([^\<]+)\<\/s:key\>'
            match = re.search(pattern, response.text)

            if match:
                shcluster_label = match.group(1)
                logging.debug(f'shcluster_label="{shcluster_label}"')

        except Exception as e:
            logging.error(
                f'failed to retrieve the shcluster_label, exception="{str(e)}"'
            )

    # define the license_identifier
    license_identifier = None

    # if we failed to identify the instance_guid, generate a unique uuid
    if not instance_guid:
        instance_guid = uuid.uuid4()

    # investigate and set
    if is_shc:
        if shcluster_label:
            license_identifier = shcluster_label
        else:
            license_identifier = instance_guid

    else:
        license_identifier = instance_guid

    # Get the RSA pub key
    with open(
        os.path.join(
            splunkhome, "etc", "apps", "trackme", "lib", "licensing", "trackme_rsa.pub"
        ),
        "r",
    ) as f:
        for line in f:
            RSAPubKey = line
            break

    # Get the trial access token
    with open(
        os.path.join(
            splunkhome,
            "etc",
            "apps",
            "trackme",
            "lib",
            "licensing",
            "trackme_trial.pub",
        ),
        "r",
    ) as f:
        for line in f:
            auth_createtrial = line
            break

    # Get the trial access token
    with open(
        os.path.join(
            splunkhome,
            "etc",
            "apps",
            "trackme",
            "lib",
            "licensing",
            "trackme_token.pub",
        ),
        "r",
    ) as f:
        for line in f:
            auth = line
            break

    try:
        trial_key = Key.create_trial_key(auth_createtrial, 18301, license_identifier)

        if trial_key[0] == None:
            logging.error(
                f"An error occurred while attempting to create a trial key: {trial_key[1]}"
            )
            raise Exception(
                f"An error occurred while attempting to create a trial key: {trial_key[1]}"
            )

        else:
            logging.info(
                f'The Trial license_key="{trial_key[0]}" was successfully generated'
            )
            result = Key.activate(
                token=auth,
                rsa_pub_key=RSAPubKey,
                product_id=18301,
                key=trial_key[0],
                machine_code=license_identifier,
            )

            # init
            response = {}
            action = None
            message = None
            license_is_valid = 0

            if result[0] == None:
                action = "failure"
                license_is_valid = 0
                message = str(result[1])

                logging.error(
                    f'Failed to activate the license key, message="{str(result[1])}"'
                )
                raise Exception(
                    f'Failed to activate the license key, message="{str(result[1])}"'
                )

            else:
                action = "success"
                license_is_valid = 1
                trial_key = trial_key[0]
                message = "The trial license was generated successfully"
                license_string = result[0].save_as_string()

            response = {
                "action": action,
                "license_is_valid": license_is_valid,
                "trial_key": trial_key,
                "message": message,
                "license_string": license_string,
                "license_type": "trial",
            }

            logging.info(f'response="{json.dumps(response, indent=2)}"')
            return response

    except Exception as e:
        logging.error(
            f'An exception occurred while attempting to generate the trial license, exception="{str(e)}"'
        )
        raise Exception(
            f'An exception occurred while attempting to generate the trial license, exception="{str(e)}"'
        )