account_setup.py

'''
    doAzureParallel Getting Started script
'''
import sys
import threading
import time
import uuid
import json
import string
import random
from azure.common import credentials
from azure.graphrbac import GraphRbacManagementClient
from azure.graphrbac.models import ApplicationCreateParameters, PasswordCredential, ServicePrincipalCreateParameters
from azure.graphrbac.models.graph_error import GraphErrorException
from azure.mgmt.authorization import AuthorizationManagementClient
from azure.mgmt.batch import BatchManagementClient
from azure.mgmt.batch.models import AutoStorageBaseProperties, BatchAccountCreateParameters
from azure.mgmt.network import NetworkManagementClient
from azure.mgmt.network.models import AddressSpace, Subnet, VirtualNetwork
from azure.mgmt.resource import ResourceManagementClient
from azure.mgmt.storage import StorageManagementClient
from azure.mgmt.storage.models import Kind, Sku, SkuName, StorageAccountCreateParameters
from azure.mgmt.subscription import SubscriptionClient
from datetime import datetime, timezone
from msrestazure.azure_cloud import AZURE_PUBLIC_CLOUD
from msrestazure.azure_exceptions import CloudError


class AccountSetupError(Exception):
    pass


class DefaultSettings():
    authentication = "sharedkey"
    resource_group = 'doazp'
    storage_account = 'doazpstorage'
    batch_account = 'doazpbatch'
    application_name = 'doazpapp'
    application_credential_name = 'doazpappcredential'
    service_principal = 'doazpsp'
    region = 'westus'


def create_resource_group(credentials, subscription_id, **kwargs):
    """
        Create a resource group
        :param credentials: msrestazure.azure_active_directory.AdalAuthentication
        :param subscription_id: str
        :param **resource_group: str
        :param **region: str
    """
    resource_client = ResourceManagementClient(credentials, subscription_id)
    resource_client.resource_groups.list()
    for i in range(3):
        try:
            resource_group = resource_client.resource_groups.create_or_update(
                resource_group_name=kwargs.get("resource_group", DefaultSettings.resource_group),
                parameters={
                    'location': kwargs.get("region", DefaultSettings.region),
                }
            )
        except CloudError as e:
            if i == 2:
                raise AccountSetupError(
                    "Unable to create resource group in region {}".format(kwargs.get("region", DefaultSettings.region)))
            print(e.message)
            print("Please try again.")
            kwargs["resource_group"] = prompt_with_default("Azure Region", DefaultSettings.region)
    return resource_group.id

def delete_resource_group(credentials, subscription_id, resource_group):
    """
        Delete a resource group
        :param credentials: msrestazure.azure_active_directory.AdalAuthentication
        :param subscription_id: str
        :param resource_group: str
    """
    resource_client = ResourceManagementClient(credentials, subscription_id)
    resource_client.resource_groups.list()
    delete_async_operation = resource_client.resource_groups.delete(
        resource_group_name=resource_group
    )
    delete_async_operation.wait()

def create_storage_account(credentials, subscription_id, **kwargs):
    """
        Create a Storage account
        :param credentials: msrestazure.azure_active_directory.AdalAuthentication
        :param subscription_id: str
        :param **resource_group: str
        :param **storage_account: str
        :param **region: str
    """
    storage_management_client = StorageManagementClient(credentials, subscription_id)
    storage_account = storage_management_client.storage_accounts.create(
        resource_group_name=kwargs.get("resource_group", DefaultSettings.resource_group),
        account_name=kwargs.get("storage_account", DefaultSettings.storage_account),
        parameters=StorageAccountCreateParameters(
            sku=Sku(SkuName.standard_lrs),
            kind=Kind.storage,
            location=kwargs.get('region', DefaultSettings.region)
        )
    )
    return storage_account.result().id

def storage_account_get_keys(credentials, subscription_id, **kwargs):
    """
        get Storage account keys
        :param credentials: msrestazure.azure_active_directory.AdalAuthentication
        :param subscription_id: str
        :param **resource_group: str
        :param **storage_account: str
        :param **region: str
    """
    storage_management_client = StorageManagementClient(credentials, subscription_id)
    storage_account_keys = storage_management_client.storage_accounts.list_keys(
        resource_group_name=kwargs.get("resource_group", DefaultSettings.resource_group),
        account_name=kwargs.get("storage_account", DefaultSettings.storage_account)
    )
    return storage_account_keys.keys[0].value

