import json
import sys
import time
import urllib.request

import jwt

from common import ISSUER, SUB_PREFIX


def fail(reason):
    print("rejected  ", reason)
    sys.exit(1)


def get(url):
    with urllib.request.urlopen(url, timeout=5) as r:
        return json.load(r)


token = open(sys.argv[1]).read().strip()
audience = sys.argv[2] if len(sys.argv) > 2 else "sts.amazonaws.com"

discovery = get(ISSUER + "/.well-known/openid-configuration")
if discovery["issuer"] != ISSUER:
    fail(f"discovery names issuer {discovery['issuer']}, expected {ISSUER}")
print("discovery  issuer", discovery["issuer"])
print("discovery  jwks_uri", discovery["jwks_uri"])

header = jwt.get_unverified_header(token)
if header.get("alg") != "ES256":
    fail(f"alg {header.get('alg')} is not ES256")
keys = get(discovery["jwks_uri"])["keys"]
match = [k for k in keys if k.get("kid") == header.get("kid")]
if not match:
    fail(f"kid {header.get('kid')} is not in the key set of {len(keys)}")
print("jwks       kid", header["kid"][:16] + "...", "matched 1 of", len(keys), "keys")
key = jwt.PyJWK.from_dict(match[0]).key

try:
    claims = jwt.decode(
        token,
        key,
        algorithms=["ES256"],
        audience=audience,
        issuer=ISSUER,
        leeway=0,
        options={"require": ["exp", "aud", "sub", "iss"]},
    )
except jwt.InvalidAudienceError:
    fail(f"aud {jwt.decode(token, options={'verify_signature': False})['aud']} is not {audience}")
except jwt.PyJWTError as e:
    fail(f"{type(e).__name__}: {e}")
print("signature  ok with the matched key")
print("iss        ok", claims["iss"])
print("aud        ok", claims["aud"])
print("exp        ok, expires in", claims["exp"] - int(time.time()), "s")
if not claims["sub"].startswith(SUB_PREFIX):
    fail(f"sub {claims['sub']} is outside {SUB_PREFIX}")
print("sub        ok", claims["sub"])
print("accepted   execution", claims["execution_id"], "of sandbox", claims["sandbox_id"])