iam-session-tag-tenant-scope

{
  "query": "select:TodoWrite",
  "max_results": 1
}

, output, 
{"type": "tool_reference", "tool_name": "TodoWrite"}

import json
import os
import re

import boto3

DATA_ROLE_ARN = os.environ["DATA_ROLE_ARN"]
ALLOWED_TENANTS = {"acme", "globex", "globex-eu"}
TENANT_RE = re.compile(r"^[a-z][a-z0-9-]*$")

_sts = boto3.client("sts")


def _bad(msg):
    return {
        "statusCode": 400,
        "headers": {"Content-Type": "application/json"},
        "body": json.dumps({"error": msg}),
    }


def handler(event, context):
    qs = event.get("queryStringParameters") or {}
    tenant = (qs.get("tenant") or "").strip()

    if not tenant or not TENANT_RE.match(tenant) or tenant not in ALLOWED_TENANTS:
        return _bad("invalid tenant")

    resp = _sts.assume_role(
        RoleArn=DATA_ROLE_ARN,
        RoleSessionName=f"tenant-{tenant}",
        DurationSeconds=900,
        Tags=[{"Key": "TenantID", "Value": tenant}],
        TransitiveTagKeys=["TenantID"],
    )
    creds = resp["Credentials"]
    body = {
        "AccessKeyId": creds["AccessKeyId"],
        "SecretAccessKey": creds["SecretAccessKey"],
        "SessionToken": creds["SessionToken"],
        "Expiration": creds["Expiration"].isoformat(),
        "Tenant": tenant,
    }
    return {
        "statusCode": 200,
        "headers": {"Content-Type": "application/json"},
        "body": json.dumps(body),
    }

"""
End-to-end build for the harbor-saas multi-tenant ABAC stack on LocalStack.

Resources created:
  - KMS CMK + alias/saas-orders-cmk
  - DynamoDB table SaasOrders (PITR on, encrypted with the CMK)
  - 6 seed rows: 2 each for acme / globex / globex-eu (with sensitive `password`)
  - IAM role TenantTokenVendorRole (lambda exec) with scoped inline log + sts:AssumeRole
  - IAM role TenantDataRole with ABAC trust (sts:AssumeRole + sts:TagSession) and
    a session-tag substitution identity policy gated to the tenant's rows
    and to non-`password` attributes
  - Lambda tenant-token-vendor (python3.11) wired to the vendor role
  - API Gateway (v1) REST api harbor-saas-api with /token GET -> AWS_PROXY,
    stage `prod` with access logging + lambda:InvokeFunction permission
    pinned to this api's SourceArn
  - SSM parameters under /harbor/saas/* so the verifier doesn't have to guess
"""

import io
import json
import os
import time
import zipfile

import boto3

REGION = os.environ.get("AWS_DEFAULT_REGION", "us-east-1")
ENDPOINT = os.environ.get("AWS_ENDPOINT_URL", "http://localstack:4566")
ACCOUNT = "000000000000"

TABLE = "SaasOrders"
KMS_ALIAS = "alias/saas-orders-cmk"
DATA_ROLE = "TenantDataRole"
VENDOR_ROLE = "TenantTokenVendorRole"
LAMBDA_NAME = "tenant-token-vendor"
LAMBDA_LOG_GROUP = f"/aws/lambda/{LAMBDA_NAME}"
APIGW_NAME = "harbor-saas-api"
APIGW_STAGE = "prod"
APIGW_ACCESS_LOG_GROUP = f"/aws/apigateway/{APIGW_NAME}-access"
ALLOWED_TENANTS = ["acme", "globex", "globex-eu"]

session = boto3.session.Session(region_name=REGION)
kw = dict(endpoint_url=ENDPOINT, region_name=REGION)
kms = session.client("kms", **kw)
ddb = session.client("dynamodb", **kw)
iam = session.client("iam", **kw)
lam = session.client("lambda", **kw)
logs = session.client("logs", **kw)
apigw = session.client("apigateway", **kw)
ssm = session.client("ssm", **kw)


