diff --git a/examples/basic_scan_progress/main.go b/examples/basic_scan_progress/main.go new file mode 100644 index 0000000..668017c --- /dev/null +++ b/examples/basic_scan_progress/main.go @@ -0,0 +1,36 @@ +package main + +import ( + "fmt" + "log" + + "github.com/Ullaakut/nmap/v2" +) + +func main() { + scanner, err := nmap.NewScanner( + nmap.WithTargets("localhost"), + nmap.WithPorts("1-4000"), + nmap.WithServiceInfo(), + nmap.WithVerbosity(3), + ) + if err != nil { + log.Fatalf("unable to create nmap scanner: %v", err) + } + + progress := make(chan float32, 1) + + // Function to listen and print the progress + go func() { + for p := range progress { + fmt.Printf("Progress: %v %%\n", p) + } + }() + + result, _, err := scanner.RunWithProgress(progress) + if err != nil { + log.Fatalf("unable to run nmap scan: %v", err) + } + + fmt.Printf("Nmap done: %d hosts up scanned in %.2f seconds\n", len(result.Hosts), result.Stats.Finished.Elapsed) +} diff --git a/go.sum b/go.sum index 56d62e7..afe7890 100644 --- a/go.sum +++ b/go.sum @@ -5,6 +5,7 @@ github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZN github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/testify v1.6.1 h1:hDPOHmpOpP40lSULcqw7IrRb/u7w6RpDC9399XyoNd0= github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c h1:dUUwHk2QECo/6vqA44rthZ8ie2QXMNeKRTHCNY2nXvo= gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/nmap.go b/nmap.go index 2c42357..072b34b 100644 --- a/nmap.go +++ b/nmap.go @@ -5,6 +5,7 @@ import ( "bufio" "bytes" "context" + "encoding/xml" "fmt" "os/exec" "strings" @@ -139,6 +140,123 @@ func (s *Scanner) Run() (result *Run, warnings []string, err error) { } } +// RunWithProgress runs nmap synchronously and returns the result of the scan. +// It needs a channel to constantly stream the progress. +func (s *Scanner) RunWithProgress(liveProgress chan<- float32) (result *Run, warnings []string, err error) { + var stdout, stderr bytes.Buffer + + // Enable XML output + s.args = append(s.args, "-oX") + + // Get XML output in stdout instead of writing it in a file + s.args = append(s.args, "-") + + // Enable progress output every second + s.args = append(s.args, "--stats-every", "1s") + + // Prepare nmap process + cmd := exec.Command(s.binaryPath, s.args...) + cmd.Stderr = &stderr + cmd.Stdout = &stdout + + // Run nmap process + err = cmd.Start() + if err != nil { + return nil, warnings, err + } + + // Make a goroutine to notify the select when the scan is done. + done := make(chan error, 1) + doneProgress := make(chan bool, 1) + go func() { + done <- cmd.Wait() + }() + + // Make goroutine to check the progress every second + // Listening for channel doneProgress + go func() { + type progress struct { + TaskProgress []TaskProgress `xml:"taskprogress" json:"task_progress"` + } + p := &progress{} + for { + select { + case <- doneProgress: + close(liveProgress) + return + default: + time.Sleep(time.Second) + _ = xml.Unmarshal(stdout.Bytes(), p) + //result, _ := Parse(stdout.Bytes()) + if len(p.TaskProgress) > 0 { + liveProgress <- p.TaskProgress[len(p.TaskProgress)-1].Percent + } + } + } + }() + + // Wait for nmap process or timeout + select { + case <-s.ctx.Done(): + + // Context was done before the scan was finished. + // The process is killed and a timeout error is returned. + _ = cmd.Process.Kill() + + return nil, warnings, ErrScanTimeout + case <-done: + + // Trigger progress function exit + close(doneProgress) + + // Process nmap stderr output containing none-critical errors and warnings + // Everyone needs to check whether one or some of these warnings is a hard issue in their use case + if stderr.Len() > 0 { + warnings = strings.Split(strings.Trim(stderr.String(), "\n"), "\n") + } + + // Check for warnings that will inevitable lead to parsing errors, hence, have priority + for _, warning := range warnings { + switch { + case strings.Contains(warning, "Malloc Failed!"): + return nil, warnings, ErrMallocFailed + // TODO: Add cases for other known errors we might want to guard. + default: + } + } + + // Parse nmap xml output. Usually nmap always returns valid XML, even if there is a scan error. + // Potentially available warnings are returned too, but probably not the reason for a broken XML. + result, err := Parse(stdout.Bytes()) + if err != nil { + warnings = append(warnings, err.Error()) // Append parsing error to warnings for those who are interested. + return nil, warnings, ErrParseOutput + } + + // Critical scan errors are reflected in the XML. + if result != nil && len(result.Stats.Finished.ErrorMsg) > 0 { + switch { + case strings.Contains(result.Stats.Finished.ErrorMsg, "Error resolving name"): + return result, warnings, ErrResolveName + // TODO: Add cases for other known errors we might want to guard. + default: + return result, warnings, fmt.Errorf(result.Stats.Finished.ErrorMsg) + } + } + + // Call filters if they are set. + if s.portFilter != nil { + result = choosePorts(result, s.portFilter) + } + if s.hostFilter != nil { + result = chooseHosts(result, s.hostFilter) + } + + // Return result, optional warnings but no error + return result, warnings, nil + } +} + // RunAsync runs nmap asynchronously and returns error. // TODO: RunAsync should return warnings as well. func (s *Scanner) RunAsync() error { diff --git a/nmap_test.go b/nmap_test.go index c322e9b..7b7ed78 100644 --- a/nmap_test.go +++ b/nmap_test.go @@ -4,6 +4,7 @@ import ( "context" "encoding/xml" "errors" + "io/ioutil" "os" "os/exec" "reflect" @@ -242,6 +243,70 @@ func TestRun(t *testing.T) { } } +func TestRunWithProgress(t *testing.T) { + // Open and parse sample result for testing + dat, err := ioutil.ReadFile("tests/xml/scan_base.xml") + if err != nil { + panic(err) + } + + r, _ := Parse(dat) + + tests := []struct { + description string + + options []func(*Scanner) + + compareWholeRun bool + + expectedResult *Run + expectedProgress []float32 + expectedErr error + expectedWarnings []string + }{ + { + description: "fake scan with slow output for progress streaming", + options: []func(*Scanner){ + WithBinaryPath("tests/scripts/fake_nmap_delay.sh"), + WithCustomArguments("tests/xml/scan_base.xml"), + }, + + compareWholeRun: true, + expectedResult: r, + expectedProgress: []float32{56.66, 81.95, 87.84, 94.43, 97.76, 97.76}, + expectedErr: nil, + }, + } + + for _, test := range tests { + t.Run(test.description, func(t *testing.T) { + s, err := NewScanner(test.options...) + if err != nil { + panic(err) // this is never supposed to err, as we are testing run and not new. + } + + progress := make(chan float32, 5) + result, _, err := s.RunWithProgress(progress) + assert.Equal(t, test.expectedErr, err) + if err != nil { + return + } + + // Test if channel data compares to given progress array + var progressOutput []float32 + for n := range progress { + progressOutput = append(progressOutput, n) + } + assert.Equal(t, test.expectedProgress, progressOutput) + + // Test if read output equals parsed xml file + if test.compareWholeRun { + assert.Equal(t, test.expectedResult.Hosts, result.Hosts) + } + }) + } +} + func TestRunAsync(t *testing.T) { tests := []struct { description string diff --git a/tests/scripts/fake_nmap_delay.sh b/tests/scripts/fake_nmap_delay.sh new file mode 100755 index 0000000..73bd33e --- /dev/null +++ b/tests/scripts/fake_nmap_delay.sh @@ -0,0 +1,7 @@ +#!/bin/bash +input=$1 +while IFS= read -r line +do + echo "$line" + sleep .5 +done < "$input"