#!/usr/bin/env python3
# Dump MAC -> friendly-name pairings from the unofficial Aruba Instant On cloud API.
# The API is undocumented and unsupported by HPE; endpoints reverse-engineered by the
# community (https://mspp.io/documenting-aruba-instant-on-sites-aruba-instant-on-api/).
# Requires a portal account WITHOUT MFA (the scripted login can't answer a challenge).

import argparse
import base64
import hashlib
import json
import os
import secrets
import sys
from urllib.parse import parse_qs, urlparse

import requests

# portal.arubainstanton.com now 308-redirects here (HPE rebrand); settings.json
# self-describes the SSO host, OAuth client id, redirect URI, and API base
PORTAL_BASE = 'https://portal.instant-on.hpe.com'
ION_API_VERSION = '7'

# Load .env file
try:
  with open('.env', 'r') as f:
    for line in f:
      if line.strip() and not line.startswith('#') and '=' in line:
        key, value = line.strip().split('=', 1)
        os.environ[key] = value
except FileNotFoundError:
  pass


def get_bearer_token(settings, username, password):
  sso_base = settings['ssoFqdn']
  redirect_uri = settings['ssoRedirectUrl']

  # PKCE: challenge is base64url(sha256(verifier)) without padding
  code_verifier = secrets.token_urlsafe(48)
  code_challenge = base64.urlsafe_b64encode(hashlib.sha256(code_verifier.encode()).digest()).rstrip(b'=').decode()

  r = requests.post(f'{sso_base}/aio/api/v1/mfa/validate/full', data={'username': username, 'password': password})
  if r.status_code != 200:
    sys.exit(f'login failed ({r.status_code}): check credentials, and that the account has MFA disabled')
  session_token = r.json()['access_token']

  # the authorization code comes back in the redirect Location, so don't follow it
  r = requests.get(f'{sso_base}/as/authorization.oauth2',
                   params={
                     'client_id': settings['ssoClientIdAuthZ'],
                     'redirect_uri': redirect_uri,
                     'response_type': 'code',
                     'scope': 'profile openid',
                     'state': secrets.token_urlsafe(32),
                     'code_challenge_method': 'S256',
                     'code_challenge': code_challenge,
                     'sessionToken': session_token,
                   },
                   allow_redirects=False)
  location = r.headers.get('Location', '')
  query = parse_qs(urlparse(location).query)
  if 'code' not in query:
    sys.exit(f'authorization failed: no code in redirect ({location or r.status_code})')

  r = requests.post(f'{sso_base}/as/token.oauth2',
                    data={
                      'client_id': settings['ssoClientIdAuthZ'],
                      'redirect_uri': redirect_uri,
                      'code': query['code'][0],
                      'code_verifier': code_verifier,
                      'grant_type': 'authorization_code',
                    })
  r.raise_for_status()
  return r.json()['access_token']


def api_get(settings, token, path):
  r = requests.get(f'{settings["restApiUrl"]}/api/{path}',
                   headers={
                     'Authorization': f'Bearer {token}',
                     'x-ion-api-version': ION_API_VERSION
                   })
  r.raise_for_status()
  data = r.json()
  # collection responses wrap the list in "elements"
  return data.get('elements', data) if isinstance(data, dict) else data


def first_key(record, *candidates):
  # field names vary slightly across API versions, so probe a few
  for key in candidates:
    if record.get(key):
      return record[key]
  return ''


def collect_clients(settings, token, site_id, site_name):
  # clientSummary covers wired and wireless alike (the old wiredClientSummary endpoint is gone)
  rows = []
  for client in api_get(settings, token, f'sites/{site_id}/clientSummary'):
    rows.append({
      'site': site_name,
      'name': first_key(client, 'name', 'hostName'),
      'mac': first_key(client, 'macAddress', 'mac').lower(),
      'ip': first_key(client, 'ipAddress', 'ip'),
      'kind': 'wireless' if client.get('wirelessNetworkId') else 'wired',
    })
  return rows


def main():
  parser = argparse.ArgumentParser(description='List MAC -> friendly-name pairings from Aruba Instant On')
  parser.add_argument('--json', action='store_true', help='output JSON instead of a table')
  args = parser.parse_args()

  try:
    username = os.environ['INSTANT_ON_USER']
    password = os.environ['INSTANT_ON_PASSWORD']
  except KeyError as e:
    sys.exit(f'missing {e.args[0]} — add it to .env or the environment')

  settings = requests.get(f'{PORTAL_BASE}/settings.json').json()
  token = get_bearer_token(settings, username, password)

  rows = []
  for site in api_get(settings, token, 'sites/'):
    rows.extend(collect_clients(settings, token, site['id'], first_key(site, 'name', 'siteName')))
  rows.sort(key=lambda r: (r['site'], r['name'].lower()))

  if args.json:
    print(json.dumps(rows, indent=2))
    return

  widths = {col: max([len(col)] + [len(r[col]) for r in rows]) for col in ('mac', 'name', 'ip', 'kind', 'site')}
  for r in [{col: col.upper() for col in widths}] + rows:
    print('  '.join(r[col].ljust(widths[col]) for col in ('mac', 'name', 'ip', 'kind', 'site')))


if __name__ == '__main__':
  main()

MAC                NAME         IP            KIND      SITE
aa:bb:cc:dd:ee:01  apple-tv     192.168.1.20  wireless  Home
aa:bb:cc:dd:ee:02  thermostat   192.168.1.31  wireless  Home
aa:bb:cc:dd:ee:03  nas          192.168.1.10  wired     Home