Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?

this python script can be used to automate updating AZ load balancer backend pool with the vms you like. this is knda custom as the dcos_vm and maint_vm are what this script is used to switch between .

#!/usr/bin/env python3
import argparse
import configparser
import logging.config
import os
import sys

# Setup..
from azure.common.credentials import ServicePrincipalCredentials
from azure.mgmt.network import (
    NetworkManagementClient,
)
from azure.mgmt.compute import (
    ComputeManagementClient
)
import requests
from requests import Request, Session

logging.basicConfig(level=logging.INFO, format='%(asctime)s %(message)s')
log = logging.getLogger(__name__)

parser = argparse.ArgumentParser('RMS(one) Load Balancer update')
parser.add_argument('-s', '--stack', action='store', dest='stack', metavar='', required=True, help='Name of stack to update LB')
parser.add_argument('-a', '--action', action='store', dest='action', metavar='', required=True, help='Action for LB', choices=['create-maint', 'delete-maint'])
args = parser.parse_args()
stack = args.stack
action = args.action
base_url = 'https://management.azure.com/'

def get_token_from_client_credentials(endpoint, client_id, client_secret):
    payload = {
        'grant_type': 'client_credentials',
        'client_id': client_id,
        'client_secret': client_secret,
        'resource': 'https://management.core.windows.net/',
    }
    #TODO add back in verify for non-fiddler
    #NOTE add Verify=False when going via a proxy with fake cert / fiddler
    response = requests.post(endpoint, data=payload).json()
    return response['access_token']

def get_virtual_machine(compute_client, resource_group_name, vm_name):
    """
    :param resource_group_name: str
    :param vm_name: str
    :return: azure.mgmt.compute.VirtualMachine
    """
    virtual_machine = compute_client.virtual_machines.get(resource_group_name, vm_name)
    logging.info('using virtual machine id: %s', virtual_machine.id)
    return virtual_machine

def get_network_interface_ip_configuration(network_client, resource_group_name, network_interface_name):
    network_interface = network_client.network_interfaces.get(resource_group_name, network_interface_name)
    return network_interface
    #for ipconfig in network_interface.network_interface.ip_configurations:
    #    return ipconfig

def get_virtual_machine_network_interface(compute_client, network_client, resource_group_name, virtual_machine_name):
    virtual_machine = get_virtual_machine(compute_client, resource_group_name, virtual_machine_name)
    for profile in virtual_machine.network_profile.network_interfaces:
        print(profile.id)
        nic_uri = profile.id

    #network_interface = get_network_interface(resource_group_name)
    label = os.path.basename(os.path.normpath(nic_uri))
    logging.info('nic on vm to use is: %s', label)

    network_interface = get_network_interface_ip_configuration(network_client, resource_group_name, label)
    logging.info('nic id is: %s', network_interface.id)
    return network_interface

def build_request(vm_object, nic_object, load_balancer=None):
    """
    :param vm_object : azure.mgmt.compute.VirtualMachine
    :param nic_object : azure.mgmt.network.networkresourceprovider.NetworkInterface
    :param load_balancer : azure.mgmt.network.LoadBalancer
    :return: dict
    """
    if load_balancer is None:
        backend_pool = []
    else:
        backend_pool = [{'id' : load_balancer.backend_address_pools[0].id}]

    request = {
        'properties': {
            'virtualMachine' : {
                'id' : vm_object.id
                },
            'ipConfigurations' : [{ #may have to build by hand
                'properties' : {
                    'loadBalancerBackendAddressPools' : backend_pool,
                    'subnet' : {
                        'id' :  nic_object.ip_configurations[0].subnet.id
                        }
                    },
                'name' : nic_object.ip_configurations[0].name,
                'id' : nic_object.ip_configurations[0].id
            }]
        },
        'id' : nic_object.id,
        'name' : nic_object.name,
        'location' : vm_object.location,
        'type' : 'Microsoft.Network/networkInterfaces'
        }

    return request

