plugin-drone-kaniko/cmd/kaniko-ecr/main.go

package main

import (
	"context"
	"fmt"
	"io/ioutil"
	"os"
	"strings"

	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/service/ecr"
	"github.com/aws/aws-sdk-go-v2/service/ecrpublic"
	"github.com/aws/smithy-go"
	"github.com/joho/godotenv"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"github.com/urfave/cli"

	kaniko "github.com/drone/drone-kaniko"
	"github.com/drone/drone-kaniko/cmd/artifact"
)

const (
	accessKeyEnv     string = "AWS_ACCESS_KEY_ID"
	secretKeyEnv     string = "AWS_SECRET_ACCESS_KEY"
	dockerConfigPath string = "/kaniko/.docker/config.json"
	ecrPublicDomain  string = "public.ecr.aws"

	defaultDigestFile string = "/kaniko/digest-file"
)

var (
	version = "unknown"
)

func main() {
	// Load env-file if it exists first
	if env := os.Getenv("PLUGIN_ENV_FILE"); env != "" {
		if err := godotenv.Load(env); err != nil {
			logrus.Fatal(err)
		}
	}

	app := cli.NewApp()
	app.Name = "kaniko docker plugin"
	app.Usage = "kaniko docker plugin"
	app.Action = run
	app.Version = version
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:   "dockerfile",
			Usage:  "build dockerfile",
			Value:  "Dockerfile",
			EnvVar: "PLUGIN_DOCKERFILE",
		},
		cli.StringFlag{
			Name:   "context",
			Usage:  "build context",
			Value:  ".",
			EnvVar: "PLUGIN_CONTEXT",
		},
		cli.StringSliceFlag{
			Name:     "tags",
			Usage:    "build tags",
			Value:    &cli.StringSlice{"latest"},
			EnvVar:   "PLUGIN_TAGS",
			FilePath: ".tags",
		},
		cli.StringSliceFlag{
			Name:   "args",
			Usage:  "build args",
			EnvVar: "PLUGIN_BUILD_ARGS",
		},
		cli.StringFlag{
			Name:   "target",
			Usage:  "build target",
			EnvVar: "PLUGIN_TARGET",
		},
		cli.StringFlag{
			Name:   "repo",
			Usage:  "docker repository",
			EnvVar: "PLUGIN_REPO",
		},
		cli.BoolFlag{
			Name:   "create-repository",
			Usage:  "create ECR repository",
			EnvVar: "PLUGIN_CREATE_REPOSITORY",
		},
		cli.StringFlag{
			Name:   "region",
			Usage:  "AWS region",
			Value:  "us-east-1",
			EnvVar: "PLUGIN_REGION",
		},
		cli.StringSliceFlag{
			Name:   "custom-labels",
			Usage:  "additional k=v labels",
			EnvVar: "PLUGIN_CUSTOM_LABELS",
		},
		cli.StringFlag{
			Name:   "registry",
			Usage:  "ECR registry",
			EnvVar: "PLUGIN_REGISTRY",
		},
		cli.StringFlag{
			Name:   "access-key",
			Usage:  "ECR access key",
			EnvVar: "PLUGIN_ACCESS_KEY",
		},
		cli.StringFlag{
			Name:   "secret-key",
			Usage:  "ECR secret key",
			EnvVar: "PLUGIN_SECRET_KEY",
		},
		cli.StringFlag{
			Name:   "snapshot-mode",
			Usage:  "Specify one of full, redo or time as snapshot mode",
			EnvVar: "PLUGIN_SNAPSHOT_MODE",
		},
		cli.BoolFlag{
			Name:   "enable-cache",
			Usage:  "Set this flag to opt into caching with kaniko",
			EnvVar: "PLUGIN_ENABLE_CACHE",
		},
		cli.StringFlag{
			Name:   "cache-repo",
			Usage:  "Remote repository that will be used to store cached layers. Cache repo should be present in specified registry. enable-cache needs to be set to use this flag",
			EnvVar: "PLUGIN_CACHE_REPO",
		},
		cli.IntFlag{
			Name:   "cache-ttl",
			Usage:  "Cache timeout in hours. Defaults to two weeks.",
			EnvVar: "PLUGIN_CACHE_TTL",
		},
		cli.StringFlag{
			Name:   "artifact-file",
			Usage:  "Artifact file location that will be generated by the plugin. This file will include information of docker images that are uploaded by the plugin.",
			EnvVar: "PLUGIN_ARTIFACT_FILE",
		},
		cli.BoolFlag{
			Name:   "no-push",
			Usage:  "Set this flag if you only want to build the image, without pushing to a registry",
			EnvVar: "PLUGIN_NO_PUSH",
		},
		cli.StringFlag{
			Name:   "verbosity",
			Usage:  "Set this flag with value as oneof <panic|fatal|error|warn|info|debug|trace> to set the logging level for kaniko. Defaults to info.",
			EnvVar: "PLUGIN_VERBOSITY",
		},
	}

	if err := app.Run(os.Args); err != nil {
		logrus.Fatal(err)
	}
}

