Moved IMDS Emulator/Proxy sources to the sts-agent folder.

sts-agent/imds_emulator.py

@@ -0,0 +1,127 @@
+#!/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()