unasuke · unasuke · Mar 9, 2017 · Mar 8, 2017 · Mar 8, 2017 · Mar 9, 2017
diff --git a/lambda.go b/lambda.go
@@ -0,0 +1,96 @@
+package main
+
+import (
+	"encoding/json"
+	"strings"
+
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/service/lambda"
+)
+
+func addPermissionToLambdaFromCloudWatchEvents(lc *lambda.Lambda, rules []Rule) error {
+	for _, rule := range rules {
+		for _, target := range rule.Targets {
+			if !IsLambdaFunction(target.Arn) {
+				continue
+			}
+
+			if result, err := isAlreadyAddPermission(lc, rule, target); err != nil {
+				return err
+			} else if result {
+				// do nothing (already granted permission)
+				continue
+			} else {
+				_, errL := lc.AddPermission(&lambda.AddPermissionInput{
+					Action:       aws.String("lambda:InvokeFunction"),
+					FunctionName: aws.String(LambdaFunctionNameFromArn(target.Arn)),
+					Principal:    aws.String("events.amazonaws.com"),
+					SourceArn:    rule.ActualRule.Arn,
+					StatementId:  aws.String(target.Id),
+				})
+
+				if errL != nil {
+					return errL
+				}
+			}
+		}
+	}
+	return nil
+}
+
+func removePermissonFromLambda(lc *lambda.Lambda, rules []Rule) error {
+	for _, rule := range rules {
+		for _, target := range rule.Targets {
+			if target.NeedDelete && IsLambdaFunction(*target.ActualTarget.Arn) {
+				_, err := lc.RemovePermission(&lambda.RemovePermissionInput{
+					FunctionName: target.ActualTarget.Arn,
+					StatementId:  target.ActualTarget.Id,
+				})
+
+				if err != nil {
+					return err
+				}
+			}
+		}
+	}
+
+	return nil
+}
+
+func isAlreadyAddPermission(lc *lambda.Lambda, rule Rule, target Target) (bool, error) {
+	var policy LambdaPolicy
+
+	policyOutput, err := lc.GetPolicy(&lambda.GetPolicyInput{
+		FunctionName: &target.Arn,
+	})
+
+	if err != nil {
+		return false, err
+	}
+
+	errJ := json.Unmarshal([]byte(*policyOutput.Policy), &policy)
+	if errJ != nil {
+		return false, errJ
+	}
+
+	for _, statement := range *policy.Statement {
+		if (statement.Resource == target.Arn &&
+			strings.HasSuffix(statement.Condition.ArnLike.AwsSourceArn, rule.Name) &&
+			statement.Effect == "Allow" &&
+			statement.Principal.Service == "events.amazonaws.com" &&
+			statement.Action == "lambda:InvokeFunction") ||
+			statement.StatementId == target.Id {
+			return true, nil
+		}
+	}
+
+	return false, nil
+}
+
+func IsLambdaFunction(arn string) bool {
+	return strings.HasPrefix(arn, "arn:aws:lambda")
+}
+
+func LambdaFunctionNameFromArn(arn string) string {
+	return strings.Split(arn, ":")[6]
+}
diff --git a/lambda_test.go b/lambda_test.go
@@ -0,0 +1,24 @@
+package main
+
+import (
+	"testing"
+)
+
+func TestIsLambdaFunction(t *testing.T) {
+	lambdaArn := "arn:aws:lambda:ap-northeast-1:000000000000:function:lambda-function-name:3"
+	if IsLambdaFunction(lambdaArn) == false {
+		t.Errorf("%s is lambda function", lambdaArn)
+	}
+
+	ec2Arn := "arn:aws:ec2:us-east-1:123456789012:instance/*"
+	if IsLambdaFunction(ec2Arn) == true {
+		t.Errorf("%s is not lambda function", ec2Arn)
+	}
+}
+
+func TestLambdaFunctionNameFromArn(t *testing.T) {
+	lambdaArn := "arn:aws:lambda:ap-northeast-1:000000000000:function:lambda-function-name:3"
+	if res := LambdaFunctionNameFromArn(lambdaArn); res != "lambda-function-name" {
+		t.Errorf("should return 'lambda-function-name', but got %s", res)
+	}
+}
diff --git a/main.go b/main.go
@@ -11,14 +11,18 @@ import (
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/cloudwatchevents"
+	"github.com/aws/aws-sdk-go/service/lambda"
 )
 
 var Version = "0.2.0"
 
 func main() {
 	var (
-		apply, dryrun   bool
-		file, awsRegion string
+		apply, dryrun                           bool
+		file, awsRegion                         string
+		err                                     error
+		sess                                    *session.Session
+		cweRulesBeforeApply, cweRulesAfterApply *cloudwatchevents.ListRulesOutput
 	)
 
 	flag.BoolVar(&apply, "apply", false, "apply to CloudWatch Events")
@@ -28,43 +32,65 @@ func main() {
 	flag.StringVar(&awsRegion, "region", os.Getenv("AWS_REGION"), "aws region")
 	flag.Parse()
 
-	sess, errS := session.NewSession(
+	sess, err = session.NewSession(
 		&aws.Config{
 			Region: aws.String(awsRegion),
 		},
 	)
-	if errS != nil {
-		fmt.Printf("Session error %v\n", errS)
+	if err != nil {
+		fmt.Printf("Session error %v\n", err)
 		os.Exit(1)
 	}
 
-	cweRulesOutput, errR := cloudwatchevents.New(sess).ListRules(nil)
-	if errR != nil {
-		fmt.Printf("API error %v\n", errR)
+	cweClient := cloudwatchevents.New(sess)
+	lambdaClient := lambda.New(sess)
+
+	cweRulesBeforeApply, err = cweClient.ListRules(nil)
+	if err != nil {
+		fmt.Printf("API error %v\n", err)
 		os.Exit(1)
 	}
 
 	describedRules := Rules{}
-	errY := loadYaml(file, &describedRules)
-	if errY != nil {
-		fmt.Printf("File error %v\n", errY)
+	err = loadYaml(file, &describedRules)
+	if err != nil {
+		fmt.Printf("File error %v\n", err)
 		os.Exit(1)
 	}
 
-	describedRules.Rules = AssociateRules(cweRulesOutput.Rules, describedRules.Rules)
+	describedRules.Rules = AssociateRules(cweRulesBeforeApply.Rules, describedRules.Rules)
 	for i, rule := range describedRules.Rules {
-		t, _ := fetchActualTargetsByRule(cloudwatchevents.New(sess), rule)
+		t, _ := fetchActualTargetsByRule(cweClient, rule)
 		describedRules.Rules[i].Targets = AssociateTargets(t, describedRules.Rules[i].Targets)
 	}
 	CheckIsNeedUpdateOrDelete(describedRules.Rules)
 	displayWhatWillChange(describedRules.Rules)
 
 	if apply && !dryrun {
-		errU := updateCloudWatchEvents(cloudwatchevents.New(sess), describedRules.Rules)
-		if errU != nil {
-			fmt.Printf("API error %v\n", errU)
+		err = updateCloudWatchEvents(cweClient, describedRules.Rules)
+		if err != nil {
+			fmt.Printf("API error %v\n", err)
 			os.Exit(1)
 		}
+
+		err = removePermissonFromLambda(lambdaClient, describedRules.Rules)
+		if err != nil {
+			fmt.Printf("API error %v\n", err)
+			os.Exit(1)
+		}
+
+		// Grant permission to invoke lambda function from CloudWatch Events
+		cweRulesAfterApply, err = cweClient.ListRules(nil)
+		describedRules.Rules = AssociateRules(cweRulesAfterApply.Rules, describedRules.Rules)
+		for i, rule := range describedRules.Rules {
+			t, _ := fetchActualTargetsByRule(cweClient, rule)
+			describedRules.Rules[i].Targets = AssociateTargets(t, describedRules.Rules[i].Targets)
+		}
+
+		err = addPermissionToLambdaFromCloudWatchEvents(lambdaClient, describedRules.Rules)
+		if err != nil {
+			fmt.Print("Grant permission error %v\n", err)
+		}
 	}
 }
 

diff --git a/structs.go b/structs.go
@@ -4,10 +4,12 @@ import (
 	cwe "github.com/aws/aws-sdk-go/service/cloudwatchevents"
 )
 
+// struct for store unmarshalized configuration yaml
 type Rules struct {
 	Rules []Rule
 }
 
+// struct for expression CloudWatch Events Rule
 type Rule struct {
 	Description        string   `yaml:"description"`
 	EventPattern       string   `yaml:"event_pattern"`
@@ -21,6 +23,7 @@ type Rule struct {
 	NeedDelete         bool
 }
 
+// struct for expression CloudWatch Events Target
 type Target struct {
 	Arn          string `yaml:"arn"`
 	Id           string `yaml:"id"`
@@ -30,3 +33,35 @@ type Target struct {
 	NeedUpdate   bool
 	NeedDelete   bool
 }
+
+// struct for JSON that return from Lambda.GetPolicy
+type LambdaPolicy struct {
+	Version   string             `json:"Version"`
+	Id        string             `json:"Id"`
+	Statement *[]PolicyStatement `json:"Statement"`
+}
+
+// part of the LambdaPolicy
+type PolicyStatement struct {
+	Resource    string           `json:"Resource"`
+	Condition   *PolicyCondition `json:"Condition"`
+	StatementId string           `json:"Sid"`
+	Effect      string           `json:"Effect"`
+	Principal   *PolicyPrincipal `json:"Principal"`
+	Action      string           `json:"Action"`
+}
+
+// part of the LambdaPolicy
+type PolicyCondition struct {
+	ArnLike *PolicyArnLike `json:"ArnLike"`
+}
+
+// part of the LambdaPolicy
+type PolicyArnLike struct {
+	AwsSourceArn string `json:"AWS:SourceArn"`
+}
+
+// part of the LambdaPolicy
+type PolicyPrincipal struct {
+	Service string `json:"Service"`
+}