udp_experimental_recorder/mumble_server_config.py

# NOTE:


import base64
import io
import shlex
from datetime import datetime

import paramiko
import yaml

# Servers to update
SERVERS = [
    "172.16.10.123",
    "172.16.10.110",
    "172.16.10.107", 
]

# SSH configuration
USERNAME = "cezen"
PASSWORD = "17lamborghini"
REMOTE_YAML = "/home/cezen/Desktop/mumble-server/docker-compose.yml"
BACKUP_YAML = REMOTE_YAML + ".bak"
VERSION_DIR = "/home/cezen/Desktop/mumble-server/config_versions"
VERSION_KEEP_LIMIT = 10
VERSION_FILE_PREFIX = "docker-compose"

# Service whose environment section will be edited
SERVICE_NAME = "mumble"


def run_sudo_command(ssh, command):
    full_command = f"echo '{PASSWORD}' | sudo -S {command}"
    stdin, stdout, stderr = ssh.exec_command(full_command)
    exit_code = stdout.channel.recv_exit_status()

    if exit_code != 0:
        raise RuntimeError(stderr.read().decode("utf-8"))

    return stdout.read().decode("utf-8") if stdout.readable() else ""


def read_remote_file(ssh, path):
    return run_sudo_command(ssh, f'cat "{path}"')


def write_remote_file(ssh, path, content):
    # Encode content to base64 to avoid all shell quoting/escaping issues
    encoded = base64.b64encode(content.encode("utf-8")).decode("ascii")
    run_sudo_command(ssh, f"bash -c 'echo {encoded} | base64 -d > \"{path}\"'")


def ensure_version_dir(ssh):
    run_sudo_command(ssh, f"mkdir -p {shlex.quote(VERSION_DIR)}")


def write_config_version(ssh, timestamp, content):
    ensure_version_dir(ssh)
    version_path = f"{VERSION_DIR}/{VERSION_FILE_PREFIX}_{timestamp}.yml"
    write_remote_file(ssh, version_path, content)
    prune_config_versions(ssh)
    return version_path


def prune_config_versions(ssh):
    keep_from = VERSION_KEEP_LIMIT + 1
    cleanup_script = (
        f"cd {shlex.quote(VERSION_DIR)} 2>/dev/null || exit 0\n"
        f"ls -1t {VERSION_FILE_PREFIX}_*.yml 2>/dev/null "
        f"| tail -n +{keep_from} "
        "| xargs -r rm --"
    )
    run_sudo_command(ssh, f"bash -c {shlex.quote(cleanup_script)}")


def backup_remote_file(ssh):
    run_sudo_command(ssh, f'cp "{REMOTE_YAML}" "{BACKUP_YAML}"')


def parse_environment(env_list):
    """
    Converts:
        ["KEY1=value1", "KEY2=value2"]
    To:
        [("KEY1", "value1"), ("KEY2", "value2")]
    """
    result = []
    for item in env_list:
        if isinstance(item, str) and "=" in item:
            key, value = item.split("=", 1)
            result.append((key, value))
    return result


def format_environment(pairs):
    """
    Converts:
        [("KEY1", "value1"), ("KEY2", "value2")]
    To:
        ["KEY1=value1", "KEY2=value2"] 
    """
    return [f"{key}={value}" for key, value in pairs]


def load_environment_from_server(host, password):
    print(f"Connecting to {host}...")

    ssh = paramiko.SSHClient()
    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())

    try:
        ssh.connect(hostname=host, username=USERNAME, password=password, timeout=10)

        content = read_remote_file(ssh, REMOTE_YAML)
        data = yaml.safe_load(content)

        env_list = data["services"][SERVICE_NAME].get("environment", [])
        env_pairs = parse_environment(env_list)

        return ssh, data, env_pairs, content

    except Exception:
        ssh.close()
        raise


def save_environment_to_server(ssh, data, env_pairs):
    data["services"][SERVICE_NAME]["environment"] = format_environment(env_pairs)

    # Create backup first
    backup_remote_file(ssh)

    # Serialize YAML
    output = io.StringIO()
    yaml.dump(data, output, default_flow_style=False, sort_keys=False)

    # Write updated file
    write_remote_file(ssh, REMOTE_YAML, output.getvalue())


def save_versions_after_success(previous_configs, timestamp):
    version_results = {}

    for host, content in previous_configs.items():
        try:
            ssh = paramiko.SSHClient()
            ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            ssh.connect(hostname=host, username=USERNAME, password=PASSWORD, timeout=10)
            try:
                version_results[host] = write_config_version(ssh, timestamp, content)
            finally:
                ssh.close()
        except Exception as e:
            version_results[host] = str(e)

    return version_results


def edit_environment(env_pairs, json_data):
    print(env_pairs,"ewejnwejnwejwej")
    updated = []
    for key, value in env_pairs:
        new_value = json_data.get(key, "")
        if new_value == "":
            updated.append((key, value))   # keep original
        else:
            updated.append((key, new_value))
    return updated


def run(json_data, servers=None):
    target_servers = list(servers or SERVERS)
    if not target_servers:
        return {
            "status": "error",
            "message": "Select at least one VM to update.",
            "results": {},
        }

    first_host = target_servers[0]

    try:
        ssh, data, env_pairs, _ = load_environment_from_server(first_host, PASSWORD)
        ssh.close()
    except Exception as e:
        return {"status": "error", "message": f"Failed to connect to {first_host}: {str(e)}"}

    updated_env = edit_environment(env_pairs, json_data)

    results = {}
    previous_configs = {}
    for host in target_servers:
        try:
            ssh, server_data, _, previous_content = load_environment_from_server(host, PASSWORD)
            try:
                save_environment_to_server(ssh, server_data, updated_env)
                results[host] = "success"
                previous_configs[host] = previous_content
            finally:
                ssh.close()
        except Exception as e:
            results[host] = str(e)

    if any(result != "success" for result in results.values()):
        return {
            "status": "error",
            "message": "Config was not saved successfully to all selected VMs. No version files were created.",
            "results": results,
        }

    timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
    version_results = save_versions_after_success(previous_configs, timestamp)

    if any(
        not str(result).startswith(f"{VERSION_DIR}/")
        for result in version_results.values()
    ):
        return {
            "status": "error",
            "message": "Config was saved to all selected VMs, but one or more version files could not be created.",
            "results": results,
            "version_results": version_results,
        }

    return {
        "status": "success",
        "results": results,
        "version_results": version_results,
    }