#!/usr/bin/python3
#
# Copyright 2023-2025 Centreon
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# For more information : contact@centreon.com
#

from os import makedirs, environ
import base64
import hashlib
import ssl
import time
from robot.libraries.BuiltIn import BuiltIn, RobotNotRunningError
from socket import gethostname
import Common
import json
from robot.api import logger


def import_robot_resources():
    global VAR_ROOT, ETC_ROOT, CONF_DIR
    try:
        BuiltIn().import_resource('db_variables.resource')
        ETC_ROOT = BuiltIn().get_variable_value("${EtcRoot}")
        VAR_ROOT = BuiltIn().get_variable_value("${VarRoot}")
        CONF_DIR = ETC_ROOT + "/centreon-engine"
    except RobotNotRunningError:
        # Handle this case if Robot Framework is not running
        print("Robot Framework is not running. Skipping resource import.")


ETC_ROOT = ""
VAR_ROOT = ""
CONF_DIR = ""

import_robot_resources()


def ctn_used_address():
    """
    ctn_used_address

    Get USED_ADDRESS env variable content
    """
    return environ.get('USED_ADDRESS', '127.0.0.1')


def ctn_host_hostname():
    """
    ctn_host_hostname

    Get HOST_HOSTNAME env variable content
    """
    return environ.get('HOST_HOSTNAME', Common.ctn_get_hostname())


def ctn_get_agent_conf_path():
    """
    ctn_get_conf_path

    Get CONF_DIR path
    """
    return CONF_DIR+"/centagent0.json"


agent_config = f"""
{{
    "log_level":"trace",
    "endpoint":"localhost:4317",
    "host":"host_1",
    "log_type":"file",
    "log_file":"/tmp/var/log/centreon-engine/centreon-agent.log",
    "custom_check_file": "{VAR_ROOT}/lib/centreon-engine/custom_checks.ini" """


agent_encrypted_config = f"""
{{
    "log_level":"trace",
    "endpoint":"{ctn_host_hostname()}:4318",
    "host":"host_1",
    "log_type":"file",
    "log_file":"{VAR_ROOT}/log/centreon-engine/centreon-agent.log" """


reversed_agent_config = f"""
{{
    "log_level":"trace",
    "endpoint":"{ctn_host_hostname()}:4320",
    "host":"host_1",
    "log_type":"file",
    "log_file":"{VAR_ROOT}/log/centreon-engine/centreon-agent.log" """

reversed_agent_encrypted_config = f"""
{{
    "log_level":"trace",
    "endpoint":"{ctn_host_hostname()}:4321",
    "host":"host_1",
    "log_type":"file",
    "log_file":"{VAR_ROOT}/log/centreon-engine/centreon-agent.log" """


def ctn_get_cert_fingerprint(cert_path: str) -> str:
    """ctn_get_cert_fingerprint
    Compute the SHA256 fingerprint of a PEM certificate, base64-encoded.
    This matches the C++ cert_tree::cert_sha() format used by the agent.
    Args:
        cert_path: path to the PEM certificate file
    Returns:
        base64-encoded SHA256 digest of the DER-encoded certificate
    """
    pem_data = open(cert_path).read()
    der_data = ssl.PEM_cert_to_DER_cert(pem_data)
    sha256_digest = hashlib.sha256(der_data).digest()
    return base64.b64encode(sha256_digest).decode('ascii')


