refactor init command for robust policy download checks

tenable · Feb 3, 2021 · 24bdc7a · 24bdc7a
1 parent 40a9ac6
commit 24bdc7a
Show file tree

Hide file tree

Showing 5 changed files with 17 additions and 21 deletions.
diff --git a/pkg/cli/init.go b/pkg/cli/init.go
@@ -29,14 +29,16 @@ var initCmd = &cobra.Command{
 
 Initializes Terrascan and clones policies from the Terrascan GitHub repository.
 `,
-	RunE:          initialInit,
+	RunE: func(cmd *cobra.Command, args []string) error {
+		return initial(cmd, args, false)
+	},
 	SilenceUsage:  true,
 	SilenceErrors: true,
 }
 
-func initialInit(cmd *cobra.Command, args []string) error {
+func initial(cmd *cobra.Command, args []string, isScanCmd bool) error {
 	// initialize terrascan
-	if err := initialize.Run(false); err != nil {
+	if err := initialize.Run(isScanCmd); err != nil {
 		zap.S().Error("failed to initialize terrascan")
 		return err
 	}

diff --git a/pkg/cli/run_test.go b/pkg/cli/run_test.go
@@ -38,7 +38,7 @@ func TestMain(m *testing.M) {
 func setup() {
 	// to download the policies for Run test
 	// downloads the policies at $HOME/.terrascan
-	initialScan(nil, nil)
+	initial(nil, nil, false)
 }
 
 func shutdown() {

diff --git a/pkg/cli/scan.go b/pkg/cli/scan.go
@@ -21,7 +21,6 @@ import (
 	"strings"
 
 	iacProvider "github.com/accurics/terrascan/pkg/iac-providers"
-	"github.com/accurics/terrascan/pkg/initialize"
 	"github.com/accurics/terrascan/pkg/policy"
 	"github.com/spf13/cobra"
 	"go.uber.org/zap"
@@ -36,21 +35,14 @@ var scanCmd = &cobra.Command{
 
 Detect compliance and security violations across Infrastructure as Code to mitigate risk before provisioning cloud native infrastructure.
 `,
-	PreRunE:       initialScan,
+	PreRunE: func(cmd *cobra.Command, args []string) error {
+		return initial(cmd, args, true)
+	},
 	RunE:          scan,
 	SilenceUsage:  true,
 	SilenceErrors: true,
 }
 
-func initialScan(cmd *cobra.Command, args []string) error {
-	// initialize terrascan
-	if err := initialize.Run(true); err != nil {
-		zap.S().Error("failed to initialize terrascan")
-		return err
-	}
-	return nil
-}
-
 func scan(cmd *cobra.Command, args []string) error {
 	zap.S().Debug("running terrascan in cli mode")
 	scanOptions.configFile = ConfigFile

diff --git a/pkg/cli/server.go b/pkg/cli/server.go
@@ -29,7 +29,7 @@ var serverCmd = &cobra.Command{
 Run Terrascan as an API server that inspects incoming IaC (Infrastructure-as-Code) files and returns the scan results.
 `,
 	PreRun: func(cmd *cobra.Command, args []string) {
-		initialScan(cmd, args)
+		initial(cmd, args, false)
 	},
 	Run: server,
 }

diff --git a/pkg/initialize/run.go b/pkg/initialize/run.go
@@ -19,6 +19,7 @@ package initialize
 import (
 	"fmt"
 	"os"
+	"path/filepath"
 
 	"github.com/accurics/terrascan/pkg/config"
 	"go.uber.org/zap"
@@ -44,9 +45,6 @@ func Run(isScanCmd bool) error {
 		if isScanCmd {
 			return nil
 		}
-		if err := os.RemoveAll(basePath); err != nil {
-			return err
-		}
 	}
 
 	// download policies
@@ -61,8 +59,10 @@ func Run(isScanCmd bool) error {
 // DownloadPolicies clones the policies to a local folder
 func DownloadPolicies() error {
 
+	tempPath := filepath.Join(os.TempDir(), "terrascan")
+
 	// clone the repo
-	r, err := git.PlainClone(basePath, false, &git.CloneOptions{
+	r, err := git.PlainClone(tempPath, false, &git.CloneOptions{
 		URL: repoURL,
 	})
 	if err != nil {
@@ -96,5 +96,7 @@ func DownloadPolicies() error {
 		return err
 	}
 
-	return nil
+	os.RemoveAll(basePath)
+
+	return os.Rename(tempPath, basePath)
 }