add guard and remove source dir option

projectdiscovery · Jul 6, 2023 · 3395237 · 3395237
1 parent 70cb3f0
commit 3395237
Show file tree

Hide file tree

Showing 3 changed files with 80 additions and 25 deletions.
diff --git a/.gitignore b/.gitignore
@@ -2,3 +2,4 @@
 .vscode
 *.exe
 
+.devcontainer
diff --git a/folder/folderutil.go b/folder/folderutil.go
@@ -1,6 +1,7 @@
 package folderutil
 
 import (
+	"errors"
 	"os"
 	"os/user"
 	"path/filepath"
@@ -168,11 +169,15 @@ func AppConfigDirOrDefault(defaultAppConfigDir string, toolName string) string {
 }
 
 // MigrateDir moves all files and non-empty directories from sourceDir to destinationDir and removes sourceDir
-func MigrateDir(sourceDir string, destinationDir string) error {
+func MigrateDir(sourceDir string, destinationDir string, removeSourceDir bool) error {
 	// trim trailing slash to avoid slash related issues
 	sourceDir = strings.TrimSuffix(sourceDir, Separator)
 	destinationDir = strings.TrimSuffix(destinationDir, Separator)
 
+	if sourceDir == destinationDir {
+		return errors.New("sourceDir and destinationDir cannot be the same")
+	}
+
 	entries, err := os.ReadDir(sourceDir)
 	if err != nil {
 		return err
@@ -193,7 +198,7 @@ func MigrateDir(sourceDir string, destinationDir string) error {
 					return err
 				}
 
-				err = MigrateDir(sourcePath, destPath)
+				err = MigrateDir(sourcePath, destPath, removeSourceDir)
 				if err != nil {
 					return err
 				}
@@ -206,9 +211,11 @@ func MigrateDir(sourceDir string, destinationDir string) error {
 		}
 	}
 
-	err = os.RemoveAll(sourceDir)
-	if err != nil {
-		return err
+	if removeSourceDir {
+		err = os.RemoveAll(sourceDir)
+		if err != nil {
+			return err
+		}
 	}
 
 	return nil

diff --git a/folder/folderutil_test.go b/folder/folderutil_test.go
@@ -2,8 +2,10 @@ package folderutil
 
 import (
 	"os"
+	"path/filepath"
 	"testing"
 
+	fileutil "github.com/projectdiscovery/utils/file"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -19,47 +21,92 @@ func TestGetFiles(t *testing.T) {
 
 func TestMigrateDir(t *testing.T) {
 	t.Run("destination folder creation error", func(t *testing.T) {
-		err := MigrateDir("/source", "/:/dest")
+		err := MigrateDir("/source", "/:/dest", true)
 		assert.Error(t, err)
 	})
 
 	t.Run("source folder not found error", func(t *testing.T) {
-		err := MigrateDir("/notExistingFolder", "/dest")
+		err := MigrateDir("/notExistingFolder", "/dest", true)
 		assert.Error(t, err)
 	})
 
-	t.Run("successful migration", func(t *testing.T) {
+	t.Run("source and destination are the same", func(t *testing.T) {
 		// setup
 		// some files in a temp dir
 		sourceDir := t.TempDir()
 		defer os.RemoveAll(sourceDir)
-		_ = os.WriteFile(sourceDir+"/file1.txt", []byte("file1"), os.ModePerm)
-		_ = os.WriteFile(sourceDir+"/file2.txt", []byte("file2"), os.ModePerm)
-		_ = os.Mkdir(sourceDir+"/dir1", os.ModePerm)
-		_ = os.WriteFile(sourceDir+"/dir1"+"/file3.txt", []byte("file3"), os.ModePerm)
-		_ = os.Mkdir(sourceDir+"/dir2", os.ModePerm)
+		_ = os.WriteFile(filepath.Join(sourceDir, "/file1.txt"), []byte("file1"), os.ModePerm)
+		_ = os.WriteFile(filepath.Join(sourceDir, "/file2.txt"), []byte("file2"), os.ModePerm)
+
+		// when: try to migrate files
+		err := MigrateDir(sourceDir, sourceDir, true)
+
+		// then: verify if files migrated successfully
+		assert.Error(t, err)
+
+		assert.True(t, fileutil.FileExists(filepath.Join(sourceDir, "/file1.txt")))
+		assert.True(t, fileutil.FileExists(filepath.Join(sourceDir, "/file2.txt")))
+	})
+
+	t.Run("successful migration with source dir removal", func(t *testing.T) {
+		// setup
+		// some files in a temp dir
+		sourceDir := t.TempDir()
+		defer os.RemoveAll(sourceDir)
+		_ = os.WriteFile(filepath.Join(sourceDir, "/file1.txt"), []byte("file1"), os.ModePerm)
+		_ = os.WriteFile(filepath.Join(sourceDir, "/file2.txt"), []byte("file2"), os.ModePerm)
+		_ = os.Mkdir(filepath.Join(sourceDir, "/dir1"), os.ModePerm)
+		_ = os.WriteFile(filepath.Join(sourceDir, "/dir1", "/file3.txt"), []byte("file3"), os.ModePerm)
+		_ = os.Mkdir(filepath.Join(sourceDir, "/dir2"), os.ModePerm)
 
 		// destination directory
 		destinationDir := t.TempDir()
 		defer os.RemoveAll(destinationDir)
 
 		// when: try to migrate files
-		err := MigrateDir(sourceDir, destinationDir)
+		err := MigrateDir(sourceDir, destinationDir, true)
 
 		// then: verify if files migrated successfully
 		assert.NoError(t, err)
 
-		_, err = os.Stat(destinationDir + "/file1.txt")
-		assert.NoError(t, err)
-		_, err = os.Stat(destinationDir + "/file2.txt")
-		assert.NoError(t, err)
-		_, err = os.Stat(destinationDir + "/dir1")
-		assert.NoError(t, err)
-		_, err = os.Stat(destinationDir + "/dir1" + "/file3.txt")
+		assert.True(t, fileutil.FileExists(filepath.Join(destinationDir, "/file1.txt")))
+		assert.True(t, fileutil.FileExists(filepath.Join(destinationDir, "/file2.txt")))
+		assert.True(t, fileutil.FolderExists(filepath.Join(destinationDir, "/dir1")))
+		assert.True(t, fileutil.FileExists(filepath.Join(destinationDir, "/dir1", "/file3.txt")))
+
+		assert.False(t, fileutil.FolderExists(filepath.Join(destinationDir, "/dir2")))
+
+		assert.False(t, fileutil.FolderExists(sourceDir))
+	})
+
+	t.Run("successful migration without source dir removal", func(t *testing.T) {
+		// setup
+		// some files in a temp dir
+		sourceDir := t.TempDir()
+		defer os.RemoveAll(sourceDir)
+		_ = os.WriteFile(filepath.Join(sourceDir, "/file1.txt"), []byte("file1"), os.ModePerm)
+		_ = os.WriteFile(filepath.Join(sourceDir, "/file2.txt"), []byte("file2"), os.ModePerm)
+		_ = os.Mkdir(filepath.Join(sourceDir, "/dir1"), os.ModePerm)
+		_ = os.WriteFile(filepath.Join(sourceDir, "/dir1", "/file3.txt"), []byte("file3"), os.ModePerm)
+		_ = os.Mkdir(filepath.Join(sourceDir, "/dir2"), os.ModePerm)
+
+		// destination directory
+		destinationDir := t.TempDir()
+		defer os.RemoveAll(destinationDir)
+
+		// when: try to migrate files
+		err := MigrateDir(sourceDir, destinationDir, false)
+
+		// then: verify if files migrated successfully
 		assert.NoError(t, err)
-		_, err = os.Stat(destinationDir + "/dir2")
-		assert.Error(t, err)
-		_, err = os.Stat(sourceDir)
-		assert.Error(t, err)
+
+		assert.True(t, fileutil.FileExists(filepath.Join(destinationDir, "/file1.txt")))
+		assert.True(t, fileutil.FileExists(filepath.Join(destinationDir, "/file2.txt")))
+		assert.True(t, fileutil.FolderExists(filepath.Join(destinationDir, "/dir1")))
+		assert.True(t, fileutil.FileExists(filepath.Join(destinationDir, "/dir1", "/file3.txt")))
+
+		assert.False(t, fileutil.FolderExists(filepath.Join(destinationDir, "/dir2")))
+
+		assert.True(t, fileutil.FolderExists(sourceDir))
 	})
 }