aztag

#!/usr/bin/env python
# aztag

import atexit
import json
import logging
import msal
import os
import requests
import sys
import time
import uuid
import yaml
from datetime import datetime

# Global variables
prgver          = "v1.0"
prgname         = "aztag"
confdir         = ""
tenant_id       = ""
client_id       = ""
client_secret   = ""
authority_url   = ""
mg_url          = "https://graph.microsoft.com"
mg_token        = {} 
mg_headers      = {}


def print_usage():
    print(prgname + " Azure SP tagging utility " + prgver + "\n"
        "      UUID                        Display Service Principal (SP) tags attribute\n"
        "      UUID \"tag1,tag2\"            Set SP tags to quoted, comma-delimited list\n"
        "      [-i] UUID                   Use interactive Azure logon to display SP tags\n"
        "      [-i] UUID \"tag1,tag2\"       Use interactive Azure logon to update SP tags\n"
        "      -xt                         Delete cached accessTokens file\n"
        "      -v                          Print this usage page")
    sys.exit(1)


def panic(s):
    print("Exception caught:\n%s" % (s))
    sys.exit(1)


def dump_variables():
    print("tenant_id:       ", tenant_id, "\n"
        "client_id:       ", client_id, "\n"
        "client_secret:   ", client_secret, "\n"
        "authority_url:   ", authority_url, "\n"
        "mg_url:          ", mg_url, "\n"
        "mg_headers:      ")
    print_json(mg_headers)
    sys.exit(1)


def remove_cache_file():
    # Remove cache file for objects of type t
    filePath = os.path.join(confdir, "automated_accessTokens.json")
    remove_file(filePath)
    filePath = os.path.join(confdir, "interactive_accessTokens.json")
    remove_file(filePath)
    sys.exit(0)


def remove_file(f):
    if file_exist(f):
        os.remove(f)


def file_exist(f):
    return os.path.exists(f)


def file_size(f):
    return os.path.getsize(f)


def print_json(r):
    print(json.dumps(r, indent=2))


def valid_uuid(id):
    try:
        uuid.UUID(str(id))
        return True
    except Exception as e:
        return False


def setup_credentials():
    # Tell this function we are updating below global variables
    global tenant_id, client_id, client_secret
    
    # Set up credentials and other variables
    creds_file = os.path.join(confdir, "credentials.json")
    if not file_exist(creds_file) or file_size(creds_file) < 1:
        print("Missing credentials file: \"%s\"" % (creds_file))
        with open(creds_file, "w") as f:
            f.write("{\n  \"tenant_id\" : \"UUID\",\n  \"client_id\" : \"UUID\",\n  \"client_secret\" : \"SECRET\"\n}\n")
        os.chmod(creds_file, 0o600)
        print("Created new skeleton file: Please edit it, fill in required values, and re-run program.")
        sys.exit(1)

    try:
        with open(creds_file) as f:
            creds = json.load(f)
            tenant_id = creds["tenant_id"]
            client_id = creds["client_id"]
            client_secret = creds["client_secret"]
    except Exception as e:
        panic(str(e))

    if not valid_uuid(tenant_id):
        print("tenant_id '%s' in '%s' is not a valid UUID" % (tenant_id, creds_file))
        sys.exit(1)
    if not valid_uuid(client_id):
        print("client_id '%s' in '%s' is not a valid UUID" % (client_id, creds_file))
        sys.exit(1)
    if not client_secret:
        print("client_secret in '%s' is blank" % (creds_file))
        sys.exit(1)


def setup_mg_token(method="automated"):
    # Tell this function we are updating below global variables
    global authority_url, mg_token, mg_headers

    authority_url = 'https://login.microsoftonline.com/' + tenant_id

    # Note that '/.defaut' uses whatever static permissions are defined for the SP in the resource domain
    mg_token = get_token([mg_url + '/.default'], method)
    mg_headers = {'Authorization': 'Bearer ' + mg_token['access_token'], 'Content-Type': 'application/json'}