def banner(msg):
    print(f"\n=== {msg} ===")


# ---------------------------------------------------------------------------
# 1. KMS CMK + alias
# ---------------------------------------------------------------------------
def create_kms():
    banner("KMS CMK")
    aliases = kms.list_aliases().get("Aliases", [])
    for a in aliases:
        if a["AliasName"] == KMS_ALIAS:
            print(f"  alias {KMS_ALIAS} -> {a.get('TargetKeyId')} (exists)")
            return a["TargetKeyId"]
    key = kms.create_key(
        Description="Customer-managed CMK for SaasOrders DDB table",
        KeyUsage="ENCRYPT_DECRYPT",
    )["KeyMetadata"]
    kms.create_alias(AliasName=KMS_ALIAS, TargetKeyId=key["KeyId"])
    print(f"  created {KMS_ALIAS} -> {key['KeyId']}")
    return key["KeyId"]


# ---------------------------------------------------------------------------
# 2. DynamoDB SaasOrders + PITR
# ---------------------------------------------------------------------------
def create_table(kms_key_id):
    banner("DynamoDB SaasOrders")
    try:
        ddb.describe_table(TableName=TABLE)
        print(f"  {TABLE} already exists")
    except ddb.exceptions.ResourceNotFoundException:
        ddb.create_table(
            TableName=TABLE,
            AttributeDefinitions=[
                {"AttributeName": "TenantID", "AttributeType": "S"},
                {"AttributeName": "OrderID", "AttributeType": "S"},
            ],
            KeySchema=[
                {"AttributeName": "TenantID", "KeyType": "HASH"},
                {"AttributeName": "OrderID", "KeyType": "RANGE"},
            ],
            BillingMode="PAY_PER_REQUEST",
            SSESpecification={
                "Enabled": True,
                "SSEType": "KMS",
                "KMSMasterKeyId": kms_key_id,
            },
        )
        waiter = ddb.get_waiter("table_exists")
        waiter.wait(TableName=TABLE)
        print(f"  created {TABLE}")

    ddb.update_continuous_backups(
        TableName=TABLE,
        PointInTimeRecoverySpecification={"PointInTimeRecoveryEnabled": True},
    )
    print("  PITR enabled")


def seed_items():
    banner("Seed rows")
    rows = []
    for t in ALLOWED_TENANTS:
        for i in (1, 2):
            rows.append(
                {
                    "TenantID": {"S": t},
                    "OrderID": {"S": f"order-{i:03d}"},
                    "amount": {"N": str(100 * i)},
                    "password": {"S": f"{t}-secret-{i}"},
                }
            )
    for r in rows:
        ddb.put_item(TableName=TABLE, Item=r)
    print(f"  wrote {len(rows)} items")


# ---------------------------------------------------------------------------
# 3. IAM roles
# ---------------------------------------------------------------------------
def _ensure_role(name, trust_doc):
    try:
        iam.get_role(RoleName=name)
        iam.update_assume_role_policy(
            RoleName=name, PolicyDocument=json.dumps(trust_doc)
        )
        print(f"  role {name} (updated trust)")
    except iam.exceptions.NoSuchEntityException:
        iam.create_role(
            RoleName=name, AssumeRolePolicyDocument=json.dumps(trust_doc)
        )
        print(f"  role {name} (created)")


