File size: 4,591 Bytes
d55322d 497bf08 d55322d 497bf08 d55322d 497bf08 d55322d 497bf08 d55322d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 |
"""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)
|