Credential dispenser¶

# lambda_function.py — VTK credential dispenser
import os
import json
import boto3

STS = boto3.client("sts")

TARGET_ROLE_ARN = os.environ["TARGET_ROLE_ARN"]      # role in your account, trusts this Lambda
DEFAULT_DURATION = int(os.environ.get("DURATION_SECONDS", "3600"))


def lambda_handler(event, context):
    # API Gateway with AWS_IAM auth puts the caller's ARN here.
    caller = (
        event.get("requestContext", {})
             .get("identity", {})
             .get("userArn", "")
    )

    # Defence in depth: only let the VTK worker role through.
    if not caller.endswith(":role/VTKWorker") and "assumed-role/VTKWorker/" not in caller:
        return {"statusCode": 403, "body": json.dumps({"error": "forbidden"})}

    # Optional: tighten the returned permissions per call.
    # Inspect event["path"] or query/headers to decide which prefix to allow,
    # then attach an inline session policy here.
    session_policy = {
        "Version": "2012-10-17",
        "Statement": [
            {
                "Effect": "Allow",
                "Action": ["s3:GetObject", "s3:PutObject"],
                "Resource": "arn:aws:s3:::acme-bucket/*",
            },
            {
                "Effect": "Allow",
                "Action": ["s3:ListBucket", "s3:GetBucketLocation"],
                "Resource": "arn:aws:s3:::acme-bucket",
            },
        ],
    }

    resp = STS.assume_role(
        RoleArn=TARGET_ROLE_ARN,
        RoleSessionName=f"vtk-{context.aws_request_id[:30]}",
        DurationSeconds=DEFAULT_DURATION,
        Policy=json.dumps(session_policy),
    )
    creds = resp["Credentials"]

    return {
        "statusCode": 200,
        "headers": {"Content-Type": "application/json"},
        "body": json.dumps({
            "AccessKeyId":     creds["AccessKeyId"],
            "SecretAccessKey": creds["SecretAccessKey"],
            "SessionToken":    creds["SessionToken"],
            "Expiration":      creds["Expiration"].isoformat(),
        }),
    }

# vtk_dispenser_stack.py
from aws_cdk import (
    Stack,
    CfnOutput,
    Duration,
    aws_iam as iam,
    aws_lambda as _lambda,
    aws_apigateway as apigw,
)
from constructs import Construct


# The two role ARNs that VTK signs requests with. Both are required if you
# ever submit jobs to the staging environment; drop the staging one in
# production-only deployments.
VTK_WORKER_ROLE_ARNS = [
    "arn:aws:iam::873682911326:role/VTKWorker",  # production
]