def vendor_role():
    banner("Vendor exec role")
    trust = {
        "Version": "2012-10-17",
        "Statement": [
            {
                "Effect": "Allow",
                "Principal": {"Service": "lambda.amazonaws.com"},
                "Action": "sts:AssumeRole",
            }
        ],
    }
    _ensure_role(VENDOR_ROLE, trust)

    # Scoped log-write to the function's own log group only.
    log_arn = f"arn:aws:logs:{REGION}:{ACCOUNT}:log-group:{LAMBDA_LOG_GROUP}:*"
    log_pol = {
        "Version": "2012-10-17",
        "Statement": [
            {
                "Effect": "Allow",
                "Action": ["logs:CreateLogStream", "logs:PutLogEvents"],
                "Resource": log_arn,
            }
        ],
    }
    iam.put_role_policy(
        RoleName=VENDOR_ROLE,
        PolicyName="LambdaOwnLogs",
        PolicyDocument=json.dumps(log_pol),
    )

    # Single sts:AssumeRole permission against TenantDataRole only (no wildcards).
    data_role_arn = f"arn:aws:iam::{ACCOUNT}:role/{DATA_ROLE}"
    assume_pol = {
        "Version": "2012-10-17",
        "Statement": [
            {
                "Effect": "Allow",
                "Action": "sts:AssumeRole",
                "Resource": data_role_arn,
            }
        ],
    }
    iam.put_role_policy(
        RoleName=VENDOR_ROLE,
        PolicyName="AssumeTenantDataRoleOnly",
        PolicyDocument=json.dumps(assume_pol),
    )
    return iam.get_role(RoleName=VENDOR_ROLE)["Role"]["Arn"]


def data_role(vendor_role_arn):
    banner("Data role")
    # Trust: principal is the vendor lambda exec role only.
    # Action set is BOTH sts:AssumeRole and sts:TagSession -- omitting
    # TagSession would silently drop the tag, so isolation would evaporate.
    # Conditions:
    #   - Null on aws:RequestTag/TenantID = "false" forces the tag to be
    #     present at all (an allowlist alone is not enough).
    #   - StringEquals on aws:RequestTag/TenantID restricts which tenant
    #     values are accepted (no wildcards).
    trust = {
        "Version": "2012-10-17",
        "Statement": [
            {
                "Sid": "VendorAssumesWithTenantTag",
                "Effect": "Allow",
                "Principal": {"AWS": vendor_role_arn},
                "Action": ["sts:AssumeRole", "sts:TagSession"],
                "Condition": {
                    "Null": {"aws:RequestTag/TenantID": "false"},
                    "StringEquals": {
                        "aws:RequestTag/TenantID": ALLOWED_TENANTS
                    },
                },
            }
        ],
    }
    _ensure_role(DATA_ROLE, trust)

    table_arn = f"arn:aws:dynamodb:{REGION}:{ACCOUNT}:table/{TABLE}"
    # Identity policy:
    #   - Action set is explicit, on the table arn -- no `*`, no Scan.
    #   - dynamodb:LeadingKeys uses ForAllValues:StringEquals (multi-valued
    #     condition key; using plain StringEquals fails closed for everyone).
    #   - The substitution is ${aws:PrincipalTag/TenantID}: PrincipalTag is
    #     the tag carried by the resulting principal. RequestTag would only
    #     exist at AssumeRole-time and silently fail open here.
    #   - dynamodb:Attributes restricts which columns are touchable, blocking
    #     the analyst from `password`. dynamodb:Select=SPECIFIC_ATTRIBUTES
    #     is enforced where it applies (Query/Scan).
    identity = {
        "Version": "2012-10-17",
        "Statement": [
            {
                "Sid": "TenantScopedRowAndColumn",
                "Effect": "Allow",
                "Action": [
                    "dynamodb:GetItem",
                    "dynamodb:BatchGetItem",
                    "dynamodb:Query",
                    "dynamodb:PutItem",
                    "dynamodb:UpdateItem",
                    "dynamodb:DeleteItem",
                    "dynamodb:BatchWriteItem",
                ],
                "Resource": table_arn,
                "Condition": {
                    "ForAllValues:StringEquals": {
                        "dynamodb:LeadingKeys": [
                            "${aws:PrincipalTag/TenantID}"
                        ],
                        "dynamodb:Attributes": [
                            "TenantID",
                            "OrderID",
                            "amount",
                        ],
                    },
                    "StringEqualsIfExists": {
                        "dynamodb:Select": "SPECIFIC_ATTRIBUTES"
                    },
                },
            }
        ],
    }
    iam.put_role_policy(
        RoleName=DATA_ROLE,
        PolicyName="TenantScopedTableAccess",
        PolicyDocument=json.dumps(identity),
    )
    return iam.get_role(RoleName=DATA_ROLE)["Role"]["Arn"]


