DEV Community

Imran Hayder
Imran Hayder

Posted on

How to update Azure load balancer backend pool via a python script

Some time back ago I was working on a very interesting problem in Azure. Lets says we have two VMs added in a load-balancer. Now the ask was to do some maintenance on one of the vm while other was still in the load-balancer.

The requirements were:

  • take the VM out of load-balancer for maintenance
  • do some work on it
  • add it back to the load-balancer

Now there wasn't any automated way to do it via Azure so I came up with following python script. Hopefully someone out there in same predicament finds it useful :)

#!/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('~/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()

The script is also added on my github gist so feel free to check out.

How to run the script

Save the script as lb.py and run as:

python lb.py -s stack_name -a create-maint

The stack_name is whatever I used in azure.ini file to get the credential from . The options create-maint and delete-maint are used to switch back and forth between the two vms. tested with python 3 and python 2
Note I follow naming convention for my resource so stack has to be passed you can totally remove the inputs as you like .

Top comments (0)