128 lines
5.2 KiB
Python
Executable File
128 lines
5.2 KiB
Python
Executable File
#!/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()
|