# ---------------------------------------------------------------------------
# 4. Lambda
# ---------------------------------------------------------------------------
def package_lambda():
    buf = io.BytesIO()
    with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as z:
        with open("/app/lambda_src/handler.py", "rb") as f:
            z.writestr("handler.py", f.read())
    return buf.getvalue()


def deploy_lambda(vendor_role_arn, data_role_arn):
    banner("Lambda")
    # Make sure the function's own log group exists -- the role is scoped to it.
    try:
        logs.create_log_group(logGroupName=LAMBDA_LOG_GROUP)
    except logs.exceptions.ResourceAlreadyExistsException:
        pass

    code = package_lambda()
    env = {"Variables": {"DATA_ROLE_ARN": data_role_arn}}

    try:
        lam.get_function(FunctionName=LAMBDA_NAME)
        lam.update_function_code(FunctionName=LAMBDA_NAME, ZipFile=code)
        lam.get_waiter("function_updated").wait(FunctionName=LAMBDA_NAME)
        lam.update_function_configuration(
            FunctionName=LAMBDA_NAME,
            Role=vendor_role_arn,
            Environment=env,
        )
        lam.get_waiter("function_updated").wait(FunctionName=LAMBDA_NAME)
        print("  updated existing function")
    except lam.exceptions.ResourceNotFoundException:
        # IAM eventual consistency tends to bite create_function -- retry briefly.
        last_exc = None
        for _ in range(10):
            try:
                lam.create_function(
                    FunctionName=LAMBDA_NAME,
                    Runtime="python3.11",
                    Role=vendor_role_arn,
                    Handler="handler.handler",
                    Code={"ZipFile": code},
                    Timeout=10,
                    Environment=env,
                )
                break
            except lam.exceptions.InvalidParameterValueException as e:
                last_exc = e
                time.sleep(1)
        else:
            raise last_exc
        lam.get_waiter("function_active").wait(FunctionName=LAMBDA_NAME)
        print("  created function")

    return lam.get_function_configuration(FunctionName=LAMBDA_NAME)[
        "FunctionArn"
    ]


# ---------------------------------------------------------------------------
# 5. API Gateway
# ---------------------------------------------------------------------------
def find_or_create_api():
    apis = apigw.get_rest_apis().get("items", [])
    for a in apis:
        if a["name"] == APIGW_NAME:
            return a["id"]
    return apigw.create_rest_api(
        name=APIGW_NAME,
        description="harbor SaaS token-vendor API",
        endpointConfiguration={"types": ["REGIONAL"]},
    )["id"]


