encrypt_ansible_yaml_string.py

import os
import yaml

from ansible.parsing.vault import VaultLib, VaultSecret
from ruamel.yaml import YAML
from ruamel.yaml.scalarstring import LiteralScalarString


def custom_literal_presenter(dumper, data):
    return dumper.represent_scalar('tag:yaml.org,2002:str', data, style='|')


def encrypt_yaml_keys(input_file, output_file, vault_password):
    with open(input_file, 'r') as f:
        data = yaml.safe_load(f)

    vault = VaultLib([(b"default", VaultSecret(vault_password.encode()))])

    encrypted_data = {}
    for key, value in data.items():
        encrypted_value = vault.encrypt(value).decode()
        encrypted_data[key] = LiteralScalarString(encrypted_value)

    with open(output_file, 'w') as f:
        yamlru = YAML()
        yamlru.representer.add_representer(LiteralScalarString, custom_literal_presenter)
        yamlru.explicit_start = True
        yamlru.indent(mapping=10, sequence=4, offset=2)
        yamlru.dump(encrypted_data, f)


def get_yaml_files_from_directory(directory_path):
    yaml_files = []
    for file in os.listdir(directory_path):
        if file.endswith("-secrets.yaml") or file.endswith("-secrets.yml"):
            yaml_files.append(os.path.join(directory_path, file))
    return yaml_files


def get_output_file_path(input_file_path, suffix):
    dirname = os.path.dirname(input_file_path)
    basename = os.path.basename(input_file_path)
    filename, ext = os.path.splitext(basename)
    return os.path.join(dirname, f"{filename}_{suffix}{ext}")


if __name__ == "__main__":
    directory_path = "automation/ansible/inventory/group_vars"
    ansible_vault_password = "<vault_pass>"
    output_suffix = "encrypted"

    yaml_files = get_yaml_files_from_directory(directory_path)

    for file_path in yaml_files:
        output_file_path = get_output_file_path(file_path, output_suffix)
        encrypt_yaml_keys(file_path, output_file_path, ansible_vault_password)