Script to regenerate all trust relationships for all regions

trust_relationships.py

#!/usr/bin/env python3

# Manages trust relationships across all regions. The script is idempotent, it will only update the trust relationships
# which need to be updated and ignore the ones which are already ok.
#
# Example of usage:
#   EC_CONFIG=qa ./bin/trust_relationships.py --region us-east-1
#
# Requirements:
#   ecl - http://github.com/elastic/cloud-cli
#
# Parameters:
#   region - (optional) Used when adding a new region to an environment with multiple regions. It will make all existing
#            regions trust this new region and this new region will trust all existing regions.
#            The list of regions can also befiltered using EC_REGIONS.
#            If not region is passed, then all regions will trust all existing regions.
#
#   dry_run - (optional) if set to any true, don't actually add the trust relationships
#
# You will need to run `ecl auth --config EC_CONFIG` to provide TOTP access to ecl, or use credentials that bypass TOTP
#
#

from distutils.version import LooseVersion
from subprocess import Popen, PIPE, check_output, CalledProcessError, run

import argparse
import json
import os
import re
import shutil
import tempfile


failcolor = '\033[91m'
regularcolor = '\033[0m'
base_url = '/api/v1/regions/{}/platform/configuration/trust-relationships'

def eclv1():
    command = ['ecl', 'version']
    version, error = run_ecl_command(command)
    return 'ecl 1.' in version

def get_regions():
    if eclv1():
        command = ['ecl', 'platform', 'region', 'list', '--output', 'json']
    else:
        command = ['ecl', 'region', 'list', '--output', 'json']

    regions, error = run_ecl_command_json(command)
    return regions

def download_local_trust_relationship(region):
    command = ['ecl', 'api', (base_url + '/local?include_certificate=true').format(region), '--output', 'json']

    response, error = run_ecl_command_json(command)
    return response["public_ca_cert"]

def run_ecl_command_json(command):
    output, error = run_ecl_command(command)
    return json.loads(output), error

def run_ecl_command(command):
    try:
        p = Popen(command, stdout=PIPE, stderr=PIPE)
        output, error = p.communicate()

        if p.returncode != 0:
            if "needs elevated permissions" in error.decode():
                print("{}ERROR:{} This script needs elevated permissions. "
                      "You can run ecl elevate-permissions --config environment \n{}".format(failcolor, regularcolor, error.decode()))
            else:
                print("{}ERROR:{} Problem running command '{}'\n{}".format(failcolor, regularcolor, ' '.join(command), error.decode()))
            exit(code=2)

    except FileNotFoundError as e:
        print("{}ERROR:{} Problem running command '{}'\n"
              "Make sure `ecl` is installed (ref: https://github.com/elastic/cloud-cli)"
              " and configured properly\n{}".format(failcolor, regularcolor, ' '.join(command), e))

        exit(code=2)
    return output if output is None else output.decode(), error if error is None else error.decode()

def download_regions_certificates(regions):
    all_certificates = {}
    for region in regions:
        print("Downloading certificate from region: [{}]".format(region))
        all_certificates[region] = download_local_trust_relationship(region)
    return all_certificates

def download_regions_trust_relationships(region):
    command = ['ecl', 'api', (base_url + '?include_certificate=true').format(region), '--output', 'json']
    response, error = run_ecl_command_json(command)

    return response["trust_relationships"]

def create_trust_relationship(from_region, to_region, certificate):
    trust_relationship = {
        "name" : to_region,
        "trust_by_default" : True,
        "public_ca_cert" : certificate
    }

    command = ['ecl', 'api', '-X', 'POST', base_url.format(from_region), '--output', 'json', '--header', 'Content-Type: application/json', '--data', json.dumps(trust_relationship)]
    print(" - creating a trust relationship in {} to trust {}.".format(from_region, to_region))
    created, error = run_ecl_command_json(command)

    if error:
        print("- {}ERROR:{} Unable to create trust relationship [{}] in [{}]. \n{}".format(failcolor, regularcolor, to_region, from_region, error))

def update_trust_relationship(tid, from_region, to_region, certificate):
    trust_relationship = {
        "name" : to_region,
        "trust_by_default" : True,
        "public_ca_cert" : certificate
    }

    command = ['ecl', 'api', '-X', 'PUT', (base_url + '/{}').format(from_region, tid), '--output', 'json', '--header', 'Content-Type: application/json', '--data', json.dumps(trust_relationship)]
    print(" - updating the existing trust relationship in [{}] to trust [{}].".format(from_region, to_region))
    created, error = run_ecl_command_json(command)

    if error:
        print("- {}ERROR:{} Unable to update trust relationship [{}] in [{}]. \n{}".format(failcolor, regularcolor, to_region, from_region, error))

def maybe_upload_trust_relationship(from_region, to_region, certificate, dry_run, existing_trust_relationships):
    existing = next((x for x in existing_trust_relationships if x["name"] == to_region), None)

    if existing is None:
        if dry_run:
            print(" - this would try to create a trust relationship in [{}] to trust [{}].".format(from_region, to_region))
        else:
            create_trust_relationship(from_region, to_region, certificate)
    else:
        if certificate == existing["public_ca_cert"]:
            print(" - skipping [{}] since the trust relationships is already up to date.".format(to_region))
        elif dry_run:
            print(" - this would try to update the trust relationship in [{}] to trust [{}].".format(from_region, to_region))
        else:
            update_trust_relationship(existing["id"], from_region, to_region, certificate)

def download_region_local_trust_relationships(maybe_region, dry_run):
    # 1. Obtain all regions
    regions = get_regions()

    # 2. Obtain certificates from all regions
    all_certificates = download_regions_certificates(regions)

    # 3. Upload certificates to all regions (or the specified one)
    for region in regions:
        existing_trust_relationships = download_regions_trust_relationships(region)

        print("Uploading certificates to region: [{}]".format(region))

        if (maybe_region is None) or (maybe_region == region):
            for r, cert in all_certificates.items():
                if r != region:
                    maybe_upload_trust_relationship(region, r, cert, dry_run, existing_trust_relationships)
        else:
            maybe_upload_trust_relationship(region, maybe_region, all_certificates[maybe_region], dry_run, existing_trust_relationships)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Manages trust relationships across all regions')
    parser.add_argument('--region', type=str, nargs='?', help='adds a new region to trust all existing regions and be trusted by them')
    parser.add_argument('--dry_run', type=bool, nargs='?', default=False, help='if set to any true, dont actually add the trust relationships')
    args = parser.parse_args()

    if args.region is None:
        print("Propagating trust relationships across all existing regions. (Dry Run: {})".format(args.dry_run))
    else:
        print("Adding trust relationships from and to region [{}]. (Dry Run: {})".format(args.region, args.dry_run))

    download_region_local_trust_relationships(args.region, args.dry_run)