Azure Runbook Python2.7 operations on vms

azure_runbook_operations_on_vms.py

"""
Azure Automation documentation : https://aka.ms/azure-automation-python-documentation
Azure Python SDK documentation : https://aka.ms/azure-python-sdk

requirements.yml

azure_common==1.1.24
azure_mgmt_compute==10.0.0
msrest==0.6.11
msrestazure==0.6.2
typing==3.7.4.1
"""

import re
import sys
import automationassets
from azure.mgmt.compute import ComputeManagementClient
from msrestazure.azure_exceptions import CloudError


def get_automation_runas_credential(runas_connection):
    from OpenSSL import crypto
    import binascii
    from msrestazure import azure_active_directory
    import adal

    # Get the Azure Automation RunAs service principal certificate
    cert = automationassets.get_automation_certificate("AzureRunAsCertificate")
    pks12_cert = crypto.load_pkcs12(cert)
    pem_pkey = crypto.dump_privatekey(crypto.FILETYPE_PEM,pks12_cert.get_privatekey())

    # Get run as connection information for the Azure Automation service principal
    application_id = runas_connection["ApplicationId"]
    thumbprint = runas_connection["CertificateThumbprint"]
    tenant_id = runas_connection["TenantId"]

    # Authenticate with service principal certificate
    resource ="https://management.core.windows.net/"
    authority_url = ("https://login.microsoftonline.com/"+tenant_id)
    context = adal.AuthenticationContext(authority_url)
    return azure_active_directory.AdalAuthentication(
    lambda: context.acquire_token_with_client_certificate(
            resource,
            application_id,
            pem_pkey,
            thumbprint)
    )


def run_command(resource_group_name, name, run_command_parameters, compute_client):
    print("DEBUG: - Run command {} on vm {}".format(run_command_parameters['script'], name))
    poller = compute_client.virtual_machines.run_command(
        resource_group_name,
        name,
        run_command_parameters
    )
    result = poller.result()
    return result


def stop_vm(resource_group_name, name, compute_client):
    print("DEBUG: - Try to stop vm {}".format(name))
    async_vm_stop = compute_client.virtual_machines.power_off(
        resource_group_name, name
    )
    async_vm_stop.wait()
    print("DEBUG: - vm {} was successfully stopped".format(name))


def start_vm(resource_group_name, name, compute_client):
    print("DEBUG: - Try to start vm {}".format(name))
    async_vm_start = compute_client.virtual_machines.start(
        resource_group_name, name
    )
    async_vm_start.wait()
    print("DEBUG: - vm {} was successfully started".format(name))


def run_command_and_stop_vm(database, resource_group_name, vm, compute_client):
    run_command_dict = {
        "dse-cassandra": {
            'command_id': 'RunShellScript',
            'script': [
                'nodetool drain && systemctl stop dse.service && systemctl stop datastax-agent.service'
            ]
        },
        "dse-opscenter": {
            'command_id': 'RunShellScript',
            'script': [
                'systemctl stop opscenter.service'
            ]
        },
        "elasticsearch": {
            'command_id': 'RunShellScript',
            'script': [
                'systemctl stop elasticsearch.service'
            ]
        }
    }
    print("DEBUG: - Database {} will be stopped before the shutdown of vm {}".format(database, vm.name))
    try:
        result = run_command(resource_group_name, vm.name, run_command_dict[database], compute_client)
    except CloudError as err:
        if "The operation requires the VM to be running (or set to run)" in str(err):
            print("DEBUG: - vm {} has been already stopped".format(vm.name))
            return
        raise
    stdout_error_regex = r'^.*?\n\[stdout\]\n(.*)?(?:\n)?\n\[stderr\]\n(.*)?(?:\n)?$'
    match = re.search(stdout_error_regex, result.value[0].message)
    error_msg = match.group(2)
    if not error_msg:
        print("vm will {} be stopped...".format(vm.name))
        stop_vm(resource_group_name, vm.name, compute_client)
    else:
        print("ERROR - Database {} was not stopped gracefully on vm {}".format(database, vm.name))
        print("ERROR - error={} while trying to stop database".format(error_msg))
        print("ERROR - vm {} will not be stopped".format(vm.name))


def stop_group_of_vms(vms, resource_group_name, compute_client):
    for vm in vms:
        print('AZUREVM operation for vm={}'.format(vm.name))
        database = vm.tags.get('database', None)
        if database:
            if database in ("dse-cassandra", "dse-opscenter", "elasticsearch"):
                run_command_and_stop_vm(database, resource_group_name, vm, compute_client)
            else:
                print("ERROR - Shutdown vm {} with database type=={} not supported.Instead Force Shutdown".format(database, vm.name))
                stop_vm(resource_group_name, vm.name, compute_client)
        else:
            stop_vm(resource_group_name, vm.name, compute_client)


def start_group_of_vms(vms, resource_group_name, compute_client):
    for vm in vms:
        print('AZUREVM operation for vm={}'.format(vm.name))
        start_vm(resource_group_name, vm.name, compute_client)


def main():
    if len(sys.argv) >= 3:
        action = sys.argv[1]
        resource_group_name = str(sys.argv[2])
        tags = str(sys.argv[3])
        tags_regex = r'^([a-zA-Z]+=\w+)(,[a-zA-Z]+=\w+)*(?<!,)$'

        if not re.search(tags_regex, tags):
            error_msg = "Tags must match this regex {}\nEg:\n\t tag1=3,tag2=4".format(tags_regex)
            print(error_msg)
            raise Exception(error_msg)
    else:
        error_msg = "Positional parameters action, resource_group_name and tags are required..."
        print(error_msg)
        raise Exception(error_msg)
    tags = tags.split(",")
    runas_connection = automationassets.get_automation_connection("AzureRunAsConnection")
    azure_credential = get_automation_runas_credential(runas_connection)
    compute_client = ComputeManagementClient(
        azure_credential,
        str(runas_connection["SubscriptionId"])
    )

    tag_regex = r'([a-zA-Z]+)=(\w+)'
    tags_dict_to_check = {}
    for tag in tags:
        match = re.search(tag_regex, tag)
        tags_dict_to_check[match.group(1)] = match.group(2)

    def tags_is_present(tags_dict, tags_dict_to_check):
        return tags_dict and all([tags_dict.get(tag_name) == tag_value for tag_name, tag_value in tags_dict_to_check.items()])

    tagged_vms = [vm for vm in compute_client.virtual_machines.list(resource_group_name) if tags_is_present(vm.tags, tags_dict_to_check)]
    if action == "start_vm":
        start_group_of_vms(tagged_vms, resource_group_name, compute_client)
    elif action == "stop_vm":
        stop_group_of_vms(tagged_vms, resource_group_name, compute_client)


if __name__ == '__main__':
    main()