package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"strings"

	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ecr"
	"gopkg.in/yaml.v2"
)

var (
	config = flag.String("config", "", "Required. Path to the file containing the list of repositories to grant permissions")
	prefix = flag.String("prefix", "", "Required. Prefix of the repository")
)

type grant struct {
	AccountArns      []string `yaml:"accountArns"`
	RepositoryPrefix string   `yaml:"repositoryPrefix"`
	Repositories     []string `yaml:"repositories"`
}

const policyFormat = `{
	"Statement": [
	  {
		"Action": [
		  "ecr:BatchCheckLayerAvailability",
		  "ecr:BatchGetImage",
		  "ecr:GetDownloadUrlForLayer"
		],
		"Principal": {
		  "AWS": [
			%s
		  ]
		},
		"Effect": "Allow",
		"Sid": "XAccount"
	  }
	],
	"Version": "2008-10-17"
  }`

func main() {
	flag.Parse()
	if *config == "" {
		flag.Usage()
		return
	}

	data, err := os.ReadFile(*config)
	if err != nil {
		log.Fatalf("Error reading file %s: %v", *config, err)
	}
	var g grant
	if err := yaml.Unmarshal(data, &g); err != nil {
		log.Fatalf("Error parsing file %s: %v", *config, err)
	}

	sess := session.Must(session.NewSession())
	e := ecr.New(sess)

	var arns []string
	for _, arn := range g.AccountArns {
		arns = append(arns, fmt.Sprintf("\"%s\"", arn))
	}
	accounts := strings.Join(arns, ",\n\t\t\t")
	for _, repo := range g.Repositories {
		p := fmt.Sprintf(policyFormat, accounts)
		fmt.Println(p)
		repository := fmt.Sprintf("%s/%s", *prefix, repo)

		_, err := e.SetRepositoryPolicy(&ecr.SetRepositoryPolicyInput{
			PolicyText:     &p,
			RepositoryName: &repository,
		})
		if err != nil {
			log.Fatalf("Error setting policy for %s: %v", repo, err)
		}
		fmt.Println("Permissions granted for repository", repository)
	}
}