python wrapper for azure cli

azure_cli.py

# -*- coding: utf-8 -*-
import subprocess
import logging
import tempfile
import json
import datetime

azure_cli = "\"C:\\Program Files (x86)\\Microsoft SDKs\\Azure\\CLI\\wbin\\azure.cmd\""

def exec_azure_cli_cmd(cmd, skip_start_lines, input = ""):
    logging.info("Executing in azure: %s" % cmd)
    test = subprocess.Popen("%s %s" % (azure_cli, cmd), stdout=subprocess.PIPE, stdin=subprocess.PIPE)
    out, err = test.communicate(input=input)
    lines = out.splitlines()
    if not lines[-1].endswith("OK"):
        raise Exception("Error in executing '%s' command: \n%s" % (cmd, lines[-1]))
    return lines[skip_start_lines:-1]

def vm_names(service_name):
    res = []
    lines = exec_azure_cli_cmd("vm list",4)
    for line in lines:
        parts = [word for word in line.split("  ") if word <> ""][1:6]
        name, status, location, dns_name, ip_address = parts
        if name.startswith(service_name):
            res.append(name)
    return res

def get_vm_endpoints(vm_name):
    res = []
    lines = exec_azure_cli_cmd("vm endpoint list %s" % vm_name, 4)
    for line in lines:
        parts = [word for word in line.split("  ") if word <> ""][1:8]
        name, protocol, public_port, private_port, virtual_ip, enable_direct_server_return, Load_balanced = parts
        res.append({
            "name" : name,
            "protocol" : protocol,
            "public_port" : int(public_port),
            "private_port" : private_port,
            "virtual_ip" : virtual_ip,
            "enable_direct_server_return" : enable_direct_server_return,
            "Load_balanced" : Load_balanced
        })
    return res

def get_vm_ssh_addr(vm_name):
    endpoints = get_vm_endpoints(vm_name)
    for d in endpoints:
        if d["name"] == "ssh":
            return d["virtual_ip"], d["public_port"]

def delete_vm(vm_name):
    exec_azure_cli_cmd("vm delete %s" % vm_name, 0, input = "y\n")


def create_main_vm(name, image, user, password, ssh_port = 22, location = "West Europe", vm_size = "Standard_D1"):
    exec_azure_cli_cmd("vm create --location \"%s\" --ssh --vm-size %s %s %s %s %s" % (location, vm_size, name, image, user, password), 0)


def create_instance_vm(name, image, user, password, ssh_port = 22, vm_size = "Standard_D1"):
    exec_azure_cli_cmd("vm create --connect --ssh %d --vm-size %s %s %s %s %s" % (ssh_port, vm_size, name, image, user, password), 0)

def get_next_ssh_port(name):
    res = 12344
    vms = vm_names(name)
    for vm in vms:
        addr, port = get_vm_ssh_addr(vm)
        res  = max(port, res)
    return res + 1

def add_instance(name, image, user, password, vm_size = "Standard_D1", ssh_port = None):
    if not ssh_port:
        logging.info("calculating new ssh port")
        ssh_port = get_next_ssh_port(name)
    create_instance_vm(name, image, user, password, ssh_port, vm_size)

def remove_instance(name):
    vms = vm_names(name)
    if len(vms) > 1:
        logging.info("removing vm %s" % vms[-1])
        delete_vm(vms[-1])

def reset_password(vm_name, user, password, exp_days = 365):
    temp = tempfile.TemporaryFile()
    try:
        exp_date = datetime.date.today() + datetime.timedelta(days=exp_days)
        opts = {
            "username":user,
            "password":password,
            "expiration":str(exp_date),
        }
        s = json.dumps(opts)
        logging.info(s)
        temp.write(s)
        temp.flush()
        logging.info("created opts file %s" % temp.name)
        cmd = "vm extension set %s VMAccessForLinux Microsoft.OSTCExtensions 1.* –-private-config-path %s" % (vm_name, temp.name)
        exec_azure_cli_cmd(cmd, 0)
    finally:
        logging.info("opts file cleanup")
        temp.close()