Azure client poc using aiohttp

azure.py

#!/usr/bin/python3.8
import asyncio
import os
from typing import List

import aiohttp
from dataclasses import dataclass, field
import typedload


@dataclass
class Oauth2Token:
    token_type: str
    expires_in: str
    ext_expires_in: str
    expires_on: str
    not_before: str
    resource: str
    access_token: str


async def fetch_token() -> Oauth2Token:
    app_id = os.getenv('AZURE_APP_ID')
    client_secret = os.getenv('AZURE_CLIENT_SECRET')
    tenant_id = os.getenv('AZURE_TENANT_ID')
    body = {
        'grant_type': 'client_credentials',
        'client_id': app_id,
        'client_secret': client_secret,
        'resource': 'https://management.azure.com/',
    }
    headers = {
        'Accept': 'application/json',
        'Content-Type': 'application/x-www-form-urlencoded',
    }
    async with aiohttp.ClientSession() as session:
        async with session.post(f'https://login.microsoftonline.com/{tenant_id}/oauth2/token',
                                headers=headers, data=body) as r:
            token_raw = await r.json()

    return typedload.load(token_raw, Oauth2Token)


async def get_network_interface_private_ips(session: aiohttp.ClientSession, network_interface_ref) -> List[str]:
    ips = []
    async with session.get('https://management.azure.com'
                           f'{network_interface_ref["id"]}?api-version=2019-12-01') as r:
        network_interface = await r.json()
        ips.extend(
            ip_configuration['properties']['privateIPAddress']
            for ip_configuration in network_interface['properties']['ipConfigurations']
        )
    return ips


async def main():
    token = await fetch_token()
    headers = {
        'Authorization': f'Bearer {token.access_token}',
        'Content-Type': 'application/json',
    }

    # List VMs and their private ips
    subscription_id = os.getenv('AZURE_SUBSCRIPTION_ID')
    async with aiohttp.ClientSession(headers=headers) as session:
        print('Listing all vm private ips using list all')
        async with session.get(f'https://management.azure.com/subscriptions/{subscription_id}/'
                               'providers/Microsoft.Compute/virtualMachines?api-version=2019-12-01') as r:
            vms = (await r.json())['value']

        for vm in vms:
            network_interfaces = vm['properties']['networkProfile']['networkInterfaces']
            ips = []
            for network_interface_ref in network_interfaces:
                ips.extend(await get_network_interface_private_ips(session, network_interface_ref))
            print(vm['name'], ', '.join(ips))

        # Now try the same with the generic resource API
        print('Listing all vm private ips using resources filter')
        url = f'https://management.azure.com/subscriptions/{subscription_id}/resources'
        params = {
            'api-version': '2019-10-01',
            # You can query by resourceType and name and resourceGroup, but you can't combine that with tags?
            '$filter': "resourceType eq 'Microsoft.Compute/virtualMachines' and "
                       "substringof('purple20', name) "
                       "and resourceGroup eq 'purple'"
            # '$filter': "tagName eq 'Name' and tagValue eq 'purple20.agi.appgate.com'"
        }
        async with session.get(url, params=params) as r:
            vm_refs = (await r.json())['value']
            for vm_ref in vm_refs:
                async with session.get('https://management.azure.com'
                                       f'{vm_ref["id"]}?api-version=2019-12-01') as r:
                    vm = await r.json()
                network_interfaces = vm['properties']['networkProfile']['networkInterfaces']
                ips = []
                for network_interface_ref in network_interfaces:
                    ips.extend(await get_network_interface_private_ips(session, network_interface_ref))
                print(vm['name'], ', '.join(ips))


if __name__ == '__main__':
    asyncio.run(main())