#!/usr/bin/python3
import logging
import requests
import jwt
import os
import time
import sys

def getIdentityToken():
    url = 'http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity?format=standard&audience=gcp'
    headers = {'Metadata-Flavor': 'Google'}
    try:
        meta_request = requests.get(url, headers=headers)
    except requests.exceptions.RequestException as e:
        raise e
    if meta_request.ok:
        return meta_request.text
    else:
        raise ValueError('Compute Engine meta data error')

def run():
    token_file = os.environ['AWS_WEB_IDENTITY_TOKEN_FILE']
    token = getIdentityToken()
    with open(token_file, "w+") as f:
        f.write(token)
    payload = jwt.decode(token, verify=False)
    return payload['exp']

def sleepUntilNextRefresh(exp: int):
    d = exp - int(time.time())
    if d > 0:
        logging.info('Issued token will be expired after {} seconds, refresh scheduled after {} seconds'.format(d, d/2))
        time.sleep(d / 2)
    else:
        logging.warning('Issued token is already expired, expiration: {}'.format(exp))

if __name__ == '__main__':
    logging.basicConfig()
    logging.root.setLevel(logging.INFO)
    while True:
        try:
            exp = run()
        except Exception as e:
            # Exit so that we can alert on this, kubelet will restart me whe I run as a sidecar
            sys.exit('Error refreshing token, exiting...'.format(e))
        sleepUntilNextRefresh(exp)