scripts/create_iam_role.py

"""Utility script to provision S3 read-only credentials for the Streamlit app.

This helper creates (or reuses) an IAM user and optional role that have
`s3:ListBucket` / `s3:GetObject` access to the `chaptive-rag` bucket. It then
prints the access key and secret so they can be set as environment variables for the Streamlit app.

Usage:

```bash
python streamlit/scripts/create_iam_role.py \
  --entity-name streamlit-cache-reader \
  --bucket chaptive-rag
```

The AWS account credentials used to run this script must already have IAM
permissions to create users/roles, attach policies, and create access keys.
"""

from __future__ import annotations

import argparse
import json
import sys
from typing import Any, Dict, Optional

import boto3
from botocore.exceptions import ClientError

S3_POLICY_NAME = "StreamlitS3ReadOnly"
DEFAULT_ENTITY_NAME = "streamlit-cache-reader"
DEFAULT_BUCKET = "chaptly-rag"


def _build_policy(bucket: str) -> Dict[str, Any]:
	resource_arn = f"arn:aws:s3:::{bucket}"
	object_arn = f"{resource_arn}/*"
	return {
		"Version": "2012-10-17",
		"Statement": [
			{
				"Effect": "Allow",
				"Action": ["s3:ListBucket"],
				"Resource": resource_arn,
			},
			{
				"Effect": "Allow",
				"Action": ["s3:GetObject"],
				"Resource": object_arn,
			},
		],
	}


def ensure_user(iam, user_name: str) -> None:
	try:
		iam.get_user(UserName=user_name)
		print(f"Reusing IAM user '{user_name}'.")
	except ClientError as exc:
		if exc.response["Error"]["Code"] == "NoSuchEntity":
			iam.create_user(UserName=user_name)
			print(f"Created IAM user '{user_name}'.")
		else:
			raise


def attach_inline_policy(iam, user_name: str, bucket: str) -> None:
	policy_doc = json.dumps(_build_policy(bucket))
	iam.put_user_policy(UserName=user_name, PolicyName=S3_POLICY_NAME, PolicyDocument=policy_doc)
	print(f"Attached inline S3 read policy '{S3_POLICY_NAME}' to user '{user_name}'.")


def create_access_key(iam, user_name: str) -> Dict[str, str]:
	response = iam.create_access_key(UserName=user_name)
	access_key = response["AccessKey"]
	return {
		"AWS_ACCESS_KEY_ID": access_key["AccessKeyId"],
		"AWS_SECRET_ACCESS_KEY": access_key["SecretAccessKey"],
	}


def ensure_role(iam, role_name: str, bucket: str, assume_policy: Dict[str, Any]) -> None:
	try:
		iam.get_role(RoleName=role_name)
		print(f"Reusing IAM role '{role_name}'.")
	except ClientError as exc:
		if exc.response["Error"]["Code"] == "NoSuchEntity":
			iam.create_role(RoleName=role_name, AssumeRolePolicyDocument=json.dumps(assume_policy))
			print(f"Created IAM role '{role_name}'.")
		else:
			raise
	policy_doc = json.dumps(_build_policy(bucket))
	iam.put_role_policy(RoleName=role_name, PolicyName=S3_POLICY_NAME, PolicyDocument=policy_doc)
	print(f"Attached inline S3 read policy '{S3_POLICY_NAME}' to role '{role_name}'.")


def parse_args() -> argparse.Namespace:
	parser = argparse.ArgumentParser(description="Provision S3 read-only IAM credentials for Streamlit.")
	parser.add_argument("--entity-name", default=DEFAULT_ENTITY_NAME, help="IAM user/role name to create or reuse.")
	parser.add_argument("--bucket", default=DEFAULT_BUCKET, help="S3 bucket name (default: chaptive-rag).")
	parser.add_argument(
		"--create-role",
		action="store_true",
		help="Also create an IAM role with the same name and inline policy (optional).",
	)
	parser.add_argument(
		"--role-trust",
		default=None,
		help="Path to a JSON file describing the assume-role trust policy (required if --create-role).",
	)
	return parser.parse_args()


def main() -> None:
	args = parse_args()
	iam = boto3.client("iam")
	ensure_user(iam, args.entity_name)
	attach_inline_policy(iam, args.entity_name, args.bucket)
	credentials = create_access_key(iam, args.entity_name)
	print("\nAdd the following values to your Streamlit secrets:")
	for key, value in credentials.items():
		print(f"{key} = {value}")
	print("AWS_REGION = ap-southeast-1")
	print(f"CHAPTIVE_S3_BUCKET = {args.bucket}")

	if args.create_role:
		if not args.role_trust:
			raise SystemExit("--role-trust JSON file is required when --create-role is set.")
		with open(args.role_trust, "r", encoding="utf-8") as handle:
			assume_policy = json.load(handle)
		ensure_role(iam, args.entity_name, args.bucket, assume_policy)
		print(
			"\nIAM role created. Attach this role to an EC2/Lambda/Streamlit Cloud runner and expose temporary credentials via environment variables."
		)


if __name__ == "__main__":
	try:
		main()
	except ClientError as exc:
		print(f"AWS error: {exc}")
		sys.exit(1)
	except KeyboardInterrupt:
		print("Aborted by user.")
		sys.exit(130)