import json
import os
import re
import logging
import subprocess
import urllib.parse

logging.basicConfig(
	filename="/var/log/netwitness/presidio/update_encrypted_password_12_5.log",
	encoding="utf-8",
	filemode="a",
	format="{asctime} - {levelname} - {message}",
	style="{",
	datefmt="%Y-%m-%d %H:%M",
	)
logger = logging.getLogger()
logger.setLevel(logging.DEBUG)


PRESIDIO_OUTPUT_SERVICE = "presidio-output"
PRESIDIO_UI_SERVICE = "presidio-ui"
CONFIG_DIR = "/etc/netwitness/presidio/configserver/configurations/"
APP_PRESIDIO_CONF_FILE_NAME = 'application-presidio.json'
ENCRYPT_UTIL_APP = "/var/lib/netwitness/presidio/install/configserver/EncryptionUtils.jar"
OLD_KEY = 'FortScale4Ever!!'
OLD_ENCRYPT_ALGO = 'AES'
ENCRYPTED_PASSWORD_PARAM = 'mongo.db.password'
OLD_ENCRYPTED_PASSWORD_PARAM = 'old.mongo.db.password'
OLD_ENCRYPTED_DATA_SOURCE_PASSWORD_PARAM = 'old.encrypted.password.source'
APP_CONF_FILE_NAME = "application.properties"
UI_CONF_FILE_NAME = "presidio-uiconf.properties"

def check_property_exist(conf_file_path, property_line):
	with open(conf_file_path, 'r') as file:
		lines = file.readlines()
	property_exists = False
	for i, line in enumerate(lines):
		if line.startswith(property_line):
			property_exists = True
			break
	return property_exists

# Function to get value of property in a file
def get_param_value(file_path, param):
    with open(file_path, 'r') as file:
        for line in file:
            if line.startswith(f"{param}="):
                return line.split('=', 1)[1].strip().replace("\=", "=")
    return None

def get_data_source_password():
    config_file_path = os.path.join(CONFIG_DIR, APP_PRESIDIO_CONF_FILE_NAME)
    password = None
    # Check if the property key exists
    with open(config_file_path, 'r') as json_file:
        application_presidio = json.load(json_file)
        source = application_presidio['dataPulling']['source']
        match = re.match('^nws://(.+):(.+)@(.+):(.+)$', source)
        password = match.group(2)
        password = urllib.parse.unquote(password)
    return password

def add_update_property_file(conf_file_name, key, value):
    value = value.replace("=", "\=")
    conf_file_path = os.path.join(CONFIG_DIR, conf_file_name)

    with open(conf_file_path, 'r') as file:
        lines = file.readlines()

    with open(conf_file_path, 'w') as file:
        updated = False
        for line in lines:
            if line.startswith(f"{key}="):
                file.write(f"{key}={value}\n")
                updated = True
            else:
                file.write(line)
        if not updated :
            file.write(f"{key}={value}\n")

def update_data_source_password(old_enc_pass, enc_pass) :
    config_file_path = os.path.join(CONFIG_DIR, APP_PRESIDIO_CONF_FILE_NAME)
    # Read JSON configuration
    with open(config_file_path, 'r') as json_file:
        application_presidio = json.load(json_file)
        source = application_presidio['dataPulling']['source']

    # Replace the captured part with the new value
    new_enc_config = source.replace(old_enc_pass, enc_pass)

    # Update the JSON configuration
    application_presidio["dataPulling"]["source"] = new_enc_config
    with open(config_file_path, 'w') as file:
        json.dump(application_presidio, file, indent=2)

def decrypt_password_with_old_key(password):
    old_original_password = None
    # Decrypt password using old key
    decrypt_cmd = ["java", "-jar", ENCRYPT_UTIL_APP, "decrypt_old_pwd", password, OLD_KEY, OLD_ENCRYPT_ALGO]
    old_original_password = subprocess.check_output(decrypt_cmd).decode().strip()
    return old_original_password

def decrypt_password_with_new_key(password):
    new_original_password = None
    # Decrypt password using old key
    decrypt_cmd = ["java", "-jar", ENCRYPT_UTIL_APP, "decrypt", password]
    new_original_password = subprocess.check_output(decrypt_cmd).decode().strip()
    return new_original_password

def encrypt_password_with_new_key(password):
    new_encrypted_password = None
    # Encrypt password using new key
    encrypt_cmd = ["java", "-jar", ENCRYPT_UTIL_APP, "encrypt", password]
    new_encrypted_password = subprocess.check_output(encrypt_cmd).decode().strip()
    return new_encrypted_password
    
def update_adapter_source_password():
    logging.info(f"updating data source password")
    old_enc_password = get_data_source_password()
    old_password = decrypt_password_with_old_key(old_enc_password)
    new_enc_password = encrypt_password_with_new_key(old_password)
    
    #save_old_password
    add_update_property_file(APP_CONF_FILE_NAME, OLD_ENCRYPTED_DATA_SOURCE_PASSWORD_PARAM, old_enc_password)

    #Update with new pwd
    update_data_source_password(old_enc_password, new_enc_password)    
    