def storage_account_get_endpoint_suffix(credentials, subscription_id, **kwargs):
    """
        get Storage account keys
        :param credentials: msrestazure.azure_active_directory.AdalAuthentication
        :param subscription_id: str
        :param **resource_group: str
        :param **storage_account: str
        :param **region: str
    """
    storage_management_client = StorageManagementClient(credentials, subscription_id)
    storage_account = storage_management_client.storage_accounts.get_properties(
        resource_group_name=kwargs.get("resource_group", DefaultSettings.resource_group),
        account_name=kwargs.get("storage_account", DefaultSettings.storage_account)
    )
    # convert https://accountname.blob.core.windows.net/ to core.windows.net
    endpoint_suffix = storage_account.primary_endpoints.blob.split(".blob.")[1].split("/")[0]
    return endpoint_suffix

def create_batch_account(credentials, subscription_id, **kwargs):
    """
        Create a Batch account
        :param credentials: msrestazure.azure_active_directory.AdalAuthentication
        :param subscription_id: str
        :param **resource_group: str
        :param **batch_account: str
        :param **region: str
        :param **storage_account_id: str
    """
    batch_management_client = BatchManagementClient(credentials, subscription_id)
    batch_account = batch_management_client.batch_account.create(
        resource_group_name=kwargs.get("resource_group", DefaultSettings.resource_group),
        account_name=kwargs.get("batch_account", DefaultSettings.batch_account),
        parameters=BatchAccountCreateParameters(
            location=kwargs.get('region', DefaultSettings.region),
            auto_storage=AutoStorageBaseProperties(
                storage_account_id=kwargs.get('storage_account_id', DefaultSettings.region)
            )
        )
    )
    return batch_account.result().id


def batch_account_get_keys(credentials, subscription_id, **kwargs):
    """
        get Batch account keys
        :param credentials: msrestazure.azure_active_directory.AdalAuthentication
        :param subscription_id: str
        :param **resource_group: str
        :param **batch_account: str
    """
    batch_management_client = BatchManagementClient(credentials, subscription_id)
    batch_account_keys = batch_management_client.batch_account.get_keys(
        resource_group_name=kwargs.get("resource_group", DefaultSettings.resource_group),
        account_name=kwargs.get("batch_account", DefaultSettings.batch_account)
    )
    return batch_account_keys.primary

def batch_account_get_url(credentials, subscription_id, **kwargs):
    """
        get Batch account url
        :param credentials: msrestazure.azure_active_directory.AdalAuthentication
        :param subscription_id: str
        :param **resource_group: str
        :param **batch_account: str
    """
    batch_management_client = BatchManagementClient(credentials, subscription_id)
    batch_account = batch_management_client.batch_account.get(
        resource_group_name=kwargs.get("resource_group", DefaultSettings.resource_group),
        account_name=kwargs.get("batch_account", DefaultSettings.batch_account)
    )
    return "https://" + batch_account.account_endpoint

def create_vnet(credentials, subscription_id, **kwargs):
    """
        Create a Batch account
        :param credentials: msrestazure.azure_active_directory.AdalAuthentication
        :param subscription_id: str
        :param **resource_group: str
        :param **virtual_network_name: str
        :param **subnet_name: str
        :param **region: str
    """
    network_client = NetworkManagementClient(credentials, subscription_id)
    resource_group_name = kwargs.get("resource_group", DefaultSettings.resource_group)
    virtual_network_name = kwargs.get("virtual_network_name", DefaultSettings.virtual_network_name)
    subnet_name = kwargs.get("subnet_name", DefaultSettings.subnet_name)
    # get vnet, and subnet if they exist
    virtual_network = subnet = None
    try:
        virtual_network = network_client.virtual_networks.get(
            resource_group_name=resource_group_name,
            virtual_network_name=virtual_network_name,
        )
    except CloudError as e:
        pass

    if virtual_network:
        confirmation_prompt = "A virtual network with the same name ({}) was found. \n"\
                             "Please note that the existing address space and subnets may be changed or destroyed. \n"\
                             "Do you want to use this virtual network? (y/n): ".format(virtual_network_name)
        deny_error = AccountSetupError("Virtual network already exists, not recreating.")
        unrecognized_input_error = AccountSetupError("Input not recognized.")
        prompt_for_confirmation(confirmation_prompt, deny_error, unrecognized_input_error)

    virtual_network = network_client.virtual_networks.create_or_update(
        resource_group_name=resource_group_name,
        virtual_network_name=kwargs.get("virtual_network_name", DefaultSettings.virtual_network_name),
        parameters=VirtualNetwork(
            location=kwargs.get("region", DefaultSettings.region),
            address_space=AddressSpace(["10.0.0.0/24"])
        )
    )
    virtual_network = virtual_network.result()
    subnet = network_client.subnets.create_or_update(
        resource_group_name=resource_group_name,
        virtual_network_name=virtual_network_name,
        subnet_name=subnet_name,
        subnet_parameters=Subnet(
            address_prefix='10.0.0.0/24'
        )
    )
    return subnet.result().id