def ctn_config_centreon_agent(key_path: str = None, cert_path: str = None, ca_path: str = None, token: str = None, ca_common_name: str = None, security_mode: str = "full", have_token: bool = True, nb_agent: int = 1, fingerprint: str = None):
    """ctn_config_centreon_agent
    Creates a default centreon agent config listening on  0.0.0.0:4317 (no encryption) or 0.0.0.0:4318 (encryption)
    Args:
        key_path: path of the private key file
        cert_path: path of public certificate file
        ca_path: path of the authority certificate file
        token:
        ca_common_name: CN
        security_mode: full, no or insecure
        have_token: add token or empty token in config file
        nb_agent: nb conf agent (several agents config files)
        fingerprint: base64-encoded SHA256 fingerprint of the CA cert; when set the
                     agent will bootstrap the CA over insecure TLS and validate it
                     using this fingerprint (implies encrypted endpoint, no ca_path needed)
    """
    # in case of wsl, agent is executed in windows host
    if environ.get("RUN_ENV", "") == "WSL":
        return

    token_default = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJjZW50cmVvbjY2MjQxIiwiaWF0IjoxNzQ0MDk3MDgxLCJleHAiOjkyMjMzNzIwMzV9.QkrT77i211-CvXoXqaBxRMzxajzA3-DK-DGVrbvJWA8"

    makedirs(CONF_DIR, mode=0o777, exist_ok=True)
    with open(f"{VAR_ROOT}/lib/centreon-engine/custom_checks.ini", "w") as ff:
        ff.write("[custom_checks]\n")
        ff.write("check_echo = " +
                 ctn_echo_command("$ARG2$ $ARG1$ from custom check") + "\n")
        ff.write("custom_check_2 = /path/to/custom_check_2 -c /arg=<value>\n")

    use_encrypted_endpoint = cert_path is not None or ca_path is not None or fingerprint is not None
    use_encryption = key_path is not None or cert_path is not None or ca_path is not None or fingerprint is not None

    for agent_index in range(nb_agent):
        with open(f"{CONF_DIR}/centagent{agent_index}.json", "w") as ff:
            if use_encrypted_endpoint:
                ff.write(agent_encrypted_config)
            else:
                ff.write(agent_config)
            if use_encryption:
                ff.write(f",\n  \"encryption\":\"{security_mode}\"")
            if key_path is not None:
                ff.write(f",\n  \"private_key\":\"{key_path}\"")
            if cert_path is not None:
                ff.write(f",\n  \"public_cert\":\"{cert_path}\"")
            if ca_path is not None:
                ff.write(f",\n  \"ca\":\"{ca_path}\"")
            if ca_common_name is not None:
                ff.write(f",\n  \"ca_common_name\":\"{ca_common_name}\"")
            if fingerprint is not None:
                ff.write(f",\n  \"fingerprint\":\"{fingerprint}\"")
            if have_token:
                if token is not None:
                    ff.write(f",\n  \"token\":\"{token}\"")
                else:
                    ff.write(f",\n  \"token\":\"{token_default}\"")

            ff.write("\n}\n")


def ctn_config_centreon_agent_set_value(key: str, value: str):
    """ctn_config_set_value
    Set a value in the centreon agent config file
    Args:
        key: key to set
        value: value to set
    """
    makedirs(CONF_DIR, mode=0o777, exist_ok=True)

    with open(f"{CONF_DIR}/centagent0.json", "r+") as ff:
        content = json.load(ff)
        print(content)
        content[key] = value
        ff.seek(0)
        ff.truncate()
        json.dump(content, ff, indent=4)


def ctn_config_reverse_centreon_agent(key_path: str = None, cert_path: str = None, ca_path: str = None, token: str = None):
    """ctn_config_centreon_agent
    Creates a default reversed centreon agent config listening on  0.0.0.0:4320 (no encryption) or 0.0.0.0:4321 (encryption)
    Args:
        key_path: path of the private key file
        cert_path: path of public certificate file
        ca_path: path of the authority certificate file
    """
    # in case of wsl, agent is executed in windows host
    if environ.get("RUN_ENV", "") == "WSL":
        return

    token_default = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJjZW50cmVvbjY2MjQxIiwiaWF0IjoxNzQ0MDk3MDgxLCJleHAiOjkyMjMzNzIwMzV9.QkrT77i211-CvXoXqaBxRMzxajzA3-DK-DGVrbvJWA8"

    makedirs(CONF_DIR, mode=0o777, exist_ok=True)
    with open(f"{CONF_DIR}/centagent0.json", "w") as ff:
        if cert_path is not None or ca_path is not None:
            ff.write(reversed_agent_encrypted_config)
        else:
            ff.write(reversed_agent_config)
        ff.write(",\n  \"reversed_grpc_streaming\":true")
        if key_path is not None or cert_path is not None or ca_path is not None:
            ff.write(",\n  \"encryption\":\"full\"")
        if key_path is not None:
            ff.write(f",\n  \"private_key\":\"{key_path}\"")
        if cert_path is not None:
            ff.write(f",\n  \"public_cert\":\"{cert_path}\"")
        if ca_path is not None:
            ff.write(f",\n  \"ca\":\"{ca_path}\"")
        if token is not None:
            ff.write(f",\n  \"token\":\"{token}\"")
        else:
            ff.write(f",\n  \"token\":\"{token_default}\"")
        ff.write("\n}\n")


def ctn_echo_command(to_echo: str):
    """
    ctn_echo_command
    returned an echo command usable by testing agent OS
    Args:
        to_echo: string to print to stdout
    """
    if environ.get("RUN_ENV", "") == "WSL":
        return '"' + environ.get('PWSH_PATH') + '"' + " C:/Users/Public/echo.ps1 " + to_echo
    else:
        return "/bin/echo " + to_echo


