Skip to content

Commit c306ae2

Browse files
committed
feat(cli): cancel the run on interrupt and add -max-time
a ctrl-c only killed the process mid-write and there was no way to bound a long run. wire a signal-cancelled context from main through App.Run, and add a -max-time flag that caps the whole run. both the interrupt and the deadline cancel the scanners that already take a context (port scan, module executor, notify) and stop the target loop between steps, so the run ends cleanly with whatever was collected still reported rather than being hard-killed. the ctx threads through scanAllTargets into scanTarget, and both the sequential and the -concurrency pool path check it before starting another target, so a cancelled run stops dispatching instead of draining the queue.
1 parent a38ba0a commit c306ae2

5 files changed

Lines changed: 61 additions & 20 deletions

File tree

cmd/sif/main.go

Lines changed: 20 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,11 @@
1313
package main
1414

1515
import (
16+
"context"
1617
"fmt"
1718
"os"
19+
"os/signal"
20+
"syscall"
1821

1922
"github.com/charmbracelet/log"
2023
"github.com/vmfunc/sif"
@@ -45,11 +48,20 @@ func main() {
4548
}
4649
}
4750

51+
if err := run(); err != nil {
52+
log.Fatal(err)
53+
}
54+
}
55+
56+
// run wires up the app and executes it, kept separate from main so its deferred
57+
// signal cleanup actually fires (main's log.Fatal calls os.Exit, which would
58+
// skip a defer placed there).
59+
func run() error {
4860
settings := config.Parse()
4961

5062
app, err := sif.New(settings)
5163
if err != nil {
52-
log.Fatal(err)
64+
return err
5365
}
5466

5567
// patchnotes print to stdout; skip them in api/silent mode so the only thing
@@ -58,8 +70,11 @@ func main() {
5870
patchnotes.ShowOnce(version)
5971
}
6072

61-
err = app.Run()
62-
if err != nil {
63-
log.Fatal(err)
64-
}
73+
// cancel the run on the first interrupt so a ctrl-c stops between scan steps
74+
// instead of only killing the process mid-write. a second interrupt still
75+
// hard-kills, since NotifyContext stops trapping once fired.
76+
ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
77+
defer stop()
78+
79+
return app.Run(ctx)
6580
}