def create_aad_user(credentials, tenant_id, **kwargs):
    """
        Create an AAD application and service principal
        :param credentials: msrestazure.azure_active_directory.AdalAuthentication
        :param tenant_id: str
        :param **application_name: str
    """
    graph_rbac_client = GraphRbacManagementClient(
        credentials,
        tenant_id,
        base_url=AZURE_PUBLIC_CLOUD.endpoints.active_directory_graph_resource_id
    )
    application_credential = uuid.uuid4()
    try:
        display_name = kwargs.get("application_name", DefaultSettings.application_name)
        application = graph_rbac_client.applications.create(
            parameters=ApplicationCreateParameters(
                available_to_other_tenants=False,
                identifier_uris=["http://{}.com".format(display_name)],
                display_name=display_name,
                password_credentials=[
                    PasswordCredential(
                        end_date=datetime(2299, 12, 31, 0, 0, 0, 0, tzinfo=timezone.utc),
                        value=application_credential,
                        key_id=uuid.uuid4()
                    )
                ]
            )
        )
        service_principal = graph_rbac_client.service_principals.create(
            ServicePrincipalCreateParameters(
                app_id=application.app_id,
                account_enabled=True
            )
        )
    except GraphErrorException as e:
        if e.inner_exception.code == "Request_BadRequest":
            application = next(graph_rbac_client.applications.list(
                filter="identifierUris/any(c:c eq 'http://{}.com')".format(display_name)))

            confirmation_prompt = "Previously created application with name {} found. "\
                                  "Would you like to use it? (y/n): ".format(application.display_name)
            prompt_for_confirmation(confirmation_prompt, e, ValueError("Response not recognized. Please try again."))

            service_principal = next(graph_rbac_client.service_principals.list(
                filter="appId eq '{}'".format(application.app_id)))
        else:
            raise e

    return application.app_id, service_principal.object_id, str(application_credential)


def create_role_assignment(credentials, subscription_id, scope, principal_id):
    """
        Gives service principal contributor role authorization on scope
        :param credentials: msrestazure.azure_active_directory.AdalAuthentication
        :param subscription_id: str
        :param scope: str
        :param principal_id: str
    """
    authorization_client = AuthorizationManagementClient(credentials, subscription_id)
    role_name = 'Contributor'
    roles = list(authorization_client.role_definitions.list(
        scope,
        filter="roleName eq '{}'".format(role_name)
    ))
    contributor_role = roles[0]
    for i in range(10):
        try:
            authorization_client.role_assignments.create(
                scope,
                uuid.uuid4(),
                {
                    'role_definition_id': contributor_role.id,
                    'principal_id': principal_id
                }
            )
            break
        except CloudError as e:
            # ignore error if service principal has not yet been created
            time.sleep(1)
            if i == 10:
                raise e


