From 711e19a4eb10cd63d25e98246db95d1927123afb Mon Sep 17 00:00:00 2001
From: Ivan Ilves <ivan.ilves@schibsted.com>
Date: Tue, 26 Sep 2017 08:25:13 +0200
Subject: [PATCH] Pull Docker images matching filters specified

---
 main.go            | 25 ++++++++++++++++++++++++-
 tag/local/local.go | 43 ++++++++++++++++++++++++++++++++++---------
 tag/tag.go         |  9 +++++++++
 3 files changed, 67 insertions(+), 10 deletions(-)

diff --git a/main.go b/main.go
index 71d3c73..6cfa54b 100644
--- a/main.go
+++ b/main.go
@@ -22,6 +22,7 @@ type options struct {
 	Password         string `short:"p" long:"password" default:"" description:"Docker registry password" env:"PASSWORD"`
 	DockerJSON       string `shord:"j" long:"docker-json" default:"~/.docker/config.json" env:"DOCKER_JSON"`
 	Concurrency      int    `short:"c" long:"concurrency" default:"32" description:"Concurrent request limit while querying registry" env:"CONCURRENCY"`
+	PullImages       bool   `short:"P" long:"pull-images" description:"Pull images matched by filter" env:"PULL_IMAGES"`
 	InsecureRegistry bool   `short:"i" long:"insecure-registry" description:"Use insecure plain-HTTP registriy" env:"INSECURE_REGISTRY"`
 	TraceRequests    bool   `short:"T" long:"trace-requests" description:"Trace registry HTTP requests" env:"TRACE_REQUESTS"`
 	Version          bool   `short:"V" long:"version" description:"Show version and exit"`
@@ -135,7 +136,7 @@ func main() {
 		os.Exit(0)
 	}
 	if o.Positional.Repository == "" {
-		suicide(errors.New("You should provide a repository name, e.g. 'nginx' or 'mesosphere/chronos'"))
+		suicide(errors.New("You should provide a repository name, e.g. 'nginx~/^1\\\\.13/' or 'mesosphere/chronos'"))
 	}
 
 	if o.InsecureRegistry {
@@ -198,4 +199,26 @@ func main() {
 			repoLocalName+":"+tg.GetName(),
 		)
 	}
+
+	if o.PullImages {
+		for _, key := range sortedKeys {
+			name := names[key]
+
+			tg := joinedTags[name]
+
+			if !matchesFilter(tg.GetName(), filter) {
+				continue
+			}
+
+			if tg.NeedsPull() {
+				ref := repoLocalName + ":" + tg.GetName()
+
+				fmt.Printf("Pulling: %s\n", ref)
+				err := local.PullImage(ref)
+				if err != nil {
+					suicide(err)
+				}
+			}
+		}
+	}
 }
diff --git a/tag/local/local.go b/tag/local/local.go
index 1f29d5a..ad85fd5 100644
--- a/tag/local/local.go
+++ b/tag/local/local.go
@@ -3,6 +3,7 @@ package local
 import (
 	"encoding/json"
 	"io"
+	"io/ioutil"
 	"net/http"
 	"strings"
 	"time"
@@ -57,6 +58,20 @@ func detectAPIVersion() (string, error) {
 	return parseAPIVersionJSON(resp.Body)
 }
 
+func newClient() (*client.Client, error) {
+	apiVersion, err := detectAPIVersion()
+	if err != nil {
+		return nil, err
+	}
+
+	cli, err := client.NewClient("unix://"+dockerSocket, apiVersion, nil, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	return cli, err
+}
+
 func newImageListOptions(repo string) (types.ImageListOptions, error) {
 	repoFilter := "reference=" + repo
 	filterArgs := filters.NewArgs()
@@ -95,14 +110,7 @@ func extractTagNames(repoTags []string, repo string) []string {
 
 // FetchTags looks up Docker repo tags and IDs present on local Docker daemon
 func FetchTags(repo string) (map[string]*tag.Tag, error) {
-	ctx := context.Background()
-
-	apiVersion, err := detectAPIVersion()
-	if err != nil {
-		return nil, err
-	}
-
-	cli, err := client.NewClient("unix://"+dockerSocket, apiVersion, nil, nil)
+	cli, err := newClient()
 	if err != nil {
 		return nil, err
 	}
@@ -111,7 +119,7 @@ func FetchTags(repo string) (map[string]*tag.Tag, error) {
 	if err != nil {
 		return nil, err
 	}
-	imageSummaries, err := cli.ImageList(ctx, listOptions)
+	imageSummaries, err := cli.ImageList(context.Background(), listOptions)
 	if err != nil {
 		return nil, err
 	}
@@ -159,3 +167,20 @@ func FormatRepoName(repository, registry string) string {
 
 	return registry + "/" + repository
 }
+
+// PullImage pulls Docker image specified locally
+func PullImage(ref string) error {
+	cli, err := newClient()
+	if err != nil {
+		return err
+	}
+
+	resp, err := cli.ImagePull(context.Background(), ref, types.ImagePullOptions{})
+	if err != nil {
+		return err
+	}
+
+	_, err = ioutil.ReadAll(resp)
+
+	return err
+}
diff --git a/tag/tag.go b/tag/tag.go
index 7f675e3..d9f4eee 100644
--- a/tag/tag.go
+++ b/tag/tag.go
@@ -77,6 +77,15 @@ func (tg *Tag) GetState() string {
 	return tg.state
 }
 
+// NeedsPull tells us if tag/image needs pull
+func (tg *Tag) NeedsPull() bool {
+	if tg.state == "ABSENT" || tg.state == "CHANGED" {
+		return true
+	}
+
+	return false
+}
+
 // SetCreated sets image creation timestamp
 func (tg *Tag) SetCreated(created int64) {
 	tg.created = created