From b208bc4c12d37cc772db023256bc461d917e9d65 Mon Sep 17 00:00:00 2001
From: knqyf263 <knqyf263@gmail.com>
Date: Wed, 8 May 2019 19:10:08 +0900
Subject: [PATCH] Add --refresh option

---
 cmd/trivy/main.go   |  4 ++++
 pkg/db/db.go        | 34 +++++++++++++++++++++++++++++++---
 pkg/git/git.go      |  6 ++++++
 pkg/run.go          | 27 ++++++++++++++++++++++++---
 pkg/scanner/scan.go | 10 ++++++++++
 5 files changed, 75 insertions(+), 6 deletions(-)

diff --git a/cmd/trivy/main.go b/cmd/trivy/main.go
index 4006fd77b1..8162fd8902 100644
--- a/cmd/trivy/main.go
+++ b/cmd/trivy/main.go
@@ -81,6 +81,10 @@ OPTIONS:
 			Name:  "ignore-unfixed",
 			Usage: "display only fixed vulnerabilities",
 		},
+		cli.BoolFlag{
+			Name:  "refresh",
+			Usage: "refresh DB (usually used after version update of trivy)",
+		},
 		cli.BoolFlag{
 			Name:  "debug, d",
 			Usage: "debug mode",
diff --git a/pkg/db/db.go b/pkg/db/db.go
index c443183906..9b32fceb8b 100644
--- a/pkg/db/db.go
+++ b/pkg/db/db.go
@@ -2,10 +2,11 @@ package db
 
 import (
 	"encoding/json"
-	"github.com/knqyf263/trivy/pkg/log"
 	"os"
 	"path/filepath"
 
+	"github.com/knqyf263/trivy/pkg/log"
+
 	"golang.org/x/xerrors"
 
 	"github.com/knqyf263/trivy/pkg/utils"
@@ -14,11 +15,11 @@ import (
 )
 
 var (
-	db *bolt.DB
+	db    *bolt.DB
+	dbDir = filepath.Join(utils.CacheDir(), "db")
 )
 
 func Init() (err error) {
-	dbDir := filepath.Join(utils.CacheDir(), "db")
 	if err = os.MkdirAll(dbDir, 0700); err != nil {
 		return xerrors.Errorf("failed to mkdir: %w", err)
 	}
@@ -32,6 +33,33 @@ func Init() (err error) {
 	return nil
 }
 
+func Reset() error {
+	if err := os.RemoveAll(dbDir); err != nil {
+		return xerrors.Errorf("failed to reset DB: %w", err)
+	}
+	return nil
+}
+
+func GetVersion() string {
+	var version string
+	value, err := Get("trivy", "metadata", "version")
+	if err != nil {
+		return ""
+	}
+	if err = json.Unmarshal(value, &version); err != nil {
+		return ""
+	}
+	return version
+}
+
+func SetVersion(version string) error {
+	err := Update("trivy", "metadata", "version", version)
+	if err != nil {
+		return xerrors.Errorf("failed to save DB version: %w", err)
+	}
+	return nil
+}
+
 func Update(rootBucket, nestedBucket, key string, value interface{}) error {
 	err := db.Update(func(tx *bolt.Tx) error {
 		return PutNestedBucket(tx, rootBucket, nestedBucket, key, value)
diff --git a/pkg/git/git.go b/pkg/git/git.go
index d576175359..d2ade7a950 100644
--- a/pkg/git/git.go
+++ b/pkg/git/git.go
@@ -5,6 +5,8 @@ import (
 	"path/filepath"
 	"strings"
 
+	"github.com/knqyf263/trivy/pkg/db"
+
 	"github.com/knqyf263/trivy/pkg/log"
 	"github.com/knqyf263/trivy/pkg/utils"
 	"golang.org/x/xerrors"
@@ -52,6 +54,10 @@ func CloneOrPull(url, repoPath string) (map[string]struct{}, error) {
 			return nil, xerrors.Errorf("failed to clone repository: %w", err)
 		}
 
+	}
+
+	// Need to refresh all vulnerabilities
+	if db.GetVersion() == "" {
 		err = filepath.Walk(repoPath, func(path string, info os.FileInfo, err error) error {
 			if info.IsDir() {
 				return nil
diff --git a/pkg/run.go b/pkg/run.go
index 219835717f..87d83326a8 100644
--- a/pkg/run.go
+++ b/pkg/run.go
@@ -23,6 +23,8 @@ import (
 )
 
 func Run(c *cli.Context) (err error) {
+	cliVersion := c.App.Version
+
 	debug := c.Bool("debug")
 	if err = log.InitLogger(debug); err != nil {
 		l.Fatal(err)
@@ -68,10 +70,22 @@ func Run(c *cli.Context) (err error) {
 		severities = append(severities, severity)
 	}
 
+	if c.Bool("refresh") {
+		log.Logger.Info("Resetting DB...")
+		if err = db.Reset(); err != nil {
+			return xerrors.Errorf("error in refresh DB: %w", err)
+		}
+	}
+
 	if err = db.Init(); err != nil {
 		return xerrors.Errorf("error in vulnerability DB initialize: %w", err)
 	}
 
+	dbVersion := db.GetVersion()
+	if dbVersion != "" && dbVersion != cliVersion {
+		log.Logger.Fatal("Detected version update of trivy. Please try again with --refresh option")
+	}
+
 	if !c.Bool("skip-update") {
 		if err = vulnsrc.Update(); err != nil {
 			return xerrors.Errorf("error in vulnerability DB update: %w", err)
@@ -103,9 +117,16 @@ func Run(c *cli.Context) (err error) {
 		return xerrors.Errorf("failed to write results: %w", err)
 	}
 
-	for _, result := range results {
-		if len(result.Vulnerabilities) > 0 {
-			os.Exit(c.Int("exit-code"))
+	if err = db.SetVersion(cliVersion); err != nil {
+		return xerrors.Errorf("unexpected error: %w", err)
+	}
+
+	exitCode := c.Int("exit-code")
+	if exitCode != 0 {
+		for _, result := range results {
+			if len(result.Vulnerabilities) > 0 {
+				os.Exit(exitCode)
+			}
 		}
 	}
 
diff --git a/pkg/scanner/scan.go b/pkg/scanner/scan.go
index 19415c47d6..8d186d5337 100644
--- a/pkg/scanner/scan.go
+++ b/pkg/scanner/scan.go
@@ -7,6 +7,8 @@ import (
 	"os"
 	"sort"
 
+	"github.com/genuinetools/reg/registry"
+
 	"github.com/knqyf263/trivy/pkg/log"
 
 	"github.com/knqyf263/trivy/pkg/report"
@@ -37,6 +39,14 @@ func ScanImage(imageName, filePath string, severities []vulnerability.Severity,
 	var err error
 	ctx := context.Background()
 
+	image, err := registry.ParseImage(imageName)
+	if err != nil {
+		return nil, xerrors.Errorf("invalid image: %w", err)
+	}
+	if image.Tag == "latest" {
+		log.Logger.Warn("You should avoid using the :latest tag as it is cached. You need to specify '--clean' option when :latest image is changed")
+	}
+
 	var target string
 	var files extractor.FileMap
 	if imageName != "" {