From 40fb583160dedfb05a00bd4f9b266216d40435b8 Mon Sep 17 00:00:00 2001
From: iwilltry42 <iwilltry42@gmail.com>
Date: Tue, 7 May 2019 09:09:05 +0200
Subject: [PATCH] better image and version flags

---
 cli/commands.go  | 24 ++++++++++++++++++++++--
 cli/container.go |  2 --
 main.go          | 12 ++++++++----
 3 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/cli/commands.go b/cli/commands.go
index 1993584a..c342a601 100644
--- a/cli/commands.go
+++ b/cli/commands.go
@@ -23,6 +23,8 @@ import (
 	"github.com/urfave/cli"
 )
 
+const defaultRegistry = "docker.io"
+
 // CheckTools checks if the docker API server is responding
 func CheckTools(c *cli.Context) error {
 	log.Print("Checking docker...")
@@ -43,6 +45,24 @@ func CheckTools(c *cli.Context) error {
 // CreateCluster creates a new single-node cluster container and initializes the cluster directory
 func CreateCluster(c *cli.Context) error {
 
+	// define image
+	image := c.String("image")
+	if c.IsSet("version") {
+		// TODO: --version to be deprecated
+		log.Println("[WARNING] The `--version` flag will be deprecated soon, please use `--image rancher/k3s:<version>` instead")
+		if c.IsSet("image") {
+			// version specified, custom image = error (to push deprecation of version flag)
+			log.Fatalln("[ERROR] Please use `--image <image>:<version>` instead of --image and --version")
+		} else {
+			// version specified, default image = ok (until deprecation of version flag)
+			image = fmt.Sprintf("%s:%s", strings.Split(image, ":")[0], c.String("version"))
+		}
+	}
+	if len(strings.Split(image, "/")) <= 2 {
+		// fallback to default registry
+		image = fmt.Sprintf("%s/%s", defaultRegistry, image)
+	}
+
 	// create cluster network
 	networkID, err := createClusterNetwork(c.String("name"))
 	if err != nil {
@@ -75,7 +95,7 @@ func CreateCluster(c *cli.Context) error {
 	log.Printf("Creating cluster [%s]", c.String("name"))
 	dockerID, err := createServer(
 		c.GlobalBool("verbose"),
-		fmt.Sprintf("%s:%s", c.String("image"), c.String("version")),
+		image,
 		c.String("port"),
 		k3sServerArgs,
 		env,
@@ -137,7 +157,7 @@ func CreateCluster(c *cli.Context) error {
 		for i := 0; i < c.Int("workers"); i++ {
 			workerID, err := createWorker(
 				c.GlobalBool("verbose"),
-				fmt.Sprintf("%s:%s", c.String("image"), c.String("version")),
+				image,
 				k3sWorkerArgs,
 				env,
 				c.String("name"),
diff --git a/cli/container.go b/cli/container.go
index a1714816..962c1346 100644
--- a/cli/container.go
+++ b/cli/container.go
@@ -165,8 +165,6 @@ func createWorker(verbose bool, image string, args []string, env []string, name
 
 // removeContainer tries to rm a container, selected by Docker ID, and does a rm -f if it fails (e.g. if container is still running)
 func removeContainer(ID string) error {
-	// TODO: first check if container is running, then try to stop it with a timeout before trying to remove it
-	// if it does not terminate gracefully, try a force remove
 	ctx := context.Background()
 	docker, err := client.NewEnvClient()
 	if err != nil {
diff --git a/main.go b/main.go
index 5621d814..a9ef14fe 100644
--- a/main.go
+++ b/main.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"fmt"
 	"log"
 	"os"
 
@@ -9,6 +10,9 @@ import (
 	"github.com/urfave/cli"
 )
 
+// defaultK3sImage specifies the default image being used for server and workers
+const defaultK3sImage = "docker.io/rancher/k3s"
+
 // main represents the CLI application
 func main() {
 
@@ -56,8 +60,8 @@ func main() {
 					Usage: "Mount one or more volumes into every node of the cluster (Docker notation: `source:destination[,source:destination]`)",
 				},
 				cli.StringFlag{
-					Name:  "version, tag",
-					Value: version.GetK3sVersion(),
+					// TODO: to be deprecated
+					Name:  "version",
 					Usage: "Choose the k3s image version",
 				},
 				cli.IntFlag{
@@ -76,8 +80,8 @@ func main() {
 				},
 				cli.StringFlag{
 					Name:  "image, i",
-					Usage: "Specify a k3s image (repo only)",
-					Value: "docker.io/rancher/k3s",
+					Usage: "Specify a k3s image (Format: <repo>/<image>:<tag>)",
+					Value: fmt.Sprintf("%s:%s", defaultK3sImage, version.GetK3sVersion()),
 				},
 				cli.StringSliceFlag{
 					Name:  "server-arg, x",