def format_secrets(**kwargs):
    '''
    Returns the secrets for the created resources to be placed in credentials.json
    The following form is returned:

    For Azure Acitve Directory Authentication:
        "service_principal": {
            "tenantId": "<AAD Diretory ID>"
            "clientId": "<AAD App Application ID>"
            "credential": "<AAD App Password>"
            "batchAccountResourceId": "</batch/account/resource/id>"
            "storageAccountResourceId": "</storage/account/resource/id>"
            "storageEndpointSuffix": "<storage account endpoint suffix>"
        }

    For Shared Key Authentication:
        "sharedKey": {
          "batchAccount": {
              "name": "<batch account name>",
              "key": "<batch account key>",
              "url": "https://batchaccount.region.batch.azure.com"
          },
          "storageAccount": {
              "name": "<storage account name>",
              "key": "<storage account key>",
              "endpointSuffix": "core.windows.net"
          }
        }
    '''
    return json.dumps(kwargs, indent = 4)


def prompt_for_confirmation(prompt, deny_error, unrecognized_input_error):
    """
        Prompt user for confirmation, 'y' for confirm, 'n' for deny
        :param prompt: str
        :param deny_error: Exception
        :param unrecognized_input_error: Exception
        :return None if prompt successful, else raises error
    """
    confirmation = input(prompt).lower()
    for i in range(3):
        if confirmation == "n":
            raise deny_error
        elif confirmation == "y":
            break
        elif confirmation != "y" and i == 2:
            raise unrecognized_input_error
        confirmation = input("Please input 'y' or 'n': ").lower()


def prompt_with_default(key, value):
    user_value = input("{0} [{1}]: ".format(key, value))
    if user_value != "":
        return user_value
    else:
        return value


def prompt_tenant_selection(tenant_ids):
    print("Multiple tenants detected. Please input the ID of the tenant you wish to use.")
    print("Tenants:", ", ".join(tenant_ids))
    given_tenant_id = input("Please input the ID of the tenant you wish to use: ")
    for i in range(3):
        if given_tenant_id in tenant_ids:
            return given_tenant_id
        if i != 2:
            given_tenant_id = input("Input not recognized, please try again: ")
    raise AccountSetupError("Tenant selection not recognized after 3 attempts.")



class Spinner:
    busy = False
    delay = 0.1

    @staticmethod
    def spinning_cursor():
        while 1:
            for cursor in '|/-\\': yield cursor

    def __init__(self, delay=None):
        self.spinner_generator = self.spinning_cursor()
        if delay and float(delay): self.delay = delay

    def __enter__(self):
        return self.start()

    def __exit__(self, exc_type, exc_val, exc_tb):
        return self.stop()

    def spinner_task(self):
        while self.busy:
            sys.stdout.write(next(self.spinner_generator))
            sys.stdout.flush()
            time.sleep(self.delay)
            sys.stdout.write('\b')
            sys.stdout.flush()

    def start(self):
        self.busy = True
        threading.Thread(target=self.spinner_task, daemon=True).start()

    def stop(self):
        self.busy = False
        time.sleep(self.delay)



