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

from google.oauth2 import service_account
from google.auth.transport.requests import Request
from azure.identity import ClientAssertionCredential
from azure.core.credentials import AccessToken

def getIdTokenByOidcOfGCPSelf():
    jwt_path = os.environ['GOOGLE_APPLICATION_CREDENTIALS']
    aud = os.environ['GOOGLE_AUDIENCE']
    credentials = service_account.IDTokenCredentials.from_service_account_file(
        jwt_path,
        target_audience=aud,
        )
    credentials.refresh(Request())
    return credentials.token

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

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

def writeTokenByOidcFromGCPToAzure():
    token = getIdentityTokenOfGCPToAzure()
    file = os.environ['AZURE_FEDERATED_TOKEN_FILE']
    with open(file, "w+") as f:
        logging.info('exchange to azure token to file: {}'.format(token))
        f.write(token)
    return token

def run():
    vendor = os.environ['VENDOR']
    if vendor == 'gcp':
        token_text = getIdentityTokenOfGCPSelf()
        writeTokenByOidcFromGCPToAzure()
    elif vendor == 'ksyun':
        token_text = getIdTokenByOidcOfGCPSelf()
    else:
        raise ValueError('Vendor env error: {}'.format(vendor))

    logging.info('Identity Text: {}'.format(token_text))
    token_file = os.environ['AWS_WEB_IDENTITY_TOKEN_FILE']
    with open(token_file, "w+") as f:
        f.write(token_text)
    unverified_header = jwt.decode(token_text, options={"verify_signature": False}) # only with header like {'aud': 'gcp', 'azp': '114667344163696279999', 'exp': 1666097306, 'iat': 1666093706, 'iss': 'https://accounts.google.com', 'sub': '114667344163696279999'}, no signature & body value
    logging.info('unverified_header: {}'.format(unverified_header))
    return unverified_header['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)