class VtkDispenserStack(Stack):
    def __init__(
        self,
        scope: Construct,
        construct_id: str,
        *,
        target_bucket: str,
        session_duration: Duration = Duration.hours(1),
        **kwargs,
    ) -> None:
        super().__init__(scope, construct_id, **kwargs)

        # 1. Target role: scoped to the bucket. The Lambda assumes this and
        #    returns its session credentials to the VTK worker.
        target_role = iam.Role(
            self,
            "VtkStorageTargetRole",
            assumed_by=iam.AccountRootPrincipal(),  # tightened below
            max_session_duration=session_duration,
            inline_policies={
                "S3Access": iam.PolicyDocument(
                    statements=[
                        iam.PolicyStatement(
                            actions=["s3:GetObject", "s3:PutObject"],
                            resources=[f"arn:aws:s3:::{target_bucket}/*"],
                        ),
                        iam.PolicyStatement(
                            actions=["s3:ListBucket", "s3:GetBucketLocation"],
                            resources=[f"arn:aws:s3:::{target_bucket}"],
                        ),
                    ]
                )
            },
        )

        # 2. Lambda handler — the same code from the "Reference Lambda
        #    implementation" section above, with the target role + duration
        #    plumbed through as environment variables.
        handler = _lambda.Function(
            self,
            "DispenserFn",
            runtime=_lambda.Runtime.PYTHON_3_12,
            handler="index.lambda_handler",
            timeout=Duration.seconds(10),
            environment={
                "TARGET_ROLE_ARN": target_role.role_arn,
                "DURATION_SECONDS": str(int(session_duration.to_seconds())),
            },
            code=_lambda.Code.from_inline(
                """
import os, json, boto3

STS = boto3.client("sts")
TARGET_ROLE_ARN = os.environ["TARGET_ROLE_ARN"]
DURATION = int(os.environ.get("DURATION_SECONDS", "3600"))


def lambda_handler(event, context):
    caller = (
        event.get("requestContext", {})
        .get("identity", {})
        .get("userArn", "")
    )
    if "assumed-role/VTKWorker/" not in caller and not caller.endswith(":role/VTKWorker"):
        return {"statusCode": 403, "body": json.dumps({"error": "forbidden"})}

    resp = STS.assume_role(
        RoleArn=TARGET_ROLE_ARN,
        RoleSessionName=f"vtk-{context.aws_request_id[:30]}",
        DurationSeconds=DURATION,
    )
    c = resp["Credentials"]
    return {
        "statusCode": 200,
        "headers": {"Content-Type": "application/json"},
        "body": json.dumps({
            "AccessKeyId":     c["AccessKeyId"],
            "SecretAccessKey": c["SecretAccessKey"],
            "SessionToken":    c["SessionToken"],
            "Expiration":      c["Expiration"].isoformat(),
        }),
    }
"""
            ),
        )

        # Allow the Lambda to assume the target role …
        handler.add_to_role_policy(
            iam.PolicyStatement(
                actions=["sts:AssumeRole"],
                resources=[target_role.role_arn],
            )
        )
        # … and make the target role trust the Lambda's execution role.
        target_role.assume_role_policy.add_statements(
            iam.PolicyStatement(
                actions=["sts:AssumeRole"],
                principals=[iam.ArnPrincipal(handler.role.role_arn)],
            )
        )

        # 3. API Gateway REST API restricted to the VTK worker role(s).
        api = apigw.RestApi(
            self,
            "DispenserApi",
            rest_api_name="vtk-credential-dispenser",
            policy=iam.PolicyDocument(
                statements=[
                    iam.PolicyStatement(
                        effect=iam.Effect.ALLOW,
                        principals=[iam.ArnPrincipal(arn) for arn in VTK_WORKER_ROLE_ARNS],
                        actions=["execute-api:Invoke"],
                        resources=["execute-api:/*/GET/credentials"],
                    ),
                    iam.PolicyStatement(
                        effect=iam.Effect.DENY,
                        principals=[iam.AnyPrincipal()],
                        actions=["execute-api:Invoke"],
                        resources=["execute-api:/*/*/*"],
                        conditions={
                            "StringNotEquals": {
                                "aws:PrincipalArn": VTK_WORKER_ROLE_ARNS
                            }
                        },
                    ),
                ]
            ),
        )

        credentials = api.root.add_resource("credentials")
        credentials.add_method(
            "GET",
            apigw.LambdaIntegration(handler),
            authorization_type=apigw.AuthorizationType.IAM,
        )

        CfnOutput(
            self,
            "HttpCredentialSource",
            value=f"{api.url}credentials",
            description="Use as `http_credential_source` on storage:get / storage:put.",
        )

# app.py
import aws_cdk as cdk
from vtk_dispenser_stack import VtkDispenserStack

app = cdk.App()
VtkDispenserStack(
    app,
    "VtkDispenser",
    target_bucket="acme-bucket",
)
app.synth()

def lambda_handler(event, context):
    caller = event["requestContext"]["identity"].get("userArn", "")
    if "assumed-role/VTKWorker/" not in caller and not caller.endswith(":role/VTKWorker"):
        return {"statusCode": 403, "body": "forbidden"}

    token = event["pathParameters"]["token"]
    record = TABLE.get_item(Key={"token": token}).get("Item")
    if not record or record["expires_at"] < int(time.time()):
        return {"statusCode": 410, "body": "gone"}

    session_policy = build_policy(record["bucket"], record["prefix"], record["actions"])
    resp = STS.assume_role(
        RoleArn=TARGET_ROLE_ARN,
        RoleSessionName=f"vtk-{token[:30]}",
        DurationSeconds=record.get("duration", 3600),
        Policy=json.dumps(session_policy),
    )
    # Optional: mark the token "used" or decrement a counter here.
    ...

credentials = api.root.add_resource("credentials")
credentials.add_resource("{token}").add_method(
    "GET",
    apigw.LambdaIntegration(handler),
    authorization_type=apigw.AuthorizationType.IAM,
)