func run(c *cli.Context) error {
	repo := c.String("repo")
	registry := c.String("registry")

	if err := checkEmptyStringFlags(repo, registry); err != nil {
		return err
	}

	if err := setupECRAuth(c.String("access-key"), c.String("secret-key"), registry); err != nil {
		return err
	}

	if c.Bool("create-repository") {
		if err := createRepository(c.String("region"), repo, registry); err != nil {
			return err
		}
	}

	plugin := kaniko.Plugin{
		Build: kaniko.Build{
			Dockerfile:   c.String("dockerfile"),
			Context:      c.String("context"),
			Tags:         c.StringSlice("tags"),
			Args:         c.StringSlice("args"),
			Target:       c.String("target"),
			Repo:         fmt.Sprintf("%s/%s", c.String("registry"), c.String("repo")),
			Labels:       c.StringSlice("custom-labels"),
			SnapshotMode: c.String("snapshot-mode"),
			EnableCache:  c.Bool("enable-cache"),
			CacheRepo:    fmt.Sprintf("%s/%s", c.String("registry"), c.String("cache-repo")),
			CacheTTL:     c.Int("cache-ttl"),
			DigestFile:   defaultDigestFile,
			NoPush:       c.Bool("no-push"),
			Verbosity:    c.String("verbosity"),
		},
		Artifact: kaniko.Artifact{
			Tags:         c.StringSlice("tags"),
			Repo:         c.String("repo"),
			Registry:     c.String("registry"),
			ArtifactFile: c.String("artifact-file"),
			RegistryType: artifact.ECR,
		},
	}
	return plugin.Exec()
}

func checkEmptyStringFlags(flags ...string) error {
	for _, flag := range flags {
		if flag == "" {
			return fmt.Errorf("%s must be specified", flag)
		}
	}

	return nil
}

func setupECRAuth(accessKey, secretKey, registry string) error {
	// If IAM role is used, access key & secret key are not required
	if accessKey != "" && secretKey != "" {
		err := os.Setenv(accessKeyEnv, accessKey)
		if err != nil {
			return errors.Wrap(err, fmt.Sprintf("failed to set %s environment variable", accessKeyEnv))
		}

		err = os.Setenv(secretKeyEnv, secretKey)
		if err != nil {
			return errors.Wrap(err, fmt.Sprintf("failed to set %s environment variable", secretKeyEnv))
		}
	}

	jsonBytes := []byte(fmt.Sprintf(`{"credStore": "ecr-login", "credHelpers": {"%s": "ecr-login", "%s": "ecr-login"}}`, ecrPublicDomain, registry))
	err := ioutil.WriteFile(dockerConfigPath, jsonBytes, 0644)
	if err != nil {
		return errors.Wrap(err, "failed to create docker config file")
	}
	return nil
}

func createRepository(region, repo, registry string) error {
	cfg, err := config.LoadDefaultConfig(context.TODO(), config.WithRegion(region))
	if err != nil {
		return errors.Wrap(err, "failed to load aws config")
	}

	var createErr error
	switch {
	// create public repo if registry string starts with public domain (ex: public.ecr.aws/example-registry)
	case strings.HasPrefix(registry, ecrPublicDomain):
		svc := ecrpublic.NewFromConfig(cfg)
		_, createErr = svc.CreateRepository(context.TODO(), &ecrpublic.CreateRepositoryInput{RepositoryName: &repo})
	// create private repo
	default:
		svc := ecr.NewFromConfig(cfg)
		_, createErr = svc.CreateRepository(context.TODO(), &ecr.CreateRepositoryInput{RepositoryName: &repo})
	}

	var apiError smithy.APIError
	if errors.As(createErr, &apiError) && apiError.ErrorCode() != "RepositoryAlreadyExistsException" {
		return errors.Wrap(createErr, "failed to create repository")
	}

	return nil
}