From 24bdc7a3318257e90da1a13ed6644b74913f5940 Mon Sep 17 00:00:00 2001
From: Devang Gaur <devang.gaur@accurics.com>
Date: Thu, 4 Feb 2021 02:05:01 +0530
Subject: [PATCH] refactor init command for robust policy download checks

---
 pkg/cli/init.go       |  8 +++++---
 pkg/cli/run_test.go   |  2 +-
 pkg/cli/scan.go       | 14 +++-----------
 pkg/cli/server.go     |  2 +-
 pkg/initialize/run.go | 12 +++++++-----
 5 files changed, 17 insertions(+), 21 deletions(-)

diff --git a/pkg/cli/init.go b/pkg/cli/init.go
index b6dcf6534..d96bd27fe 100644
--- 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
index 8f231f703..d7b5949c4 100644
--- 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
index a4eac4eaf..1d755426b 100644
--- 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
index 6031b4225..edb46279c 100644
--- 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
index 998a2b14a..c49482aca 100644
--- 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)
 }