if __name__ == "__main__":
    print("\nGetting credentials.")
    # get credentials and tenant_id
    creds, subscription_id = credentials.get_azure_cli_credentials()
    subscription_client = SubscriptionClient(creds)
    tenant_ids = [tenant.id for tenant in subscription_client.tenants.list()]
    if len(tenant_ids) != 1:
        tenant_id = prompt_tenant_selection(tenant_ids)
    else:
        tenant_id = tenant_ids[0]

    if len(sys.argv) > 1 and sys.argv[1] == "deleteresourcegroup":
        resource_group = input("Resource Group Name: ")
        with Spinner():
          delete_resource_group(creds, subscription_id, resource_group)
        print("Deleted resource group.")
        sys.exit()
   
    if len(sys.argv) > 1 and sys.argv[1] == "serviceprincipal":
        authentication = "serviceprincipal"
    else:
        authentication = DefaultSettings.authentication

    print("Input the desired names and values for your Azure resources. "\
          "Default values are provided in the brackets. "\
          "Hit enter to use default.")

    chars = string.ascii_lowercase
    suffix = "".join(random.choice(chars) for i in range(4))
    DefaultSettings.storage_account += suffix
    DefaultSettings.batch_account += suffix
    if authentication == DefaultSettings.authentication:
        kwargs = {
            "region": prompt_with_default("Azure Region", DefaultSettings.region),
            "resource_group": prompt_with_default("Resource Group Name", DefaultSettings.resource_group),
            "storage_account": prompt_with_default("Storage Account Name", DefaultSettings.storage_account),
            "batch_account": prompt_with_default("Batch Account Name", DefaultSettings.batch_account)
        }
    else:
        kwargs = {
            "region": prompt_with_default("Azure Region", DefaultSettings.region),
            "resource_group": prompt_with_default("Resource Group Name", DefaultSettings.resource_group),
            "storage_account": prompt_with_default("Storage Account Name", DefaultSettings.storage_account),
            "batch_account": prompt_with_default("Batch Account Name", DefaultSettings.batch_account),
            # "virtual_network_name": prompt_with_default("Virtual Network Name", DefaultSettings.virtual_network_name),
            # "subnet_name": prompt_with_default("Subnet Name", DefaultSettings.subnet_name),
            "application_name": prompt_with_default("Active Directory Application Name", DefaultSettings.application_name),
            "application_credential_name": prompt_with_default("Active Directory Application Credential Name", DefaultSettings.resource_group),
            "service_principal": prompt_with_default("Service Principal Name", DefaultSettings.service_principal)
        }

    print("Creating the Azure resources.")

    # create resource group
    with Spinner():
        resource_group_id = create_resource_group(creds, subscription_id, **kwargs)
        kwargs["resource_group_id"] = resource_group_id
    print("Created Resource Group.")

    # create storage account
    with Spinner():
        storage_account_id = create_storage_account(creds, subscription_id, **kwargs)
        kwargs["storage_account_id"] = storage_account_id
    print("Created Storage Account.")

    with Spinner():
        storage_account_endpoint_suffix = storage_account_get_endpoint_suffix(creds, subscription_id, **kwargs)
        kwargs["storage_account_endpoint_suffix"] = storage_account_endpoint_suffix
    print("Retrieved Storage Account endpoint suffix.")

    # create batch account
    with Spinner():
        batch_account_id = create_batch_account(creds, subscription_id, **kwargs)
    print("Created Batch Account.")

    if authentication == DefaultSettings.authentication:
        # retrieve batch account key
        with Spinner():
            batch_account_key = batch_account_get_keys(creds, subscription_id, **kwargs)
            kwargs["batch_account_key"] = batch_account_key
        print("Retrieved Batch Account key.")

        # retrieve batch account url
        with Spinner():
            batch_account_url = batch_account_get_url(creds, subscription_id, **kwargs)
            kwargs["batch_account_url"] = batch_account_url
        print("Retrieved Batch Account url.")

        with Spinner():
            storage_account_keys = storage_account_get_keys(creds, subscription_id, **kwargs)
            kwargs["storage_account_key"] = storage_account_keys
        print("Retrieved Storage Account key.")

        secrets = format_secrets(
            **{
                "batchAccount": {
                  "name": kwargs["batch_account"],
                  "key": kwargs["batch_account_key"],
                  "url": kwargs["batch_account_url"]
                },
                "storageAccount": {
                  "name": kwargs["storage_account"],
                  "key": kwargs["storage_account_key"],
                  "endpointSuffix": kwargs["storage_account_endpoint_suffix"]
                }
            }
        )
    else:
        # create vnet with a subnet
        # subnet_id = create_vnet(creds, subscription_id)

        # create AAD application and service principal
        with Spinner():
            profile = credentials.get_cli_profile()
            aad_cred, subscirption_id, tenant_id = profile.get_login_credentials(
                resource=AZURE_PUBLIC_CLOUD.endpoints.active_directory_graph_resource_id
            )

            application_id, service_principal_object_id, application_credential = create_aad_user(aad_cred, tenant_id, **kwargs)
        print("Created Azure Active Directory service principal.")

        with Spinner():
            create_role_assignment(creds, subscription_id, resource_group_id, service_principal_object_id)
        print("Configured permsisions.")

        secrets = format_secrets(
            **{
                "servicePrincipal": {
                  "tenantId": tenant_id,
                  "clientId": application_id,
                  "credential": application_credential,
                  "batchAccountResourceId": batch_account_id,
                  "storageAccountResourceId": storage_account_id,
                  "storageEndpointSuffix": kwargs["storage_account_endpoint_suffix"]
                }
            }
        )

    print("\n# Copy the following into your credentials.json file\n{}".format(secrets))