Feature/run scan with progress sync function (#60)

Co-authored-by: Elias Flötzinger <[email protected]>

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)
+}

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=

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 {

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

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"