generate-security-policy.py

import subprocess
import base64
import os
import argparse
import json


def is_valid_json(file_path):
    try:
        with open(file_path, 'r') as file:
            json.load(file)
        return True
    except (json.JSONDecodeError, FileNotFoundError):
        return False

def main():
    # Create argument parser
    parser = argparse.ArgumentParser(description='Wrapper around az confcom tooling to support CS2 ContainerGroupProfiles and generate secuirty policy.')

    # Add argument for the JSON file path using --filename
    parser.add_argument('--filename', required=True, help='Specify path to the ARM JSON file')

    # Parse the command-line arguments
    args = parser.parse_args()

    # Check if the provided file is a valid JSON file
    if args.filename and is_valid_json(args.filename):
        # Open the JSON file in read-only mode
        with open(args.filename, 'r') as file:
            original_json_data = json.load(file)
    else:
        print("Error: Please provide a valid JSON file using --filename option.")

    # Replace containerGroupProfiles with containerGroups so that az confcom can generate the policy
    for each in original_json_data["resources"]:
        if each["type"] == "Microsoft.ContainerInstance/containerGroupProfiles":
            each["type"] = "Microsoft.ContainerInstance/containerGroups"
            break

    # Read the original JSON file into a temporary modified file
    modified_json_path = "modified.json"
    with open(modified_json_path, 'w') as file:
        json.dump(original_json_data, file)

    # Specify the az confcom command to execute
    command = ["az", "confcom", "acipolicygen", "-a", modified_json_path, "--debug-mode"]

    # Execute the command
    try:
        subprocess.run(command, check=True)
    except subprocess.CalledProcessError as e:
        print("Please fix the issue highlighted above and rerun the command.")
    else:
        # copy the generated policy from the modified json
        print("Fetching cce policy from modified JSON...")
        with open(modified_json_path, 'r') as file:
            modified_json_data = json.load(file)
            for each in modified_json_data["resources"]:
                if each["type"] == "Microsoft.ContainerInstance/containerGroups":
                    cce_policy = each["properties"]["confidentialComputeProperties"]["ccePolicy"]
                    break

        # copy the policy into input ARM template
        print("Copying cce policy into original JSON...")
        for each in original_json_data["resources"]:
            if each["type"] == "Microsoft.ContainerInstance/containerGroups":
                each["type"] = "Microsoft.ContainerInstance/containerGroupProfiles"
                each["properties"]["confidentialComputeProperties"]["ccePolicy"] = cce_policy
                break

        # Save the input ARM template
        with open(args.filename, "w") as jsonFile:
            json.dump(original_json_data, jsonFile, indent=4)

        print(f"Successfully executed command and updated JSON. Removed {modified_json_path}.")

    # Remove the temporary modified JSON file
    os.remove(modified_json_path)


if __name__ == "__main__":
    main()