136 lines
4.6 KiB
Python
136 lines
4.6 KiB
Python
# vim: ts=4 et:
|
|
|
|
import json
|
|
import logging
|
|
import os
|
|
import sys
|
|
import time
|
|
import urllib.error
|
|
|
|
from datetime import datetime
|
|
from email.utils import parsedate
|
|
from urllib.request import Request, urlopen
|
|
|
|
|
|
class IdentityBrokerClient:
|
|
"""Client for identity broker
|
|
|
|
Export IDENTITY_BROKER_ENDPOINT to override the default broker endpoint.
|
|
Export IDENTITY_BROKER_API_KEY to specify an API key for the broker.
|
|
|
|
See README_BROKER.md for more information and a spec.
|
|
"""
|
|
|
|
_DEFAULT_ENDPOINT = 'https://aws-access.crute.us/api/account'
|
|
_DEFAULT_ACCOUNT = 'alpine-amis-user'
|
|
_LOGFORMAT = '%(name)s - %(levelname)s - %(message)s'
|
|
|
|
def __init__(self, endpoint=None, key=None, account=None, debug=False):
|
|
# log to STDOUT so that it's not all red when executed by Packer
|
|
self._logger = logging.getLogger('identity-broker')
|
|
self._logger.setLevel(logging.DEBUG if debug else logging.INFO)
|
|
console = logging.StreamHandler(sys.stdout)
|
|
console.setFormatter(logging.Formatter(self._LOGFORMAT))
|
|
self._logger.addHandler(console)
|
|
|
|
self._endpoint = os.environ.get('IDENTITY_BROKER_ENDPOINT') or endpoint \
|
|
or self._DEFAULT_ENDPOINT
|
|
self._key = os.environ.get('IDENTITY_BROKER_API_KEY') or key
|
|
self._account = account or self._DEFAULT_ACCOUNT
|
|
if not self._key:
|
|
raise Exception('No identity broker key found')
|
|
|
|
self._headers = {
|
|
'Accept': 'application/vnd.broker.v2+json',
|
|
'X-API-Key': self._key
|
|
}
|
|
self._cache = {}
|
|
self._expires = {}
|
|
self._default_region = {}
|
|
|
|
def _is_cache_valid(self, path):
|
|
if path not in self._cache:
|
|
return False
|
|
|
|
# path is subject to expiry AND its time has passed
|
|
if self._expires[path] and self._expires[path] < datetime.utcnow():
|
|
return False
|
|
|
|
return True
|
|
|
|
def _get(self, path):
|
|
self._logger.debug("request: %s", path)
|
|
if not self._is_cache_valid(path):
|
|
while True: # to handle rate limits
|
|
try:
|
|
res = urlopen(Request(path, headers=self._headers))
|
|
except urllib.error.HTTPError as ex:
|
|
if ex.status == 401:
|
|
raise Exception('Expired or invalid identity broker token')
|
|
|
|
if ex.status == 406:
|
|
raise Exception('Invalid or malformed identity broker token')
|
|
|
|
# TODO: will this be entirely handled by the 401 above?
|
|
if ex.headers.get('Location') == '/logout':
|
|
raise Exception('Identity broker token is expired')
|
|
|
|
if ex.status == 429:
|
|
self._logger.warning(
|
|
'Rate-limited by identity broker, sleeping 30 seconds')
|
|
time.sleep(30)
|
|
continue
|
|
|
|
raise ex
|
|
|
|
if res.status not in {200, 429}:
|
|
raise Exception(res.reason)
|
|
|
|
# never expires without valid RFC 1123 Expires header
|
|
if expires := res.getheader('Expires'):
|
|
expires = parsedate(expires)
|
|
# convert RFC 1123 to datetime, if parsed successfully
|
|
expires = datetime(*expires[:6])
|
|
|
|
self._expires[path] = expires
|
|
self._cache[path] = json.load(res)
|
|
break
|
|
|
|
self._logger.debug("response: %s", self._cache[path])
|
|
return self._cache[path]
|
|
|
|
def get_credentials_url(self, vendor):
|
|
accounts = self._get(self._endpoint)
|
|
if vendor not in accounts:
|
|
raise Exception(f'No {vendor} credentials found')
|
|
|
|
for account in accounts[vendor]:
|
|
if account['short_name'] == self._account:
|
|
return account['credentials_url']
|
|
|
|
raise Exception('No account credentials found')
|
|
|
|
def get_regions(self, vendor):
|
|
out = {}
|
|
|
|
for region in self._get(self.get_credentials_url(vendor)):
|
|
if region['enabled']:
|
|
out[region['name']] = region['credentials_url']
|
|
|
|
if region['default']:
|
|
self._default_region[vendor] = region['name']
|
|
|
|
return out
|
|
|
|
def get_default_region(self, vendor):
|
|
if vendor not in self._default_region:
|
|
self.get_regions(vendor)
|
|
|
|
return self._default_region.get(vendor)
|
|
|
|
def get_credentials(self, vendor, region=None):
|
|
if not region:
|
|
region = self.get_default_region(vendor)
|
|
|
|
return self._get(self.get_regions(vendor)[region])
|