def wire_apigw(lambda_arn):
    banner("API Gateway")
    api_id = find_or_create_api()
    print(f"  api id {api_id}")

    resources = apigw.get_resources(restApiId=api_id)["items"]
    root_id = next(r["id"] for r in resources if r["path"] == "/")
    token_resource = next(
        (r for r in resources if r["path"] == "/token"), None
    )
    if token_resource is None:
        token_resource = apigw.create_resource(
            restApiId=api_id, parentId=root_id, pathPart="token"
        )
    token_id = token_resource["id"]

    try:
        apigw.put_method(
            restApiId=api_id,
            resourceId=token_id,
            httpMethod="GET",
            authorizationType="NONE",
        )
    except apigw.exceptions.ConflictException:
        pass

    integration_uri = (
        f"arn:aws:apigateway:{REGION}:lambda:path/2015-03-31/functions/"
        f"{lambda_arn}/invocations"
    )
    apigw.put_integration(
        restApiId=api_id,
        resourceId=token_id,
        httpMethod="GET",
        type="AWS_PROXY",
        integrationHttpMethod="POST",
        uri=integration_uri,
    )

    # access log group + stage with access logging
    try:
        logs.create_log_group(logGroupName=APIGW_ACCESS_LOG_GROUP)
    except logs.exceptions.ResourceAlreadyExistsException:
        pass

    # deploy + stage. APIGW config changes don't promote without redeploy.
    deployment = apigw.create_deployment(
        restApiId=api_id, stageName=APIGW_STAGE
    )
    print(f"  deployment {deployment['id']} -> stage {APIGW_STAGE}")

    access_log_arn = (
        f"arn:aws:logs:{REGION}:{ACCOUNT}:log-group:{APIGW_ACCESS_LOG_GROUP}"
    )
    access_log_format = json.dumps(
        {
            "requestId": "$context.requestId",
            "ip": "$context.identity.sourceIp",
            "method": "$context.httpMethod",
            "resourcePath": "$context.resourcePath",
            "status": "$context.status",
            "responseLength": "$context.responseLength",
        }
    )
    apigw.update_stage(
        restApiId=api_id,
        stageName=APIGW_STAGE,
        patchOperations=[
            {
                "op": "replace",
                "path": "/accessLogSettings/destinationArn",
                "value": access_log_arn,
            },
            {
                "op": "replace",
                "path": "/accessLogSettings/format",
                "value": access_log_format,
            },
        ],
    )
    print(f"  access logging -> {APIGW_ACCESS_LOG_GROUP}")

    # invoke permission, pinned to THIS api's resource arn
    source_arn = (
        f"arn:aws:execute-api:{REGION}:{ACCOUNT}:{api_id}/*/GET/token"
    )
    statement_id = "apigw-invoke-token"
    try:
        lam.remove_permission(
            FunctionName=LAMBDA_NAME, StatementId=statement_id
        )
    except lam.exceptions.ResourceNotFoundException:
        pass
    lam.add_permission(
        FunctionName=LAMBDA_NAME,
        StatementId=statement_id,
        Action="lambda:InvokeFunction",
        Principal="apigateway.amazonaws.com",
        SourceArn=source_arn,
    )
    print(f"  invoke permission pinned to {source_arn}")

    invoke_url = (
        f"http://localhost:4566/restapis/{api_id}/{APIGW_STAGE}/_user_request_/token"
    )
    return api_id, invoke_url


# ---------------------------------------------------------------------------
# 6. SSM pointers
# ---------------------------------------------------------------------------
def write_ssm(values):
    banner("SSM pointers")
    for name, value in values.items():
        ssm.put_parameter(
            Name=name, Value=value, Type="String", Overwrite=True
        )
        print(f"  {name} = {value}")


def main():
    kms_key_id = create_kms()
    create_table(kms_key_id)
    seed_items()
    vendor_arn = vendor_role()
    data_arn = data_role(vendor_arn)
    # Lambda needs both roles to exist. IAM is eventually consistent in
    # LocalStack; the deploy_lambda helper retries on InvalidParameterValue.
    lambda_arn = deploy_lambda(vendor_arn, data_arn)
    api_id, invoke_url = wire_apigw(lambda_arn)

    write_ssm(
        {
            "/harbor/saas/table-name": TABLE,
            "/harbor/saas/role-arn": data_arn,
            "/harbor/saas/vendor-role-arn": vendor_arn,
            "/harbor/saas/lambda-arn": lambda_arn,
            "/harbor/saas/api-id": api_id,
            "/harbor/saas/api-url": invoke_url,
        }
    )

    print("\nDONE.")
    print(f"  invoke: curl {invoke_url}?tenant=acme")


if __name__ == "__main__":
    main()

"""End-to-end smoke test for the harbor-saas ABAC stack."""
import json
import os
import urllib.request

import boto3
from botocore.exceptions import ClientError

