#!/usr/bin/env python3
"""Expose EC2 instance-profile credentials through the STS web-identity shape."""

import json
import os
import sys
import urllib.parse
import urllib.request
import uuid
from http.server import BaseHTTPRequestHandler, HTTPServer
from xml.sax.saxutils import escape

IMDS = "http://169.254.169.254/latest"
EXPECTED_ROLE = os.environ.get("CP9_EXPECTED_ROLE_ARN", "")
EXPECTED_TOKEN = "cp9-remote-ec2-k3d"


def metadata(path):
    token_request = urllib.request.Request(
        f"{IMDS}/api/token",
        method="PUT",
        headers={"X-aws-ec2-metadata-token-ttl-seconds": "60"},
    )
    token = urllib.request.urlopen(token_request, timeout=5).read().decode()
    request = urllib.request.Request(
        f"{IMDS}/meta-data/{path}",
        headers={"X-aws-ec2-metadata-token": token},
    )
    return urllib.request.urlopen(request, timeout=5).read().decode()


def assume_role_response(role_arn, session_name, web_identity_token):
    if not EXPECTED_ROLE or role_arn != EXPECTED_ROLE:
        raise ValueError("unexpected role")
    if web_identity_token != EXPECTED_TOKEN:
        raise ValueError("unexpected web identity token")

    profile = metadata("iam/security-credentials/").strip()
    credentials = json.loads(metadata(f"iam/security-credentials/{profile}"))
    required = ("AccessKeyId", "SecretAccessKey", "Token", "Expiration")
    if credentials.get("Code") != "Success" or not all(credentials.get(key) for key in required):
        raise RuntimeError("EC2 instance-profile credentials are unavailable")

    account = role_arn.split(":")[4]
    role_name = role_arn.rsplit("/", 1)[-1]
    assumed_role_arn = f"arn:aws:sts::{account}:assumed-role/{role_name}/{session_name}"
    return f"""<?xml version="1.0" encoding="UTF-8"?>
<AssumeRoleWithWebIdentityResponse xmlns="https://sts.amazonaws.com/doc/2011-06-15/">
  <AssumeRoleWithWebIdentityResult>
    <Audience>sts.amazonaws.com</Audience>
    <AssumedRoleUser>
      <Arn>{escape(assumed_role_arn)}</Arn>
      <AssumedRoleId>cp9:{escape(session_name)}</AssumedRoleId>
    </AssumedRoleUser>
    <Credentials>
      <AccessKeyId>{escape(credentials["AccessKeyId"])}</AccessKeyId>
      <SecretAccessKey>{escape(credentials["SecretAccessKey"])}</SecretAccessKey>
      <SessionToken>{escape(credentials["Token"])}</SessionToken>
      <Expiration>{escape(credentials["Expiration"])}</Expiration>
    </Credentials>
    <Provider>cp9-remote-ec2-k3d</Provider>
    <SubjectFromWebIdentityToken>{EXPECTED_TOKEN}</SubjectFromWebIdentityToken>
  </AssumeRoleWithWebIdentityResult>
  <ResponseMetadata><RequestId>{uuid.uuid4()}</RequestId></ResponseMetadata>
</AssumeRoleWithWebIdentityResponse>
"""


class Handler(BaseHTTPRequestHandler):
    def do_GET(self):
        self.send_response(200 if self.path == "/healthz" else 404)
        self.end_headers()

    def do_POST(self):
        try:
            length = int(self.headers.get("Content-Length", "0"))
            if length <= 0 or length > 16384:
                raise ValueError("invalid request size")
            form = urllib.parse.parse_qs(self.rfile.read(length).decode())
            if form.get("Action") != ["AssumeRoleWithWebIdentity"]:
                raise ValueError("unexpected action")
            payload = assume_role_response(
                form.get("RoleArn", [""])[0],
                form.get("RoleSessionName", [""])[0],
                form.get("WebIdentityToken", [""])[0],
            ).encode()
            self.send_response(200)
            self.send_header("Content-Type", "text/xml")
            self.send_header("Content-Length", str(len(payload)))
            self.end_headers()
            self.wfile.write(payload)
        except Exception as exc:
            payload = str(exc).encode()
            self.send_response(500)
            self.send_header("Content-Length", str(len(payload)))
            self.end_headers()
            self.wfile.write(payload)

    def log_message(self, _format, *_args):
        pass


def self_test():
    global EXPECTED_ROLE, metadata
    EXPECTED_ROLE = "arn:aws:iam::123456789012:role/cp9-runtime"
    values = {
        "iam/security-credentials/": "cp9-runtime",
        "iam/security-credentials/cp9-runtime": json.dumps(
            {
                "Code": "Success",
                "AccessKeyId": "AKIA_TEST",
                "SecretAccessKey": "secret",
                "Token": "token",
                "Expiration": "2030-01-01T00:00:00Z",
            }
        ),
    }
    metadata = values.__getitem__
    response = assume_role_response(EXPECTED_ROLE, "proof", EXPECTED_TOKEN)
    assert "<AccessKeyId>AKIA_TEST</AccessKeyId>" in response
    assert "assumed-role/cp9-runtime/proof" in response


if __name__ == "__main__":
    if sys.argv[1:] == ["--self-test"]:
        self_test()
    else:
        HTTPServer(("127.0.0.1", 8085), Handler).serve_forever()