#!/usr/bin/env python3
"""
Tests the IMDS emulator: the metadata endpoints, the request checks, the
Compute Engine credentials of google-auth on top of the emulator, and the
validity of the issued access token at Google.
"""
import argparse
import json
import os
import sys
import urllib.error
import urllib.parse
import urllib.request

parser = argparse.ArgumentParser(description="Test the IMDS emulator.")
parser.add_argument("--host", default="localhost", help="host and optional port of the emulator")
parser.add_argument("--bucket", help="bucket to list with the token, optional")
args = parser.parse_args()

root = f"http://{args.host}/computeMetadata/v1/"
failures = 0


def check(name, condition, detail=""):
    """
    Prints the result of a check and counts the failures.
    """
    global failures
    print(f"{'PASS' if condition else 'FAIL'} {name}{': ' + detail if detail else ''}")
    if not condition:
        failures += 1


def get(url, headers=None):
    """
    Returns the status, the headers and the body of a GET request.
    """
    request = urllib.request.Request(url, headers=headers or {})
    try:
        with urllib.request.urlopen(request, timeout=10) as response:
            return response.status, response.headers, response.read().decode()
    except urllib.error.HTTPError as error:
        return error.code, error.headers, error.read().decode()


flavor = {"Metadata-Flavor": "Google"}

# Detection ping and request checks
status, headers, _ = get(f"http://{args.host}/", flavor)
check("ping", status == 200 and headers.get("Metadata-Flavor") == "Google", f"status {status}")
status, _, _ = get(root + "project/project-id")
check("request without Metadata-Flavor is rejected", status == 403, f"status {status}")
status, _, _ = get(root + "project/project-id", {**flavor, "X-Forwarded-For": "192.0.2.1"})
check("request with X-Forwarded-For is rejected", status == 403, f"status {status}")

# Metadata
status, _, project_id = get(root + "project/project-id", flavor)
check("project ID", status == 200 and project_id != "", project_id)
status, _, email = get(root + "instance/service-accounts/default/email", flavor)
check("service account email", status == 200 and "@" in email, email)
status, _, body = get(root + "instance/service-accounts/default/?recursive=true", flavor)
check("service account info", status == 200 and json.loads(body).get("email") == email, body)
status, _, body = get(root + f"instance/service-accounts/{email}/email", flavor)
check("service account by email", status == 200 and body == email, body)

# Token
status, _, body = get(root + "instance/service-accounts/default/token", flavor)
token = json.loads(body) if status == 200 else {}
check("token", status == 200 and token.get("access_token") and token.get("expires_in", 0) > 0, f"expires in {token.get('expires_in')} s")

# Compute Engine credentials of google-auth, the path used by the STS agent
for name in ("GOOGLE_APPLICATION_CREDENTIALS", "GOOGLE_API_CERTIFICATE_CONFIG"):
    os.environ.pop(name, None)
os.environ["GCE_METADATA_HOST"] = args.host
os.environ["GCE_METADATA_IP"] = args.host

import google.auth
import google.auth.transport.requests

credentials, adc_project_id = google.auth.default(scopes=["https://www.googleapis.com/auth/cloud-platform"])
credentials.refresh(google.auth.transport.requests.Request())
check("google-auth uses Compute Engine credentials", type(credentials).__module__ == "google.auth.compute_engine.credentials", type(credentials).__module__)
check("google-auth project ID", adc_project_id == project_id, str(adc_project_id))
check("google-auth service account", credentials.service_account_email == email, credentials.service_account_email)
check("google-auth token", credentials.valid, f"expires {credentials.expiry} UTC")

# Validity of the token at Google
status, _, body = get("https://oauth2.googleapis.com/tokeninfo?" + urllib.parse.urlencode({"access_token": credentials.token}))
info = json.loads(body) if status == 200 else {}
check("token accepted by Google", status == 200, f"scope {info.get('scope')}, expires in {info.get('expires_in')} s")

# Use of the token with Cloud Storage
if args.bucket:
    status, _, body = get(
        f"https://storage.googleapis.com/storage/v1/b/{urllib.parse.quote(args.bucket)}/o?maxResults=1",
        {"Authorization": f"Bearer {credentials.token}"},
    )
    check(f"list objects of {args.bucket}", status == 200, f"status {status}")

print(f"{failures} check(s) failed" if failures else "All checks passed")
sys.exit(1 if failures else 0)