#!/usr/bin/env python3
"""
Compute Engine metadata server (IMDS) emulator.

Serves access tokens of the credentials found by Application Default
Credentials (ADC) on the IMDS token endpoint, together with the metadata
that Google client libraries read when they authenticate on Compute Engine.
"""
import datetime
import http.server
import json
import os
import socket
import sys
import threading
import urllib.parse

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

DEFAULT_SCOPE = "https://www.googleapis.com/auth/cloud-platform"
ROOT = "/computeMetadata/v1/"

port = int(os.environ.get("IMDS_PORT", "80"))
hostname = os.environ.get("IMDS_HOSTNAME", socket.gethostname())
project_number = os.environ.get("PROJECT_NUMBER", "")

base_credentials, adc_project_id = google.auth.default(scopes=[DEFAULT_SCOPE])
project_id = os.environ.get("PROJECT_ID") or adc_project_id or ""
email = getattr(base_credentials, "service_account_email", None) or os.environ.get("SERVICE_ACCOUNT_EMAIL", "")
if not email:
    sys.exit("The service account email is unknown, set SERVICE_ACCOUNT_EMAIL")

credentials_by_scopes = {}
lock = threading.Lock()


def token(scopes):
    """
    Returns the access token and its remaining lifetime in seconds for the
    scopes, refreshing the cached credentials of the scopes when needed.
    """
    key = tuple(sorted(scopes))
    with lock:
        credentials = credentials_by_scopes.get(key)
        if credentials is None:
            credentials = google.auth.credentials.with_scopes_if_required(base_credentials, list(key))
            credentials_by_scopes[key] = credentials
        if not credentials.valid:
            credentials.refresh(google.auth.transport.requests.Request())
        expires_in = int((credentials.expiry - datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)).total_seconds())
        return credentials.token, expires_in


class Handler(http.server.BaseHTTPRequestHandler):
    """
    Answers the IMDS paths used by the Google client libraries, and 404 for
    any other path.
    """

    server_version = "Metadata Server for VM"
    sys_version = ""

    def reply(self, status, body, content_type="application/text"):
        data = body.encode() if isinstance(body, str) else json.dumps(body).encode()
        self.send_response(status)
        self.send_header("Metadata-Flavor", "Google")
        self.send_header("Content-Type", content_type)
        self.send_header("Content-Length", str(len(data)))
        self.end_headers()
        self.wfile.write(data)

    def do_GET(self):
        url = urllib.parse.urlsplit(self.path)
        query = urllib.parse.parse_qs(url.query)

        # The real IMDS rejects requests without the flavor header and proxied requests.
        if self.headers.get("Metadata-Flavor") != "Google" or self.headers.get("X-Forwarded-For"):
            return self.reply(403, "Forbidden")

        if url.path in ("/", ROOT):
            return self.reply(200, "")

        if not url.path.startswith(ROOT):
            return self.reply(404, "Not Found")
        path = url.path[len(ROOT):]

        # Both "default" and the email address name the service account.
        for alias in ("default", email):
            prefix = f"instance/service-accounts/{alias}/"
            if path.startswith(prefix):
                path = "instance/service-accounts/default/" + path[len(prefix):]
                break

        account = {"aliases": ["default"], "email": email, "scopes": [DEFAULT_SCOPE]}
        if path == "instance/service-accounts/default/token":
            scopes = query.get("scopes", [DEFAULT_SCOPE])[0].split(",")
            access_token, expires_in = token(scopes)
            return self.reply(200, {"access_token": access_token, "expires_in": expires_in, "token_type": "Bearer"}, "application/json")
        if path == "instance/service-accounts/default/email":
            return self.reply(200, email)
        if path == "instance/service-accounts/default/scopes":
            return self.reply(200, DEFAULT_SCOPE + "\n")
        if path == "instance/service-accounts/default/" and query.get("recursive") == ["true"]:
            return self.reply(200, account, "application/json")
        if path == "instance/service-accounts/" and query.get("recursive") == ["true"]:
            return self.reply(200, {"default": account, email: account}, "application/json")
        if path == "instance/service-accounts/":
            return self.reply(200, f"default/\n{email}/\n")
        if path == "project/project-id" and project_id:
            return self.reply(200, project_id)
        if path == "project/numeric-project-id" and project_number:
            return self.reply(200, project_number)
        if path == "instance/hostname":
            return self.reply(200, hostname)
        if path == "universe/universe-domain":
            return self.reply(200, "googleapis.com")

        return self.reply(404, "Not Found")

    def log_request(self, code="-", size="-"):
        print(f"{self.command} {urllib.parse.urlsplit(self.path).path} {code}", flush=True)


print(f"Serving tokens of {email} for project {project_id or '-'} on port {port}", flush=True)
http.server.ThreadingHTTPServer(("", port), Handler).serve_forever()