REGION = os.environ.get("AWS_DEFAULT_REGION", "us-east-1")
ENDPOINT = os.environ.get("AWS_ENDPOINT_URL", "http://localstack:4566")

ssm = boto3.client("ssm", endpoint_url=ENDPOINT, region_name=REGION)
api_url = ssm.get_parameter(Name="/harbor/saas/api-url")["Parameter"]["Value"]


def http_get(url):
    with urllib.request.urlopen(url) as r:
        return r.status, json.loads(r.read().decode())


def with_creds(creds):
    return boto3.client(
        "dynamodb",
        endpoint_url=ENDPOINT,
        region_name=REGION,
        aws_access_key_id=creds["AccessKeyId"],
        aws_secret_access_key=creds["SecretAccessKey"],
        aws_session_token=creds["SessionToken"],
    )


print("[1] GET /token?tenant=acme")
status, body = http_get(f"{api_url}?tenant=acme")
assert status == 200, status
assert {"AccessKeyId", "SecretAccessKey", "SessionToken"} <= set(body), body
print(f"    OK -- 200, got AccessKeyId={body['AccessKeyId'][:8]}...")
acme_creds = body

print("[2] GET /token?tenant=../../etc/passwd  (rejected)")
import urllib.parse

req = urllib.request.Request(
    f"{api_url}?tenant={urllib.parse.quote('../../etc/passwd')}"
)
try:
    urllib.request.urlopen(req)
    raise AssertionError("expected 400")
except urllib.error.HTTPError as e:
    assert e.code == 400, e.code
    print("    OK -- 400")

print("[3] acme creds Query its own rows, projecting only allowed cols")
ddb_acme = with_creds(acme_creds)
out = ddb_acme.query(
    TableName="SaasOrders",
    KeyConditionExpression="TenantID = :t",
    ExpressionAttributeValues={":t": {"S": "acme"}},
    ProjectionExpression="TenantID, OrderID, amount",
    Select="SPECIFIC_ATTRIBUTES",
)
items = out["Items"]
assert len(items) == 2, items
assert all("password" not in it for it in items), items
print(f"    OK -- got {len(items)} items, no password leaked")

print("[4] acme creds tries to read globex rows (cross-tenant)")
try:
    ddb_acme.query(
        TableName="SaasOrders",
        KeyConditionExpression="TenantID = :t",
        ExpressionAttributeValues={":t": {"S": "globex"}},
        ProjectionExpression="TenantID, OrderID, amount",
        Select="SPECIFIC_ATTRIBUTES",
    )
    raise AssertionError("cross-tenant query unexpectedly succeeded")
except ClientError as e:
    code = e.response["Error"]["Code"]
    assert code in {"AccessDeniedException", "AccessDenied"}, code
    print(f"    OK -- denied ({code})")

print("[5] acme creds tries to read its own password column")
try:
    ddb_acme.get_item(
        TableName="SaasOrders",
        Key={"TenantID": {"S": "acme"}, "OrderID": {"S": "order-001"}},
        ProjectionExpression="password",
    )
    raise AssertionError("password projection unexpectedly succeeded")
except ClientError as e:
    code = e.response["Error"]["Code"]
    assert code in {"AccessDeniedException", "AccessDenied"}, code
    print(f"    OK -- denied ({code})")

print("[6] globex-eu round-trip")
status, body = http_get(f"{api_url}?tenant=globex-eu")
assert status == 200
ddb_eu = with_creds(body)
out = ddb_eu.query(
    TableName="SaasOrders",
    KeyConditionExpression="TenantID = :t",
    ExpressionAttributeValues={":t": {"S": "globex-eu"}},
    ProjectionExpression="TenantID, OrderID, amount",
    Select="SPECIFIC_ATTRIBUTES",
)
assert len(out["Items"]) == 2
print("    OK -- globex-eu sees 2 of its own rows")

print("\nALL SMOKE TESTS PASSED.")