internal/config/config.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ type Settings struct {
4444
Nuclei bool
4545
JavaScript bool
4646
Timeout time.Duration
47+
MaxTime time.Duration // abort the whole run after this long (0 = no limit)
4748
URLs goflags.StringSlice
4849
File string
4950
ApiMode bool
@@ -177,6 +178,7 @@ func registerFlags(settings *Settings) *goflags.FlagSet {
177178
flagSet.CreateGroup("runtime", "Runtime",
178179
flagSet.BoolVarP(&settings.Debug, "debug", "d", false, "Enable debug logging"),
179180
flagSet.DurationVarP(&settings.Timeout, "timeout", "t", 10*time.Second, "HTTP request timeout"),
181+
flagSet.DurationVar(&settings.MaxTime, "max-time", 0, "Abort the whole run after this duration (0 = no limit)"),
180182
flagSet.StringVarP(&settings.LogDir, "log", "l", "", "Directory to store logs in"),
181183
flagSet.IntVar(&settings.Threads, "threads", 10, "Number of threads to run scans on"),
182184
flagSet.IntVar(&settings.Concurrency, "concurrency", 1, "Number of targets to scan in parallel (>1 interleaves console output)"),

sif.go

Lines changed: 33 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -234,7 +234,7 @@ func normalizeTarget(target string) (string, error) {
234234

235235
// Run runs the pentesting suite, with the targets specified, according to the
236236
// settings specified.
237-
func (app *App) Run() error {
237+
func (app *App) Run(ctx context.Context) error {
238238
// Handle --list-modules before any other processing
239239
if app.settings.ListModules {
240240
loader, err := modules.NewLoader()
@@ -298,6 +298,15 @@ func (app *App) Run() error {
298298
}
299299
}
300300

301+
// bound the whole run when -max-time is set; the deadline rides on the same
302+
// ctx as the interrupt handler, so either one cancels the in-flight scanners
303+
// that take a context and stops the target loop between steps.
304+
if app.settings.MaxTime > 0 {
305+
var cancel context.CancelFunc
306+
ctx, cancel = context.WithTimeout(ctx, app.settings.MaxTime)
307+
defer cancel()
308+
}
309+
301310
scansRun := make([]string, 0, 16)
302311

303312
// accumulate every module result across targets so the report writers can
@@ -323,7 +332,7 @@ func (app *App) Run() error {
323332
}
324333
}
325334

326-
results, err := app.scanAllTargets(storeDir, wantReport)
335+
results, err := app.scanAllTargets(ctx, storeDir, wantReport)
327336
if err != nil {
328337
return err
329338
}
@@ -339,7 +348,7 @@ func (app *App) Run() error {
339348
}
340349
}
341350

342-
return app.finishRun(scansRun, allFindings, reportResults, wantReport)
351+
return app.finishRun(ctx, scansRun, allFindings, reportResults, wantReport)
343352
}
344353

345354
// targetScan holds one target's isolated scan output: its findings, report rows,
@@ -358,7 +367,7 @@ type targetScan struct {
358367
// console, which output.SetConcurrent serializes and de-animates. Results are
359368
// indexed by target position, so the caller merges them in a stable order no
360369
// matter which worker finished first.
361-
func (app *App) scanAllTargets(storeDir string, wantReport bool) ([]targetScan, error) {
370+
func (app *App) scanAllTargets(ctx context.Context, storeDir string, wantReport bool) ([]targetScan, error) {
362371
results := make([]targetScan, len(app.targets))
363372

364373
concurrency := app.settings.Concurrency
@@ -371,7 +380,13 @@ func (app *App) scanAllTargets(storeDir string, wantReport bool) ([]targetScan,
371380

372381
if concurrency <= 1 {
373382
for i, url := range app.targets {
374-
ts, err := app.scanTarget(url, storeDir, wantReport)
383+
// stop cleanly on interrupt or -max-time rather than starting another
384+
// target; whatever was collected so far still gets reported.
385+
if ctx.Err() != nil {
386+
log.Warnf("scan cancelled, not starting further targets: %v", ctx.Err())
387+
break
388+
}
389+
ts, err := app.scanTarget(ctx, url, storeDir, wantReport)
375390
if err != nil {
376391
return nil, err
377392
}
@@ -386,6 +401,10 @@ func (app *App) scanAllTargets(storeDir string, wantReport bool) ([]targetScan,
386401
sem := make(chan struct{}, concurrency)
387402
var wg sync.WaitGroup
388403
for i, url := range app.targets {
404+
if ctx.Err() != nil {
405+
log.Warnf("scan cancelled, not starting further targets: %v", ctx.Err())
406+
break
407+
}
389408
wg.Add(1)
390409
sem <- struct{}{}
391410
go func(i int, url string) {
@@ -399,7 +418,7 @@ func (app *App) scanAllTargets(storeDir string, wantReport bool) ([]targetScan,
399418
errs[i] = fmt.Errorf("panic scanning %s: %v", url, r)
400419
}
401420
}()
402-
results[i], errs[i] = app.scanTarget(url, storeDir, wantReport)
421+
results[i], errs[i] = app.scanTarget(ctx, url, storeDir, wantReport)
403422
}(i, url)
404423
}
405424
wg.Wait()
@@ -414,7 +433,7 @@ func (app *App) scanAllTargets(storeDir string, wantReport bool) ([]targetScan,
414433

415434
// scanTarget runs the full scanner set for one target and returns its isolated
416435
// accumulators without mutating run-wide state.
417-
func (app *App) scanTarget(url, storeDir string, wantReport bool) (targetScan, error) {
436+
func (app *App) scanTarget(ctx context.Context, url, storeDir string, wantReport bool) (targetScan, error) {
418437
var scansRun []string
419438
var logFiles []string
420439

@@ -501,7 +520,7 @@ func (app *App) scanTarget(url, storeDir string, wantReport bool) (targetScan, e
501520
}
502521

503522
if app.settings.Ports != "none" {
504-
result, err := scan.Ports(context.Background(), app.settings.Ports, url, app.settings.Timeout, app.settings.Threads, app.settings.LogDir)
523+
result, err := scan.Ports(ctx, app.settings.Ports, url, app.settings.Timeout, app.settings.Threads, app.settings.LogDir)
505524
if err != nil {
506525
log.Errorf("Error while running port scan: %s", err)
507526
} else {
@@ -769,6 +788,9 @@ func (app *App) scanTarget(url, storeDir string, wantReport bool) (targetScan, e
769788
}
770789

771790
for _, m := range toRun {
791+
if ctx.Err() != nil {
792+
break
793+
}
772794
switch m.Info().ID {
773795
case "nuclei-scan":
774796
if app.settings.Nuclei {
@@ -789,7 +811,7 @@ func (app *App) scanTarget(url, storeDir string, wantReport bool) (targetScan, e
789811
}
790812
modLog := output.Module(m.Info().ID)
791813
modLog.Start()
792-
result, err := m.Execute(context.Background(), url, opts)
814+
result, err := m.Execute(ctx, url, opts)
793815
if err != nil {
794816
modLog.Error("failed: %v", err)
795817
continue
@@ -846,15 +868,15 @@ func (app *App) scanTarget(url, storeDir string, wantReport bool) (targetScan, e
846868
// finishRun performs the run-wide steps after every target has been scanned:
847869
// notify, the silent findings stream, report files and the summary. It consumes
848870
// the merged accumulators so per-target scanning stays isolated in scanTarget.
849-
func (app *App) finishRun(scansRun []string, allFindings []finding.Finding, reportResults []report.Result, wantReport bool) error {
871+
func (app *App) finishRun(ctx context.Context, scansRun []string, allFindings []finding.Finding, reportResults []report.Result, wantReport bool) error {
850872
// the normalized findings are the handoff point for notify/diff; surface the
851873
// count now so the path is live and observable without changing output.
852874
log.Debugf("normalized %d findings across %d targets", len(allFindings), len(app.targets))
853875

854876
// notify: ship the severity-filtered findings to any configured provider.
855877
// kept as an isolated block so it merges cleanly with the diff-store bundle.
856878
if app.settings.Notify {
857-
if err := app.notifyFindings(context.Background(), allFindings); err != nil {
879+
if err := app.notifyFindings(ctx, allFindings); err != nil {
858880
log.Errorf("notify: %v", err)
859881
}
860882
}

sif_concurrency_test.go

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
package sif
1414

1515
import (
16+
"context"
1617
"testing"
1718

1819
"github.com/vmfunc/sif/internal/output"
@@ -36,7 +37,7 @@ func TestScanAllTargetsConcurrentIsolation(t *testing.T) {
3637
app.targets = []string{srvA.URL, srvB.URL, srvC.URL}
3738
app.settings.Concurrency = 3
3839

39-
results, err := app.scanAllTargets("", false)
40+
results, err := app.scanAllTargets(context.Background(), "", false)
4041
if err != nil {
4142
t.Fatalf("scanAllTargets: %v", err)
4243
}
@@ -62,7 +63,7 @@ func TestScanAllTargetsSequentialMatchesInputOrder(t *testing.T) {
6263
app.targets = []string{srvA.URL, srvB.URL}
6364
app.settings.Concurrency = 1
6465

65-
results, err := app.scanAllTargets("", false)
66+
results, err := app.scanAllTargets(context.Background(), "", false)
6667
if err != nil {
6768
t.Fatalf("scanAllTargets: %v", err)
6869
}

sif_scantarget_test.go

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
package sif
1414

1515
import (
16+
"context"
1617
"net/http"
1718
"net/http/httptest"
1819
"testing"
@@ -55,15 +56,15 @@ func TestScanTargetIsolatesPerTargetState(t *testing.T) {
5556

5657
app := headersOnlyApp()
5758

58-
tsA, err := app.scanTarget(srvA.URL, "", false)
59+
tsA, err := app.scanTarget(context.Background(), srvA.URL, "", false)
5960
if err != nil {
6061
t.Fatalf("scanTarget(A): %v", err)
6162
}
6263
if len(tsA.scansRun) != 1 || tsA.scansRun[0] != "HTTP Headers" {
6364
t.Fatalf("target A scansRun = %v, want exactly [HTTP Headers]", tsA.scansRun)
6465
}
6566

66-
tsB, err := app.scanTarget(srvB.URL, "", false)
67+
tsB, err := app.scanTarget(context.Background(), srvB.URL, "", false)
6768
if err != nil {
6869
t.Fatalf("scanTarget(B): %v", err)
6970
}

0 commit comments

Comments
 (0)