diff --git a/main.go b/main.go index 65ad69a9..f36d2e80 100644 --- a/main.go +++ b/main.go @@ -6,8 +6,12 @@ package main import ( "log" + "os" + "os/signal" + "syscall" "github.com/schollz/croc/v10/src/cli" + "github.com/schollz/croc/v10/src/utils" ) func main() { @@ -27,7 +31,25 @@ func main() { // fmt.Println("wrote profile") // } // }() - if err := cli.Run(); err != nil { - log.Fatalln(err) - } + + // Create a channel to receive OS signals + sigs := make(chan os.Signal, 1) + signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM) + + go func() { + if err := cli.Run(); err != nil { + log.Fatalln(err) + } + }() + + // Wait for a termination signal + sig := <-sigs + log.Println("Received signal:", sig) + + // Perform any necessary cleanup here + log.Println("Performing cleanup...") + utils.CleanupTempData() + + // Exit the program gracefully + os.Exit(0) } diff --git a/src/cli/cli.go b/src/cli/cli.go index d65281aa..fa856082 100644 --- a/src/cli/cli.go +++ b/src/cli/cli.go @@ -446,9 +446,15 @@ 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) { - f, err := os.CreateTemp(".", "croc-stdin-") + makeTempFolder() + f, err := os.CreateTemp("temp", "croc-stdin-") if err != nil { return } diff --git a/src/utils/utils.go b/src/utils/utils.go index 1f09c2ac..3538aa8c 100644 --- a/src/utils/utils.go +++ b/src/utils/utils.go @@ -610,3 +610,13 @@ func ValidFileName(fname string) (err error) { } return } +func CleanupTempData() { + path := "temp" + // Remove the directory and its contents + err := os.RemoveAll(path) + if err != nil { + log.Fatal(err) + } else { + log.Println("temp directory and its contents deleted successfully") + } +}