def send_loadbalancer_request(payload, auth_token, resource_id, max_retries=20):
    endpoint = base_url + resource_id + '?api-version=2019-06-01'
    header = {'Authorization' : 'Bearer ' + auth_token}
    while max_retries > 0:
        session = Session()
        request = Request('PUT', endpoint, json=payload, headers=header)
        prepared = session.prepare_request(request)

        log.debug('raw body sent')
        log.debug(prepared.body)

        response = session.send(prepared)
        print(response.status_code)
        print(response.text)
        if response.status_code == 200:
            break
        elif response.status_code == 429:
            log.info('retrying an HTTP send due to 429 retryable response')
            log.info('this will be try# %s', max_retries)
        max_retries = max_retries - 1
    return response

def main():
    ini_config = configparser.ConfigParser()
    ini_config.read('../rms_roc/azure.ini')
    stack_data = ini_config[stack]
    tenant_id = stack_data['tenant']
    client_id = stack_data['client_id']
    client_secret = stack_data['secret']
    sub_id = stack_data['subscription_id']
    endpoint = 'https://login.microsoftonline.com/' + tenant_id + '/oauth2/token'
    auth_token = get_token_from_client_credentials(endpoint, client_id, client_secret)
    # now the Azure management credentials
    credentials = ServicePrincipalCredentials(client_id=client_id,
                                              secret=client_secret,
                                              tenant=tenant_id)
    # now the specific compute, network resource type clients
    compute_client = ComputeManagementClient(credentials, sub_id)
    network_client = NetworkManagementClient(credentials, sub_id)
    # Resources
    resources = {"vmnames":{"dcos_vms": [f"{stack}-dcos-extpublicslave1", f"{stack}-dcos-extpublicslave2"], "maint_vm" : f"{stack}-maintpage"},
                 "vmResourceGroup": f"{stack}-dcos",
                 "netResourceGroup": f"{stack}-Network-Infrastructure",
                 "loadBalancerName": f"{stack}-dcos-extpublicslave",
                 "subnetName": f"{stack}-SubNet1",
                 "virtualNetworkName" : f"{stack}-Vnet1"}

    #TODO modify this to mach your specific settings
    vm_resource_group = resources['vmResourceGroup']
    load_balancer_name = resources['loadBalancerName']
    #TODO - end - only the "above" should need to change.
    dcos_vms_res = {}
    maint_vm_res = {}
    for dcos_vm_name in resources["vmnames"]["dcos_vms"]:
        dcos_vms_res[dcos_vm_name] = {"vm":"", "nic":""}
        dcos_vms_res[dcos_vm_name]["vm"] = compute_client.virtual_machines.get(vm_resource_group, dcos_vm_name)
        dcos_vms_res[dcos_vm_name]["nic"] = get_virtual_machine_network_interface(compute_client, network_client, vm_resource_group, dcos_vm_name)
    maint_vm_res["vm"] = compute_client.virtual_machines.get(vm_resource_group, resources["vmnames"]["maint_vm"])
    maint_vm_res["nic"] = get_virtual_machine_network_interface(compute_client, network_client, vm_resource_group, resources["vmnames"]["maint_vm"])
    #the load balancer
    load_balancer = network_client.load_balancers.get(vm_resource_group, load_balancer_name)

    # running maint on/off action
    if action == "create-maint":
        log.info("Running maintenance ON ")
        maint_vm_lb = load_balancer
        dcos_vm_lb = None
    elif action == "delete-maint":
        log.info("Running maintenance OFF ")
        maint_vm_lb = None
        dcos_vm_lb = load_balancer
    else:
        log.error("ERROR: invalid action specified")
        sys.exit(1)

    maint_lb_request = build_request(maint_vm_res["vm"], maint_vm_res["nic"], maint_vm_lb)
    send_loadbalancer_request(maint_lb_request, auth_token, maint_vm_res["nic"].id)
    for dcos_vm in dcos_vms_res:
        dcos_lb_request = build_request(dcos_vms_res[dcos_vm]["vm"], dcos_vms_res[dcos_vm]["nic"], dcos_vm_lb)
        send_loadbalancer_request(dcos_lb_request, auth_token, dcos_vms_res[dcos_vm]["nic"].id)


if __name__ == "__main__":
    main()

run as python lb.py -s whatever_name -a create-maint i follow naming convention for my resource so stack has to be passed you can totally remove the inputs as u like . also create-maint and delete-maint is used to switch back and forth between the two vms. tested with python 3 and python 2

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.