"""End-to-end smoke for the harbor-saas ABAC stack.

Two halves, mirroring the verifier description:
  1. Walk the trust + identity policy JSON to catch shape mistakes that an
     auditor would flag.
  2. Confirm the wiring works end-to-end via a real GET /token?tenant=acme
     round-trip and exercise the resulting credentials against DDB.

Note on LocalStack: ENFORCE_IAM honors action+resource matching for assumed
roles (e.g., Scan and CreateTable on the data role correctly fail), but the
DynamoDB-specific context keys (dynamodb:LeadingKeys, dynamodb:Attributes)
aren't evaluated at runtime in this environment. The auditor's check is on
policy SHAPE, which is what the assertions below cover.
"""
import json
import os
import urllib.request
import urllib.parse

import boto3
from botocore.exceptions import ClientError

REGION = os.environ.get("AWS_DEFAULT_REGION", "us-east-1")
ENDPOINT = os.environ.get("AWS_ENDPOINT_URL", "http://localstack:4566")
ACCOUNT = "000000000000"

ssm = boto3.client("ssm", endpoint_url=ENDPOINT, region_name=REGION)
iam = boto3.client("iam", endpoint_url=ENDPOINT, region_name=REGION)


def need(name):
    return ssm.get_parameter(Name=name)["Parameter"]["Value"]


def section(s):
    print(f"\n--- {s} ---")


# ---------------------------------------------------------------------------
# Half 1: SSM pointers exist
# ---------------------------------------------------------------------------
section("SSM pointers")
required = [
    "/harbor/saas/table-name",
    "/harbor/saas/role-arn",
    "/harbor/saas/vendor-role-arn",
    "/harbor/saas/lambda-arn",
    "/harbor/saas/api-id",
    "/harbor/saas/api-url",
]
for n in required:
    print(f"  {n} = {need(n)}")

# ---------------------------------------------------------------------------
# Half 1: Walk policy JSON for the audited shape
# ---------------------------------------------------------------------------
section("TenantDataRole trust policy shape")
trust = iam.get_role(RoleName="TenantDataRole")["Role"]["AssumeRolePolicyDocument"]
stmt = trust["Statement"][0]

actions = stmt["Action"] if isinstance(stmt["Action"], list) else [stmt["Action"]]
assert set(actions) == {"sts:AssumeRole", "sts:TagSession"}, actions
print("  trap 1: actions == {sts:AssumeRole, sts:TagSession} OK")

cond = stmt["Condition"]
assert cond["Null"]["aws:RequestTag/TenantID"] in ("false", False), cond["Null"]
print("  trap 2: Null/aws:RequestTag/TenantID = false (tag must exist) OK")

allowed = cond["StringEquals"]["aws:RequestTag/TenantID"]
allowed = allowed if isinstance(allowed, list) else [allowed]
assert "*" not in allowed and set(allowed) == {"acme", "globex", "globex-eu"}, allowed
print(f"  trap 3: tenant allowlist = {sorted(allowed)} (no wildcard) OK")

vendor_arn = need("/harbor/saas/vendor-role-arn")
principal = stmt["Principal"]
assert principal == {"AWS": vendor_arn}, principal
print(f"  principal pinned to vendor exec role only OK")

section("TenantDataRole identity policy shape")
ident = iam.get_role_policy(
    RoleName="TenantDataRole", PolicyName="TenantScopedTableAccess"
)["PolicyDocument"]
istmt = ident["Statement"][0]
table_arn = f"arn:aws:dynamodb:{REGION}:{ACCOUNT}:table/SaasOrders"
assert istmt["Resource"] == table_arn, istmt["Resource"]
assert "*" not in istmt["Action"], istmt["Action"]
assert "dynamodb:Scan" not in istmt["Action"], "Scan must not be permitted"
print("  scoped to table arn, no '*', no Scan OK")

icond = istmt["Condition"]
fav = icond["ForAllValues:StringEquals"]
assert fav["dynamodb:LeadingKeys"] == ["${aws:PrincipalTag/TenantID}"], fav
print("  ForAllValues:StringEquals on dynamodb:LeadingKeys with ${aws:PrincipalTag/TenantID} OK")