def update_mongo_password(conf_file_name):
    logging.info(f"updating mongo password in {conf_file_name} file")
    config_file_path = os.path.join(CONFIG_DIR, conf_file_name)
    old_enc_password = get_param_value(config_file_path, ENCRYPTED_PASSWORD_PARAM)
    old_enc_password = urllib.parse.unquote(old_enc_password)
    old_password = decrypt_password_with_old_key(old_enc_password)
    new_enc_password = encrypt_password_with_new_key(old_password)

    #save_old_password
    add_update_property_file(conf_file_name, OLD_ENCRYPTED_PASSWORD_PARAM, old_enc_password)

    #Update with new pwd
    add_update_property_file(conf_file_name, ENCRYPTED_PASSWORD_PARAM, new_enc_password)

    if conf_file_name == APP_CONF_FILE_NAME :
        logging.info(f"updating mongo password in airflow/workflows-default.json file")
        config_file_path = os.path.join(CONFIG_DIR, "airflow/workflows-default.json")
        with open(config_file_path, 'r') as json_file:
            workflows_default = json.load(json_file)
            workflows_default['mongo']['db']['password'] = new_enc_password
            with open(config_file_path, 'w') as file:
                json.dump(workflows_default, file, indent=2)

def check_update_adapter_source_password():
        conf_file_path = os.path.join(CONFIG_DIR, APP_CONF_FILE_NAME)
        encrypted_password_source_flag_exist = check_property_exist(conf_file_path, OLD_ENCRYPTED_DATA_SOURCE_PASSWORD_PARAM)
        logging.debug(f"encrypted_password_source_flag_exist : {encrypted_password_source_flag_exist}")
        if encrypted_password_source_flag_exist :
                old_encrypted_password = get_param_value(conf_file_path, OLD_ENCRYPTED_DATA_SOURCE_PASSWORD_PARAM)
                new_encrypted_password = get_data_source_password()
                logging.debug(f"old_encrypted_password : {old_encrypted_password}")
                logging.debug(f"new_encrypted_password : {new_encrypted_password}")
                if old_encrypted_password != None and old_encrypted_password == new_encrypted_password :
                    update_adapter_source_password()
        else :
            update_adapter_source_password()

def check_update_mongo_password_presidio():
        conf_file_path = os.path.join(CONFIG_DIR, APP_CONF_FILE_NAME)
        old_mongo_db_password_flag_exist = check_property_exist(conf_file_path, OLD_ENCRYPTED_PASSWORD_PARAM)
        logging.debug(f"old_mongo_db_password_flag_exist {old_mongo_db_password_flag_exist}")
        if old_mongo_db_password_flag_exist :
                old_encrypted_password = get_param_value(conf_file_path, OLD_ENCRYPTED_PASSWORD_PARAM)
                new_encrypted_password = get_param_value(conf_file_path, ENCRYPTED_PASSWORD_PARAM)
                logging.debug(f"old_encrypted_password : {old_encrypted_password}")
                logging.debug(f"new_encrypted_password : {new_encrypted_password}")
                if old_encrypted_password != None and old_encrypted_password == new_encrypted_password :
                    update_mongo_password(APP_CONF_FILE_NAME)
        else :
            update_mongo_password(APP_CONF_FILE_NAME)

def check_update_mongo_password_presidio_ui():
        conf_file_path = os.path.join(CONFIG_DIR, UI_CONF_FILE_NAME)
        old_mongo_db_password_flag_exist = check_property_exist(conf_file_path, OLD_ENCRYPTED_PASSWORD_PARAM)
        logging.debug(f"old_mongo_db_password_flag_exist : {old_mongo_db_password_flag_exist}")
        if old_mongo_db_password_flag_exist :
                old_encrypted_password = get_param_value(conf_file_path, OLD_ENCRYPTED_PASSWORD_PARAM)
                new_encrypted_password = get_param_value(conf_file_path, ENCRYPTED_PASSWORD_PARAM)
                logging.debug(f"old_encrypted_password : {old_encrypted_password}")
                logging.debug(f"new_encrypted_password : {new_encrypted_password}")
                if old_encrypted_password != None and old_encrypted_password == new_encrypted_password :
                    update_mongo_password(UI_CONF_FILE_NAME)
        else :
            update_mongo_password(UI_CONF_FILE_NAME)

def update_password():
        conf_file_path = os.path.join(CONFIG_DIR, APP_CONF_FILE_NAME)
        upgrade_property = 'upgrade.password.encrypted=true'
        upgrade_property_flag_exist = check_property_exist(conf_file_path, upgrade_property)
        logging.debug(f"upgrade_property_flag_exist : {upgrade_property_flag_exist}")

        if upgrade_property_flag_exist :
            check_update_adapter_source_password()
            check_update_mongo_password_presidio()
            check_update_mongo_password_presidio_ui()
        else :
            logging.info(f"Looks like rerun_ueba_server_config.py script has not yet been executed as part of UEBA post upgrade task. \nKindly execute rerun_ueba_server_config.py before running this script")
            print(f"Looks like rerun_ueba_server_config.py script has not yet been executed as part of UEBA post upgrade task. \nKindly execute rerun_ueba_server_config.py before running this script")

version = subprocess.getoutput("rpm -qa | grep rsa-nw-presidio-core | cut -d \'-\' -f5")
if version == '12.5.0.0':
    logging.info("********************* Starting update of encrypted password *************");
    logging.info(f"UEBA version : {version}");
    print(f"UEBA version : {version}");
    update_password()
    print(f"Successfully executed the script.")
else :
    logging.info(f"This script is required only for 12.5.0.0 version")
    print(f"This script is required only for 12.5.0.0 version")