From ec480844a817498a83fd9e75024e7d3c7c1fb282 Mon Sep 17 00:00:00 2001
From: Joona Hoikkala <joohoi@users.noreply.github.com>
Date: Sat, 16 Nov 2019 16:51:29 +0200
Subject: [PATCH] Add SIGTERM monitoring (#101)

---
 pkg/ffuf/job.go | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/pkg/ffuf/job.go b/pkg/ffuf/job.go
index 7203dd2..0d409a3 100644
--- a/pkg/ffuf/job.go
+++ b/pkg/ffuf/job.go
@@ -4,7 +4,10 @@ import (
 	"fmt"
 	"log"
 	"math/rand"
+	"os"
+	"os/signal"
 	"sync"
+	"syscall"
 	"time"
 )
 
@@ -66,6 +69,8 @@ func (j *Job) Start() {
 		j.Output.Banner()
 	}
 	j.Running = true
+	// Monitor for SIGTERM and do cleanup properly (writing the output files etc)
+	j.interruptMonitor()
 	var wg sync.WaitGroup
 	wg.Add(1)
 	go j.runProgress(&wg)
@@ -105,6 +110,17 @@ func (j *Job) Start() {
 	return
 }
 
+func (j *Job) interruptMonitor() {
+	sigChan := make(chan os.Signal, 2)
+	signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
+	go func() {
+		for _ = range sigChan {
+			j.Error = "Caught keyboard interrupt (Ctrl-C)\n"
+			j.Stop()
+		}
+	}()
+}
+
 func (j *Job) runProgress(wg *sync.WaitGroup) {
 	defer wg.Done()
 	j.startTime = time.Now()