# https://cloud.google.com/iap/docs/authentication-howto

import json
import time

from jwt import JWT, jwk_from_pem
import requests

jwt = JWT()

OAUTH_TOKEN_URI = "https://www.googleapis.com/oauth2/v4/token"
JWT_BEARER_TOKEN_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:jwt-bearer"

CLIENT_ID = (
    ""
)  # https://console.cloud.google.com/apis/credentials
SERVICE_ACCOUNT_SECRET_KEY = json.loads(
    open("service-account.json", "r").read()
)  # https://console.cloud.google.com/iam-admin/serviceaccounts (Actions -> Create key -> JSON)


def get_jwt_assertion():
    message = {
        "kid": SERVICE_ACCOUNT_SECRET_KEY["private_key_id"],
        "iss": SERVICE_ACCOUNT_SECRET_KEY["client_email"],
        "sub": SERVICE_ACCOUNT_SECRET_KEY["client_email"],
        "aud": OAUTH_TOKEN_URI,
        "iat": int(time.time()),
        "exp": int(time.time()) + 60 * 65,
        "target_audience": CLIENT_ID,
    }

    return jwt.encode(
        message,
        jwk_from_pem(SERVICE_ACCOUNT_SECRET_KEY["private_key"].encode("utf-8")),
        "RS256",
    )


def get_google_open_id_connect_token():
    body = {"assertion": get_jwt_assertion(), "grant_type": JWT_BEARER_TOKEN_GRANT_TYPE}
    r = requests.post(OAUTH_TOKEN_URI, data=body)
    r.raise_for_status()

    return r.json()["id_token"]


print(get_google_open_id_connect_token())