def ctn_check_pl_command(arg: str):
    """
    ctn_check_pl_command
    returned an check.pl command usable by testing agent OS
    Args:
        arg: arguments to pass to check.pl or check.ps1 command
    """
    if environ.get("RUN_ENV", "") == "WSL":
        return '"' + environ.get('PWSH_PATH') + '"' + " C:/Users/Public/check.ps1 " + arg + " " + environ.get("WINDOWS_PROJECT_PATH")
    else:
        return "/tmp/var/lib/centreon-engine/check.pl " + arg


def ctn_get_drive_statistics(drive_name_format: str):
    """
    ctn_get_drive_statistics
    return a dictionary of drive statistics indexed by expected perfdata names
    Args:
        drive_name_format: format of the drive name to search for
    """
    if environ.get("RUN_ENV", "") == "WSL":
        drive_dict = {}
        json_test_args = environ.get("JSON_TEST_PARAMS")
        test_args = json.loads(json_test_args)
        for drive in test_args["drive"]:
            if drive['Free'] is not None:
                drive_dict[drive_name_format.format(drive['Name'])] = (
                    100 * drive['Free']) / (drive['Used'] + drive['Free'])
        return drive_dict
    else:
        return None


def ctn_get_uptime():
    """
    ctn_get_uptime
    return a dict with only one element: uptime => uptime value
    """
    if environ.get("RUN_ENV", "") == "WSL":
        uptime_dict = {}
        json_test_args = environ.get("JSON_TEST_PARAMS")
        test_args = json.loads(json_test_args)
        if test_args["uptime"] is not None:
            uptime_dict['uptime'] = time.time() - test_args["uptime"]
            return uptime_dict
    return None


def ctn_get_memory():
    """
    ctn_get_memory statistics
    return a dict with these elements (expected perfdata):
    - memory.free.bytes
    - memory.usage.bytes
    - memory.usage.percentage
    - swap.free.bytes
    - swap.usage.bytes
    - swap.usage.percentage
    - virtual-memory.free.bytes
    - virtual-memory.usage.bytes
    - virtual-memory.usage.percentage
    """

    if environ.get("RUN_ENV", "") == "WSL":
        memory_dict = {'swap.free.bytes': None,
                       'swap.usage.bytes': None, 'swap.usage.percentage': None}
        json_test_args = environ.get("JSON_TEST_PARAMS")
        test_args = json.loads(json_test_args)
        if test_args["mem_info"] is not None:
            # values of systeminfo are given in Mb
            virtual_free = int(test_args["mem_info"]["virtual_free"].replace(
                ",", "").split()[0]) * 1024 * 1024
            virtual_max = int(test_args["mem_info"]["virtual_max"].replace(
                ",", "").split()[0])*1024 * 1024
            free = int(test_args["mem_info"]["free"].replace(
                ",", "").split()[0])*1024 * 1024
            total = int(test_args["mem_info"]["total"].replace(
                ",", "").split()[0])*1024 * 1024
            memory_dict['virtual-memory.free.bytes'] = virtual_free
            memory_dict['virtual-memory.usage.bytes'] = virtual_max - virtual_free
            memory_dict['virtual-memory.usage.percentage'] = 100 - \
                (100.0 * virtual_free) / virtual_max

            memory_dict['memory.free.bytes'] = free
            memory_dict['memory.usage.bytes'] = total - free
            memory_dict['memory.usage.percentage'] = 100 - \
                (100.0 * free) / total
            return memory_dict
    return None


def ctn_get_service():
    """
    ctn_get_service statistics
    return a dict with these elements (expected perfdata):
    - services.stopped.count
    - services.starting.count
    - services.stopping.count
    - services.running.count
    - services.continuing.count
    - services.pausing.count
    - services.paused.count
    """

    if environ.get("RUN_ENV", "") == "WSL":
        service_dict = {'services.stopped.count': 0, 'services.starting.count': None, 'services.stopping.count': None, 'services.running.count': 0,
                        'services.continuing.count': None, 'services.pausing.count': None, 'services.paused.count': None}
        json_test_args = environ.get("JSON_TEST_PARAMS")
        test_args = json.loads(json_test_args)
        if test_args["serv_stat"] is not None:
            service_dict["services.stopped.count"] = test_args["serv_stat"]["services.stopped.count"]
            service_dict["services.running.count"] = test_args["serv_stat"]["services.running.count"]
            return service_dict
    return None