def get_token(scopes, method="automated"):
    token = None

    # Set up cache, as per https://msal-python.readthedocs.io/en/latest/#msal.SerializableTokenCache
    cache = msal.SerializableTokenCache()
    cache_file = os.path.join(confdir, method + "_accessTokens.json") # Use unique cache file for each method
    if file_exist(cache_file):
        cache.deserialize(open(cache_file, 'r').read())
    
    atexit.register(lambda:
        open(cache_file, 'w').write(cache.serialize())
        # Hint: The following option line persists only when state changed
        if cache.has_state_changed else None
        and os.chmod(cache_file, 0o600)
    )

    # Acquire token
    app = None
    if method == "interactive":
        username = input('Azure AD Username: ')
        tenant_id = username.split('@')[1]
        try:
            app = msal.PublicClientApplication(
                client_id="1950a258-227b-4e31-a9cf-717495945fc2",
                # Important: Note we are using standard Microsoft Azure PowerShell client_id
                client_credential=None,
                authority = authority_url + tenant_id,
                token_cache = cache
            )
        except Exception as e:
            panic(str(e))

        try:
            # Use 'Device code flow'
            # First, try getting possible cached token for this account
            for x in app.get_accounts():
                if x["username"] == username:
                    token = app.acquire_token_silent(scopes, account=x)
            if not token:
                flow = app.initiate_device_flow(scopes=scopes)
                if "user_code" not in flow:
                    raise ValueError("Fail to create device flow. Err: %s" % json.dumps(flow, indent=4))
                print(flow["message"])
                sys.stdout.flush()
                try:
                    # Else, just get a new token
                    token = app.acquire_token_by_device_flow(flow)
                    #token = app.acquire_token_interactive(scopes=scopes)
                    if 'access_token' not in token:
                        panic(str(e))
                except:
                    pass            
        except Exception as e:
            panic(str(e))
    else:
        try:
            app = msal.ConfidentialClientApplication(
                client_id,
                client_credential = client_secret,
                authority = authority_url,
                token_cache = cache
            )
        except Exception as e:
            panic(str(e))           
        try:
            # Try getting cached token first
            token = app.acquire_token_silent(scopes, account=None)
            if not token:
                try:
                    # Else, just get a new token
                    token = app.acquire_token_for_client(scopes=scopes)
                    if 'access_token' not in token:
                        panic(str(e))
                except:
                    pass
        except Exception as e:
            panic(str(e))

    return token


def api_get(resource, headers=None, params=None, verbose=False):
    if headers == None:
        headers = mg_headers
    try:
        r = requests.get(resource, headers=headers, params=params).json()
        if isinstance(r, int):  # Handle $count filter integer returns
            return r
        if 'error' in r:
            if verbose:
                print("API CALL: %s\nPARAMS  : %s\nHEADERS : %s" % (resource, params, headers))
            print(r['error']['message'])
        return r
    except Exception as e:
        panic(str(e))


def api_patch(resource, headers=None, params=None, verbose=False, data=None):
    if headers == None:
        headers = mg_headers
    try:
        if verbose:
            print("API CALL DETAILS\nURL     : %s\nPARAMS  : %s\nHEADERS : %s\nDATA\n%s" % (resource, params, headers, data))
        r = requests.patch(resource, headers=headers, params=params, json=data)
        return r
    except Exception as e:
        panic(str(e))


def update_tags(id, tag_str):
    # Build tag list
    payload = {}
    if tag_str == "":
        payload = { "tags": [] }
    else:
        tags = []
        for t in tag_str.split(','):
            tags.append(t.strip())
        payload = { "tags": tags }

    # Update SP
    r = api_patch(mg_url + "/v1.0/servicePrincipals/" + id, data=payload)
    if r.status_code != 204:
        print_json(r.json())


def print_tags(id):
    # Get SP object
    x = api_get(mg_url + "/v1.0/servicePrincipals/" + id)

    # Print its tags
    if "tags" in x:
        if len(x["tags"]) > 0:
            print("tags:")
            for tag in x["tags"]:
                print("  - %s" % (tag))
        else:        
            print("tags: []")


def main(args = None):
    number_of_args = len(sys.argv[1:]) # Not including the program itself
    if number_of_args not in [1, 2, 3]:
        # Don't accept less than 1 or more than 3 arguments
        print_usage()  
    
    global confdir # We're updating a global variable
    if os.environ['HOME']:
        confdir = os.path.join(os.environ['HOME'], "." + prgname)
        if not file_exist(confdir):
            os.makedirs(confdir)
            os.chmod(confdir, 0o700)
    else:
        print("Missing HOME environment variable")
        sys.exit(1)

    if number_of_args == 1: # One-argument requests
        arg1 = sys.argv[1]

        setup_credentials() # Set up tenant ID and credentials
        setup_mg_token()    # Remaining requests need API tokens

        if valid_uuid(arg1):
            print_tags(arg1)
        elif arg1.lower() == "-z":
            dump_variables()
        elif arg1.lower() == "-xt":
            remove_cache_file() # Remove token file
        else:
            print_usage()

    elif number_of_args == 2: # Two-argument requests
        arg1 = sys.argv[1]
        arg2 = sys.argv[2]

        if arg1.lower() == "-i" and valid_uuid(arg2):
            setup_mg_token("interactive") # Remaining requests need API tokens
            print_tags(arg2)
            sys.exit(1)
        else:
            setup_credentials() # Set up tenant ID and credentials
            setup_mg_token()    # Remaining requests need API tokens

        if valid_uuid(arg1):
            update_tags(arg1, arg2)
        else:
            print_usage()
    else:    # Three-argument requests
        arg1 = sys.argv[1]
        arg2 = sys.argv[2]
        arg3 = sys.argv[3]

        if arg1.lower() == "-i" and valid_uuid(arg2):
            setup_mg_token("interactive") # Remaining requests need API tokens
            update_tags(arg2, arg3)
        else:
            print_usage()


if __name__ == '__main__':
    main()