diff --git a/cmd/sif/main.go b/cmd/sif/main.go index 3f676fcb..3feb07fd 100644 --- a/cmd/sif/main.go +++ b/cmd/sif/main.go @@ -13,8 +13,11 @@ package main import ( + "context" "fmt" "os" + "os/signal" + "syscall" "github.com/charmbracelet/log" "github.com/vmfunc/sif" @@ -45,11 +48,20 @@ func main() { } } + if err := run(); err != nil { + log.Fatal(err) + } +} + +// run wires up the app and executes it, kept separate from main so its deferred +// signal cleanup actually fires (main's log.Fatal calls os.Exit, which would +// skip a defer placed there). +func run() error { settings := config.Parse() app, err := sif.New(settings) if err != nil { - log.Fatal(err) + return err } // patchnotes print to stdout; skip them in api/silent mode so the only thing @@ -58,8 +70,11 @@ func main() { patchnotes.ShowOnce(version) } - err = app.Run() - if err != nil { - log.Fatal(err) - } + // cancel the run on the first interrupt so a ctrl-c stops between scan steps + // instead of only killing the process mid-write. a second interrupt still + // hard-kills, since NotifyContext stops trapping once fired. + ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) + defer stop() + + return app.Run(ctx) } diff --git a/internal/config/config.go b/internal/config/config.go index dabc82f0..a15801bf 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -44,6 +44,7 @@ type Settings struct { Nuclei bool JavaScript bool Timeout time.Duration + MaxTime time.Duration // abort the whole run after this long (0 = no limit) URLs goflags.StringSlice File string ApiMode bool @@ -177,6 +178,7 @@ func registerFlags(settings *Settings) *goflags.FlagSet { flagSet.CreateGroup("runtime", "Runtime", flagSet.BoolVarP(&settings.Debug, "debug", "d", false, "Enable debug logging"), flagSet.DurationVarP(&settings.Timeout, "timeout", "t", 10*time.Second, "HTTP request timeout"), + flagSet.DurationVar(&settings.MaxTime, "max-time", 0, "Abort the whole run after this duration (0 = no limit)"), flagSet.StringVarP(&settings.LogDir, "log", "l", "", "Directory to store logs in"), flagSet.IntVar(&settings.Threads, "threads", 10, "Number of threads to run scans on"), flagSet.IntVar(&settings.Concurrency, "concurrency", 1, "Number of targets to scan in parallel (>1 interleaves console output)"), diff --git a/sif.go b/sif.go index 481e4fb2..af91204d 100644 --- a/sif.go +++ b/sif.go @@ -234,7 +234,7 @@ func normalizeTarget(target string) (string, error) { // Run runs the pentesting suite, with the targets specified, according to the // settings specified. -func (app *App) Run() error { +func (app *App) Run(ctx context.Context) error { // Handle --list-modules before any other processing if app.settings.ListModules { loader, err := modules.NewLoader() @@ -298,6 +298,15 @@ func (app *App) Run() error { } } + // bound the whole run when -max-time is set; the deadline rides on the same + // ctx as the interrupt handler, so either one cancels the in-flight scanners + // that take a context and stops the target loop between steps. + if app.settings.MaxTime > 0 { + var cancel context.CancelFunc + ctx, cancel = context.WithTimeout(ctx, app.settings.MaxTime) + defer cancel() + } + scansRun := make([]string, 0, 16) // accumulate every module result across targets so the report writers can @@ -323,7 +332,7 @@ func (app *App) Run() error { } } - results, err := app.scanAllTargets(storeDir, wantReport) + results, err := app.scanAllTargets(ctx, storeDir, wantReport) if err != nil { return err } @@ -339,7 +348,7 @@ func (app *App) Run() error { } } - return app.finishRun(scansRun, allFindings, reportResults, wantReport) + return app.finishRun(ctx, scansRun, allFindings, reportResults, wantReport) } // targetScan holds one target's isolated scan output: its findings, report rows, @@ -358,7 +367,7 @@ type targetScan struct { // console, which output.SetConcurrent serializes and de-animates. Results are // indexed by target position, so the caller merges them in a stable order no // matter which worker finished first. -func (app *App) scanAllTargets(storeDir string, wantReport bool) ([]targetScan, error) { +func (app *App) scanAllTargets(ctx context.Context, storeDir string, wantReport bool) ([]targetScan, error) { results := make([]targetScan, len(app.targets)) concurrency := app.settings.Concurrency @@ -371,7 +380,13 @@ func (app *App) scanAllTargets(storeDir string, wantReport bool) ([]targetScan, if concurrency <= 1 { for i, url := range app.targets { - ts, err := app.scanTarget(url, storeDir, wantReport) + // stop cleanly on interrupt or -max-time rather than starting another + // target; whatever was collected so far still gets reported. + if ctx.Err() != nil { + log.Warnf("scan cancelled, not starting further targets: %v", ctx.Err()) + break + } + ts, err := app.scanTarget(ctx, url, storeDir, wantReport) if err != nil { return nil, err } @@ -386,6 +401,10 @@ func (app *App) scanAllTargets(storeDir string, wantReport bool) ([]targetScan, sem := make(chan struct{}, concurrency) var wg sync.WaitGroup for i, url := range app.targets { + if ctx.Err() != nil { + log.Warnf("scan cancelled, not starting further targets: %v", ctx.Err()) + break + } wg.Add(1) sem <- struct{}{} go func(i int, url string) { @@ -399,7 +418,7 @@ func (app *App) scanAllTargets(storeDir string, wantReport bool) ([]targetScan, errs[i] = fmt.Errorf("panic scanning %s: %v", url, r) } }() - results[i], errs[i] = app.scanTarget(url, storeDir, wantReport) + results[i], errs[i] = app.scanTarget(ctx, url, storeDir, wantReport) }(i, url) } wg.Wait() @@ -414,7 +433,7 @@ func (app *App) scanAllTargets(storeDir string, wantReport bool) ([]targetScan, // scanTarget runs the full scanner set for one target and returns its isolated // accumulators without mutating run-wide state. -func (app *App) scanTarget(url, storeDir string, wantReport bool) (targetScan, error) { +func (app *App) scanTarget(ctx context.Context, url, storeDir string, wantReport bool) (targetScan, error) { var scansRun []string var logFiles []string @@ -501,7 +520,7 @@ func (app *App) scanTarget(url, storeDir string, wantReport bool) (targetScan, e } if app.settings.Ports != "none" { - result, err := scan.Ports(context.Background(), app.settings.Ports, url, app.settings.Timeout, app.settings.Threads, app.settings.LogDir) + result, err := scan.Ports(ctx, app.settings.Ports, url, app.settings.Timeout, app.settings.Threads, app.settings.LogDir) if err != nil { log.Errorf("Error while running port scan: %s", err) } else { @@ -769,6 +788,9 @@ func (app *App) scanTarget(url, storeDir string, wantReport bool) (targetScan, e } for _, m := range toRun { + if ctx.Err() != nil { + break + } switch m.Info().ID { case "nuclei-scan": if app.settings.Nuclei { @@ -789,7 +811,7 @@ func (app *App) scanTarget(url, storeDir string, wantReport bool) (targetScan, e } modLog := output.Module(m.Info().ID) modLog.Start() - result, err := m.Execute(context.Background(), url, opts) + result, err := m.Execute(ctx, url, opts) if err != nil { modLog.Error("failed: %v", err) continue @@ -846,7 +868,7 @@ func (app *App) scanTarget(url, storeDir string, wantReport bool) (targetScan, e // finishRun performs the run-wide steps after every target has been scanned: // notify, the silent findings stream, report files and the summary. It consumes // the merged accumulators so per-target scanning stays isolated in scanTarget. -func (app *App) finishRun(scansRun []string, allFindings []finding.Finding, reportResults []report.Result, wantReport bool) error { +func (app *App) finishRun(ctx context.Context, scansRun []string, allFindings []finding.Finding, reportResults []report.Result, wantReport bool) error { // the normalized findings are the handoff point for notify/diff; surface the // count now so the path is live and observable without changing output. log.Debugf("normalized %d findings across %d targets", len(allFindings), len(app.targets)) @@ -854,7 +876,7 @@ func (app *App) finishRun(scansRun []string, allFindings []finding.Finding, repo // notify: ship the severity-filtered findings to any configured provider. // kept as an isolated block so it merges cleanly with the diff-store bundle. if app.settings.Notify { - if err := app.notifyFindings(context.Background(), allFindings); err != nil { + if err := app.notifyFindings(ctx, allFindings); err != nil { log.Errorf("notify: %v", err) } } diff --git a/sif_concurrency_test.go b/sif_concurrency_test.go index fbed38d4..ef6520c9 100644 --- a/sif_concurrency_test.go +++ b/sif_concurrency_test.go @@ -13,6 +13,7 @@ package sif import ( + "context" "testing" "github.com/vmfunc/sif/internal/output" @@ -36,7 +37,7 @@ func TestScanAllTargetsConcurrentIsolation(t *testing.T) { app.targets = []string{srvA.URL, srvB.URL, srvC.URL} app.settings.Concurrency = 3 - results, err := app.scanAllTargets("", false) + results, err := app.scanAllTargets(context.Background(), "", false) if err != nil { t.Fatalf("scanAllTargets: %v", err) } @@ -62,7 +63,7 @@ func TestScanAllTargetsSequentialMatchesInputOrder(t *testing.T) { app.targets = []string{srvA.URL, srvB.URL} app.settings.Concurrency = 1 - results, err := app.scanAllTargets("", false) + results, err := app.scanAllTargets(context.Background(), "", false) if err != nil { t.Fatalf("scanAllTargets: %v", err) } diff --git a/sif_scantarget_test.go b/sif_scantarget_test.go index 72af8b0f..fc741bd8 100644 --- a/sif_scantarget_test.go +++ b/sif_scantarget_test.go @@ -13,6 +13,7 @@ package sif import ( + "context" "net/http" "net/http/httptest" "testing" @@ -55,7 +56,7 @@ func TestScanTargetIsolatesPerTargetState(t *testing.T) { app := headersOnlyApp() - tsA, err := app.scanTarget(srvA.URL, "", false) + tsA, err := app.scanTarget(context.Background(), srvA.URL, "", false) if err != nil { t.Fatalf("scanTarget(A): %v", err) } @@ -63,7 +64,7 @@ func TestScanTargetIsolatesPerTargetState(t *testing.T) { t.Fatalf("target A scansRun = %v, want exactly [HTTP Headers]", tsA.scansRun) } - tsB, err := app.scanTarget(srvB.URL, "", false) + tsB, err := app.scanTarget(context.Background(), srvB.URL, "", false) if err != nil { t.Fatalf("scanTarget(B): %v", err) }