attrs = fav["dynamodb:Attributes"]
assert "password" not in attrs, attrs
print(f"  dynamodb:Attributes allowlist = {attrs} (password absent) OK")

section("TenantTokenVendorRole shape")
vendor_pols = iam.list_role_policies(RoleName="TenantTokenVendorRole")["PolicyNames"]
assert "AssumeTenantDataRoleOnly" in vendor_pols
assume_doc = iam.get_role_policy(
    RoleName="TenantTokenVendorRole", PolicyName="AssumeTenantDataRoleOnly"
)["PolicyDocument"]
assume_stmt = assume_doc["Statement"][0]
assert assume_stmt["Action"] == "sts:AssumeRole"
assert assume_stmt["Resource"] == need("/harbor/saas/role-arn")
assert "*" not in assume_stmt["Resource"]
print("  assume policy points at TenantDataRole only, no wildcard OK")
attached = iam.list_attached_role_policies(RoleName="TenantTokenVendorRole")[
    "AttachedPolicies"
]
assert all(
    p["PolicyName"] != "AWSLambdaBasicExecutionRole" for p in attached
), attached
print("  AWSLambdaBasicExecutionRole NOT attached (scoped log-write inline) OK")

# ---------------------------------------------------------------------------
# Half 2: real round-trip
# ---------------------------------------------------------------------------
section("API Gateway round-trip")
api_url = need("/harbor/saas/api-url")
with urllib.request.urlopen(f"{api_url}?tenant=acme") as r:
    body = json.loads(r.read())
assert {"AccessKeyId", "SecretAccessKey", "SessionToken", "Expiration"} <= set(
    body
), body
print(f"  GET /token?tenant=acme -> 200 with creds (key {body['AccessKeyId'][:8]}...)")
acme_creds = body

bad_url = f"{api_url}?tenant={urllib.parse.quote('../../etc/passwd')}"
try:
    urllib.request.urlopen(bad_url)
    raise AssertionError("expected 400")
except urllib.error.HTTPError as e:
    assert e.code == 400
    print("  GET /token?tenant=<bad> -> 400 OK")

section("ENFORCE_IAM checks (action set on assumed role)")
ddb = boto3.client(
    "dynamodb",
    endpoint_url=ENDPOINT,
    region_name=REGION,
    aws_access_key_id=acme_creds["AccessKeyId"],
    aws_secret_access_key=acme_creds["SecretAccessKey"],
    aws_session_token=acme_creds["SessionToken"],
)
for label, call in [
    ("Scan", lambda: ddb.scan(TableName="SaasOrders", Limit=1)),
    (
        "CreateTable",
        lambda: ddb.create_table(
            TableName="ShouldNotExist",
            AttributeDefinitions=[{"AttributeName": "pk", "AttributeType": "S"}],
            KeySchema=[{"AttributeName": "pk", "KeyType": "HASH"}],
            BillingMode="PAY_PER_REQUEST",
        ),
    ),
]:
    try:
        call()
        raise AssertionError(f"{label} unexpectedly allowed")
    except ClientError as e:
        assert e.response["Error"]["Code"] in {
            "AccessDeniedException",
            "AccessDenied",
        }, e
        print(f"  {label} denied OK")

section("Caller identity is the assumed role (session-tag aware)")
sts = boto3.client(
    "sts",
    endpoint_url=ENDPOINT,
    region_name=REGION,
    aws_access_key_id=acme_creds["AccessKeyId"],
    aws_secret_access_key=acme_creds["SecretAccessKey"],
    aws_session_token=acme_creds["SessionToken"],
)
caller = sts.get_caller_identity()["Arn"]
assert caller.endswith(":assumed-role/TenantDataRole/tenant-acme"), caller
print(f"  {caller} OK")

print("\nALL CHECKS PASSED.")