aws/s3_quota/setup.sh

   1 #!/bin/bash
   2
   3
   4 # Set your preferred profile name.
   5 #AWS_DEFAULT_PROFILE=default
   6 # Set S3 bucket names in which quota should be enabled.
   7 # Leaving empty will enable quota on all buckets.
   8 #Buckets="your_bucket_name1 your_bucket_name2 ..."
   9 Buckets=
  10 Rolename=lambda-s3quota
  11 Funcname=quota-s3
  12 lastlog=
  13
  14
  15 logstart()
  16 {
  17         echo "\e[33m★ $*\e[m" >&2
  18         lastlog="$*"
  19 }
  20 logend()
  21 {
  22         local c=2 m=✔
  23         if [ "$1" -gt 0 ] 2>/dev/null || [ -z "$1" ]
  24         then
  25                 c=1
  26                 m=✘
  27         fi
  28         [ "$1" -ge 0 ] 2>/dev/null && shift
  29         echo "\e[3$c;1m$m $lastlog\e[m \e[3${c}m$*\e[m" >&2
  30 }
  31
  32
  33 ## Create IAM Role
  34
  35 logstart Create Role $Rolename
  36 aws iam create-role --role-name "$Rolename" --assume-role-policy-document file://assumerole-lambda.json
  37 rolearn=`aws iam get-role --role-name "$Rolename" --query Role.Arn --output text`
  38 logend "$rolearn"
  39
  40 ## Create IAM Policies
  41
  42 logstart Create Policy IAMListUsers
  43 aws iam create-policy --policy-name IAMListUsers --policy-document file://policy-listusers.json
  44 logstart Create Policy S3PutObjectTagging
  45 aws iam create-policy --policy-name S3PutObjectTagging --policy-document file://policy-s3puttag.json
  46
  47
  48 ## Attach Policies to Role
  49
  50 get_policy_arn()
  51 {
  52         local policy=$1
  53         aws iam list-policies --query "Policies[?PolicyName==\`$policy\`].Arn" --output text
  54 }
  55
  56 for policy in IAMListUsers S3PutObjectTagging AWSLambdaBasicExecutionRole
  57 do
  58         logstart Attach Policy $policy to Role $Rolename
  59         arn=`get_policy_arn "$policy"`
  60         aws iam attach-role-policy --role-name "$Rolename" --policy-arn "$arn"
  61         logend "$arn"
  62 done
  63
  64 aws iam list-attached-role-policies --role-name "$Rolename" --output text
  65
  66
  67 ## Create Cloudwatch Log Group
  68 logstart Create Cloudwatch Log Group
  69 aws logs create-log-group --log-group-name /aws/lambda/"$Funcname"
  70 aws logs put-retention-policy --log-group-name /aws/lambda/"$Funcname" --retention-in-days 7
  71
  72
  73 ## Create Lambda function
  74
  75 logstart Create Lambda $Funcname
  76 tmpfile=`mktemp --suffix=.zip -u`
  77 trap 'rm "$tmpfile"' EXIT INT
  78 zip "$tmpfile" lambda_function.py
  79
  80 aws lambda create-function --function-name "$Funcname" \
  81         --description "Listens to S3 write events and put 'Creator' tag on objects." \
  82         --role "$rolearn" \
  83         --runtime python2.7 \
  84         --handler lambda_function.lambda_handler \
  85         --zip-file "fileb://$tmpfile"
  86
  87 lambdaarn=`aws lambda get-function --function-name "$Funcname" --query Configuration.FunctionArn --output text`
  88 logend "$lambdaarn"
  89
  90
  91 ## Grant S3 to invoke Lambda
  92
  93 logstart Permit S3 service to Lambda
  94 aws lambda add-permission --function-name "$Funcname" --statement-id quota-s3 --action lambda:InvokeFunction --principal s3.amazonaws.com
  95
  96
  97 ## Create Lambda Trigger within S3
  98
  99 merge_notification_config()
 100 {
 101         local bucket=$1
 102         export lambdaarn=$2
 103         export eventId="create-all-lambda"
 104         aws s3api get-bucket-notification-configuration --bucket "$bucket" |\
 105         python -c '
 106 from __future__ import print_function
 107 import json
 108 import sys
 109 import os
 110 eventId = os.environ["eventId"]
 111 cnfstr = "".join(sys.stdin.readlines())
 112 sys.stderr.write("Current Notification Config: %s\n" % (cnfstr))
 113 if cnfstr:
 114         config = json.loads(cnfstr)
 115 else:
 116         config = {"LambdaFunctionConfigurations":[]}
 117 Lambda = dict(map(lambda x: (x["Id"], x), config["LambdaFunctionConfigurations"]))
 118 newLambda = {eventId: {
 119  "Id": eventId,
 120  "LambdaFunctionArn": os.environ["lambdaarn"],
 121  "Events": ["s3:ObjectCreated:*"]
 122 }}
 123 Lambda.update(newLambda)
 124 config.update({"LambdaFunctionConfigurations": Lambda.values()})
 125 print(json.dumps(config))
 126 '
 127 }
 128
 129 for bucket in `[ -n "$Bucket" ] && echo "$Buckets" || aws s3api list-buckets --query 'Buckets[*].Name' --output text`
 130 do
 131         logstart Set Notification Config on bucket $bucket
 132         aws s3api put-bucket-notification-configuration --bucket "$bucket" --notification-configuration "$(merge_notification_config "$bucket" "$lambdaarn")"
 133         logend $?
 134 done
 135