Merge pull request #801 from schollz:schollz/issue799

Schollz/issue799
schollz · Sep 3, 2024 · f9241b7 · f9241b7
2 parents cdf3aa0 + bb74eaf
commit f9241b7
Show file tree

Hide file tree

Showing 4 changed files with 57 additions and 37 deletions.
diff --git a/main.go b/main.go
@@ -5,13 +5,13 @@ package main
 //go:generate git tag -af v$VERSION -m "v$VERSION"
 
 import (
-	"log"
 	"os"
 	"os/signal"
 	"syscall"
 
 	"github.com/schollz/croc/v10/src/cli"
 	"github.com/schollz/croc/v10/src/utils"
+	log "github.com/schollz/logger"
 )
 
 func main() {
@@ -38,17 +38,17 @@ func main() {
 
 	go func() {
 		if err := cli.Run(); err != nil {
-			log.Fatalln(err)
+			log.Error(err)
 		}
+		// Exit the program gracefully
+		utils.RemoveMarkedFiles()
+		os.Exit(0)
 	}()
 
 	// Wait for a termination signal
 	sig := <-sigs
-	log.Println("Received signal:", sig)
-
-	// Perform any necessary cleanup here
-	log.Println("Performing cleanup...")
-	utils.CleanupTempData()
+	log.Debugf("Received signal:", sig)
+	utils.RemoveMarkedFiles()
 
 	// Exit the program gracefully
 	os.Exit(0)

diff --git a/src/cli/cli.go b/src/cli/cli.go
@@ -358,6 +358,7 @@ func send(c *cli.Context) (err error) {
 		if err != nil {
 			return
 		}
+		utils.MarkFileForRemoval(fnames[0])
 		defer func() {
 			e := os.Remove(fnames[0])
 			if e != nil {
@@ -369,6 +370,7 @@ func send(c *cli.Context) (err error) {
 		if err != nil {
 			return
 		}
+		utils.MarkFileForRemoval(fnames[0])
 		defer func() {
 			e := os.Remove(fnames[0])
 			if e != nil {
@@ -446,15 +448,9 @@ func getStdin() (fnames []string, err error) {
 	fnames = []string{f.Name()}
 	return
 }
-func makeTempFolder() {
-	path := "temp"
-	if _, err := os.Stat(path); os.IsNotExist(err) {
-		os.Mkdir(path, os.ModePerm)
-	}
-}
+
 func makeTempFileWithString(s string) (fnames []string, err error) {
-	makeTempFolder()
-	f, err := os.CreateTemp("temp", "croc-stdin-")
+	f, err := os.CreateTemp(".", "croc-stdin-")
 	if err != nil {
 		return
 	}

diff --git a/src/croc/croc.go b/src/croc/croc.go
@@ -388,6 +388,7 @@ func GetFilesInfo(fnames []string, zipfolder bool, ignoreGit bool) (filesInfo []
 			fpath = filepath.Dir(fpath)
 			dest := filepath.Base(fpath) + ".zip"
 			utils.ZipDirectory(dest, fpath)
+			utils.MarkFileForRemoval(dest)
 			stat, errStat = os.Lstat(dest)
 			if errStat != nil {
 				err = errStat

diff --git a/src/utils/utils.go b/src/utils/utils.go
@@ -11,7 +11,6 @@ import (
 	"encoding/hex"
 	"fmt"
 	"io"
-	"log"
 	"math"
 	"math/big"
 	"net"
@@ -26,6 +25,7 @@ import (
 	"github.com/kalafut/imohash"
 	"github.com/minio/highwayhash"
 	"github.com/pion/stun"
+	log "github.com/schollz/logger"
 	"github.com/schollz/mnemonicode"
 	"github.com/schollz/progressbar/v3"
 )
@@ -276,7 +276,8 @@ func PublicIP() (ip string, err error) {
 func LocalIP() string {
 	conn, err := net.Dial("udp", "8.8.8.8:80")
 	if err != nil {
-		log.Fatal(err)
+		log.Error(err)
+		return ""
 	}
 	defer conn.Close()
 
@@ -477,12 +478,12 @@ func IsLocalIP(ipaddress string) bool {
 
 func ZipDirectory(destination string, source string) (err error) {
 	if _, err = os.Stat(destination); err == nil {
-		log.Fatalf("%s file already exists!\n", destination)
+		log.Errorf("%s file already exists!\n", destination)
 	}
 	fmt.Fprintf(os.Stderr, "Zipping %s to %s\n", source, destination)
 	file, err := os.Create(destination)
 	if err != nil {
-		log.Fatalln(err)
+		log.Error(err)
 	}
 	defer file.Close()
 	writer := zip.NewWriter(file)
@@ -493,30 +494,30 @@ func ZipDirectory(destination string, source string) (err error) {
 	defer writer.Close()
 	err = filepath.Walk(source, func(path string, info os.FileInfo, err error) error {
 		if err != nil {
-			log.Fatalln(err)
+			log.Error(err)
 		}
 		if info.Mode().IsRegular() {
 			f1, err := os.Open(path)
 			if err != nil {
-				log.Fatalln(err)
+				log.Error(err)
 			}
 			defer f1.Close()
 			zipPath := strings.ReplaceAll(path, source, strings.TrimSuffix(destination, ".zip"))
 			zipPath = filepath.ToSlash(zipPath)
 			w1, err := writer.Create(zipPath)
 			if err != nil {
-				log.Fatalln(err)
+				log.Error(err)
 			}
 			if _, err := io.Copy(w1, f1); err != nil {
-				log.Fatalln(err)
+				log.Error(err)
 			}
 			fmt.Fprintf(os.Stderr, "\r\033[2K")
 			fmt.Fprintf(os.Stderr, "\rAdding %s", zipPath)
 		}
 		return nil
 	})
 	if err != nil {
-		log.Fatalln(err)
+		log.Error(err)
 	}
 	fmt.Fprintf(os.Stderr, "\n")
 	return nil
@@ -525,7 +526,7 @@ func ZipDirectory(destination string, source string) (err error) {
 func UnzipDirectory(destination string, source string) error {
 	archive, err := zip.OpenReader(source)
 	if err != nil {
-		log.Fatalln(err)
+		log.Error(err)
 	}
 	defer archive.Close()
 
@@ -537,15 +538,15 @@ func UnzipDirectory(destination string, source string) error {
 		// make sure the filepath does not have ".."
 		filePath = filepath.Clean(filePath)
 		if strings.Contains(filePath, "..") {
-			log.Fatalf("Invalid file path %s\n", filePath)
+			log.Errorf("Invalid file path %s\n", filePath)
 		}
 		if f.FileInfo().IsDir() {
 			os.MkdirAll(filePath, os.ModePerm)
 			continue
 		}
 
 		if err := os.MkdirAll(filepath.Dir(filePath), os.ModePerm); err != nil {
-			log.Fatalln(err)
+			log.Error(err)
 		}
 
 		// check if file exists
@@ -560,16 +561,16 @@ func UnzipDirectory(destination string, source string) error {
 
 		dstFile, err := os.OpenFile(filePath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, f.Mode())
 		if err != nil {
-			log.Fatalln(err)
+			log.Error(err)
 		}
 
 		fileInArchive, err := f.Open()
 		if err != nil {
-			log.Fatalln(err)
+			log.Error(err)
 		}
 
 		if _, err := io.Copy(dstFile, fileInArchive); err != nil {
-			log.Fatalln(err)
+			log.Error(err)
 		}
 
 		dstFile.Close()
@@ -610,13 +611,35 @@ func ValidFileName(fname string) (err error) {
 	}
 	return
 }
-func CleanupTempData() {
-	path := "temp"
-	// Remove the directory and its contents
-	err := os.RemoveAll(path)
+
+const crocRemovalFile = "croc-marked-files.txt"
+
+func MarkFileForRemoval(fname string) {
+	// append the fname to the list of files to remove
+	f, err := os.OpenFile(crocRemovalFile, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0o600)
 	if err != nil {
-		log.Fatal(err)
-	} else {
-		log.Println("temp directory and its contents deleted successfully")
+		log.Debug(err)
+		return
+	}
+	defer f.Close()
+	_, err = f.WriteString(fname + "\n")
+}
+
+func RemoveMarkedFiles() (err error) {
+	// read the file and remove all the files
+	f, err := os.Open(crocRemovalFile)
+	if err != nil {
+		return
 	}
+	defer f.Close()
+	scanner := bufio.NewScanner(f)
+	for scanner.Scan() {
+		fname := scanner.Text()
+		err = os.Remove(fname)
+		if err == nil {
+			log.Tracef("Removed %s", fname)
+		}
+	}
+	os.Remove(crocRemovalFile)
+	return
 }