#!/usr/bin/env python3
# Copyright (C) 2024 Spearhead Systems SRL

from urllib   import request, parse, error
from datetime import datetime, timezone, timedelta
import json
import sys
import re


VAULT_METRICS = [
    'Availability',
    'SaturationShoebox',
    'ServiceApiLatency',
    'ServiceApiHit',
    'ServiceApiResult',
]

FIREWALL_METRICS = [
    'FirewallHealth',
    'Throughput',
    'FirewallLatencyPng',
]

REGION_RE = re.compile('/locations/(.+?)/')
RESOURCE_GROUP_RE = re.compile('/resourceGroups/(.+?)/')


# https://learn.microsoft.com/en-us/azure/azure-resource-manager/management/request-limits-and-throttling
def get_url(req, default):
    #import http.client
    #http.client.HTTPConnection.debuglevel = 1
    try:
        res = request.urlopen(req)
        return res.read()
    except error.HTTPError as e:
        if e.code == 429:
            return default
        else:
            raise e


def set_proxy(req, proxy):
    if proxy is None or proxy == '':
        return

    match = re.match('(https?)://(.+?)/?$', proxy, re.I)
    req.set_proxy(match[2], match[1].lower())

    # The explicit Host header is required for this to also work with a proxy.
    # If we don't include it, Python sends the proxy's Host to Microsoft
    # instead! So we have to set the Host to the Microsoft domain manually.
    match = re.match('https://(.+?)/', req.full_url, re.I)
    req.add_header('Host', match[1] + ":443")


def get_token(tenant, username, password, proxy):
    data = parse.urlencode({
        'client_id': username,
        'client_secret': password,
        'grant_type': 'client_credentials',
        'claims': '{"access_token": {"xms_cc": {"values": ["CP1"]}}}',
        'scope': 'https://management.core.windows.net//.default offline_access openid profile',
        'client_info': 1,
    })

    req = request.Request(f'https://login.microsoftonline.com/{tenant}/oauth2/v2.0/token',
                          data=str.encode(data))
    set_proxy(req, proxy)

    res = get_url(req, None)
    if res is None:
        return

    token_data = json.loads(res)
    token = token_data['access_token']
    return token


def get_json(token, proxy, path, version='2023-07-01'):
    url = f"https://management.azure.com{path}{'?' in path and '&' or '?'}api-version={version}"
    req = request.Request(url, headers={'Authorization': f'Bearer {token}'})
    set_proxy(req, proxy)
    res = get_url(req, "[]")
    data = json.loads(res)
    return data['value']


def list_subscriptions(token, proxy):
    return get_json(token, proxy, '/subscriptions')


def list_vaults(token, proxy, subscription):
    return get_json(token, proxy, f'/subscriptions/{subscription}/resources?$filter=resourceType%20eq%20%27Microsoft.KeyVault%2Fvaults%27')


def list_firewalls(token, proxy, subscription):
    return get_json(token, proxy, f'/subscriptions/{subscription}/resources?$filter=resourceType%20eq%20%27Microsoft.Network%2FazureFirewalls%27')


def list_defender_alerts(token, proxy, subscription):
    return get_json(token, proxy, f'/subscriptions/{subscription}/providers/Microsoft.Security/alerts', '2022-01-01')


def get_recent_metrics(token, proxy, path, metrics):
    end   = datetime.now()
    start = end - timedelta(minutes=2)

    start_str   = start.isoformat().split('.')[0] + 'Z'
    end_str     = end.isoformat().split('.')[0] + 'Z'
    metrics_str = ','.join(metrics)

    return get_json(token, proxy, f'{path}/providers/microsoft.insights/metrics?metricnames={metrics_str}&timespan={start_str}/{end_str}', '2023-10-01')


def metrics_to_lookup(metrics):
    lookup = {}

    for metric in metrics:
        name = metric['name']['value']
        series = metric['timeseries']
        if series:
            value = series[0]['data'][-1]
            key = next(filter(lambda foo: foo != 'timeStamp', value), None)
            lookup[name] = value.get(key)

    return lookup


def get_args(argv):
    if (len(argv) != 5 and len(argv) != 6) or argv[1] not in ['keyvault', 'firewall', 'defender']:
        print(f"{sys.argv[0]} <command> <tenant ID> <username> <password> <proxy>", file=sys.stderr)
        print(f"Valid commands are: 'keyvault', 'firewall', 'defender'", file=sys.stderr)
        print(f"Proxy is an optional argument", file=sys.stderr)
        exit(1)
    return argv[1], argv[2], argv[3], argv[4], (argv[5] if len(argv) == 6 else None)


def print_json(obj):
    print(json.dumps(obj))


def get_resource_group(obj):
    found = re.search(RESOURCE_GROUP_RE, obj['id'])
    if found:
        return found[1]
    return None


resource_groups = {}
command, tenant, username, password, proxy = get_args(sys.argv)
token = get_token(tenant, username, password, proxy)

for subscription in list_subscriptions(token, proxy):
    subscription_id = subscription['subscriptionId']

    if command == 'defender':
        for alert in list_defender_alerts(token, proxy, subscription_id):
            properties = alert['properties']
            status = properties['status']

            if not status in ['Active', 'InProgress']:
                continue

            group = get_resource_group(alert)

            resource_groups.setdefault(group, []).append({
                'type': command,
                'name': alert['name'],
                'location': re.search(REGION_RE, alert['id'])[1],
                'resource_group': group,
                'alert': {
                    'status': status,
                    'severity': properties['severity'],
                    'url':      properties['alertUri'],
                    'info':     properties['alertDisplayName']

                }
            })

    elif command == 'firewall':
        for firewall in list_firewalls(token, proxy, subscription_id):
            metrics = get_recent_metrics(token, proxy, firewall['id'], FIREWALL_METRICS)
            group = get_resource_group(firewall)

            resource_groups.setdefault(group, []).append({
                'type': command,
                'name': firewall['name'],
                'location': firewall['location'],
                'resource_group': get_resource_group(firewall),
                'metrics':  metrics_to_lookup(metrics),
            })

    elif command == 'keyvault':
        for vault in list_vaults(token, proxy, subscription_id):
            metrics = get_recent_metrics(token, proxy, vault['id'], VAULT_METRICS)
            group = get_resource_group(vault)

            resource_groups.setdefault(group, []).append({
                'type':     command,
                'name':     vault['name'],
                'location': vault['location'],
                'resource_group': group,
                'metrics':  metrics_to_lookup(metrics),
            })

for group, results in resource_groups.items():
    if group is None:
        print(f"<<<<>>>>")
    else:
        print(f"<<<<{group}>>>>")

    print(f"<<<azure_{command}:sep(0)>>>")
    for result in results:
        print_json(result)
