Add context version of main functions

gomap.go

@@ -2,9 +2,11 @@ package gomap
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"fmt"
 	"net"
+	"net/netip"
 )
 
 // IPScanResult contains the results of a scan on a single ip
@@ -55,32 +57,52 @@ type RangeScanResult []*IPScanResult
 
 // ScanIP scans a single IP for open ports
 func ScanIP(hostname string, proto string, fastscan bool, stealth bool) (*IPScanResult, error) {
+	return ScanIPCtx(context.Background(), hostname, proto, fastscan, stealth)
+}
+
+// ScanIPCtx scans a single IP for open ports
+func ScanIPCtx(
+	ctx context.Context, hostname string, proto string, fastscan bool, stealth bool, ports ...int,
+) (*IPScanResult, error) {
 	laddr, err := getLocalIP()
 	if err != nil {
 		return nil, err
 	}
 
 	if stealth {
-		if canSocketBind(laddr) == false {
+		if !canSocketBind(laddr) {
 			return nil, fmt.Errorf("socket: operation not permitted")
 		}
 	}
-	return scanIPPorts(hostname, laddr, proto, fastscan, stealth)
+
+	return scanIPPorts(ctx, hostname, laddr, proto, fastscan, stealth, ports...)
 }
 
 // ScanRange scans every address on a CIDR for open ports
-func ScanRange(proto string, fastscan bool, stealth bool) (RangeScanResult, error) {
+func ScanRange(prefix netip.Prefix, proto string, fastscan bool, stealth bool) (RangeScanResult, error) {
+	return ScanRangeCtx(context.Background(), prefix, proto, fastscan, stealth)
+}
+
+// ScanRangeCtx scans every address on a CIDR for open ports
+func ScanRangeCtx(
+	ctx context.Context, prefix netip.Prefix, proto string, fastscan bool, stealth bool, ports ...int,
+) (RangeScanResult, error) {
 	laddr, err := getLocalIP()
 	if err != nil {
 		return nil, err
 	}
 
 	if stealth {
-		if canSocketBind(laddr) == false {
+		if !canSocketBind(laddr) {
 			return nil, fmt.Errorf("socket: operation not permitted")
 		}
 	}
-	return scanIPRange(laddr, proto, fastscan, stealth)
+
+	return scanIPRange(ctx, prefix, laddr, proto, fastscan, stealth, ports...)
+}
+
+func ScanLocal(ctx context.Context, proto string, fastscan bool, stealth bool, ports ...int) (RangeScanResult, error) {
+	return ScanRangeCtx(ctx, getLocalRange(), proto, fastscan, stealth, ports...)
 }
 
 // String with the results of a single scanned IP
@@ -146,7 +168,7 @@ func (results RangeScanResult) String() string {
 func (results *IPScanResult) Json() (string, error) {
 	var ipdata JsonIP
 	fmt.Println(results.IP)
-	ipdata.IP = fmt.Sprintf("%s", results.IP[len(results.IP)-1])
+	ipdata.IP = results.IP[len(results.IP)-1].String()
 	ipdata.Hostname = results.Hostname
 
 	active := false
@@ -180,7 +202,7 @@ func (results RangeScanResult) Json() (string, error) {
 
 	for _, r := range results {
 		var ipdata JsonIP
-		ipdata.IP = fmt.Sprintf("%s", r.IP[len(r.IP)-1])
+		ipdata.IP = r.IP[len(r.IP)-1].String()
 		ipdata.Hostname = r.Hostname
 
 		active := false

gomap_funcs.go

@@ -5,7 +5,7 @@ import (
 	"fmt"
 	"log"
 	"net"
-	"strings"
+	"net/netip"
 )
 
 func canSocketBind(laddr string) bool {
@@ -25,8 +25,8 @@ func canSocketBind(laddr string) bool {
 }
 
 // createHostRange converts a input ip addr string to a slice of ips on the cidr
-func createHostRange(netw string) []string {
-	_, ipv4Net, err := net.ParseCIDR(netw)
+func createHostRange(netw netip.Prefix) []string {
+	_, ipv4Net, err := net.ParseCIDR(netw.Masked().String())
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -36,7 +36,7 @@ func createHostRange(netw string) []string {
 	finish := (start & mask) | (mask ^ 0xffffffff)
 
 	var hosts []string
-	for i := start + 1; i <= finish-1; i++ {
+	for i := start; i <= finish; i++ {
 		ip := make(net.IP, 4)
 		binary.BigEndian.PutUint32(ip, i)
 		hosts = append(hosts, ip.String())
@@ -46,21 +46,22 @@ func createHostRange(netw string) []string {
 }
 
 // getLocalRange returns local ip range or defaults on error to most common
-func getLocalRange() string {
+func getLocalRange() netip.Prefix {
 	addrs, err := net.InterfaceAddrs()
 	if err != nil {
-		return "192.168.1.0/24"
+		return netip.MustParsePrefix("192.168.1.0/24")
 	}
+
 	for _, address := range addrs {
 		// check the address type and if it is not a loopback the display it
 		if ipnet, ok := address.(*net.IPNet); ok && !ipnet.IP.IsLoopback() {
-			if ipnet.IP.To4() != nil {
-				split := strings.Split(ipnet.IP.String(), ".")
-				return split[0] + "." + split[1] + "." + split[2] + ".0/24"
+			if v4 := ipnet.IP.To4(); v4 != nil {
+				netip.PrefixFrom(netip.AddrFrom4([4]byte{v4[0], v4[1], v4[2], v4[3]}), 24)
 			}
 		}
 	}
-	return "192.168.1.0/24"
+
+	return netip.MustParsePrefix("192.168.1.0/24")
 }
 
 // getLocalRange returns local ip range or defaults on error to most common
@@ -77,5 +78,5 @@ func getLocalIP() (string, error) {
 			}
 		}
 	}
-	return "", fmt.Errorf("No IP Found")
+	return "", fmt.Errorf("no IP found")
 }

gomap_scan.go

@@ -1,35 +1,40 @@
 package gomap
 
 import (
+	"context"
 	"fmt"
 	"net"
+	"net/netip"
 	"strconv"
+	"sync"
 	"time"
 )
 
 // scanIPRange scans an entire cidr range for open ports
 // I am fairly happy with this code since its just iterating
 // over scanIPPorts. Most issues are deeper in the code.
-func scanIPRange(laddr string, proto string, fastscan bool, stealth bool) (RangeScanResult, error) {
-	iprange := getLocalRange()
+func scanIPRange(
+	ctx context.Context, iprange netip.Prefix, laddr string, proto string, fastscan bool, stealth bool, ports ...int,
+) (RangeScanResult, error) {
 	hosts := createHostRange(iprange)
 
 	var results RangeScanResult
 	for _, h := range hosts {
-		scan, err := scanIPPorts(h, laddr, proto, fastscan, stealth)
+		scan, err := scanIPPorts(ctx, h, laddr, proto, fastscan, stealth, ports...)
 		if err != nil {
 			continue
 		}
+
 		results = append(results, scan)
 	}
 
 	return results, nil
 }
 
 // scanIPPorts scans a list of ports on <hostname> <protocol>
-func scanIPPorts(hostname string, laddr string, proto string, fastscan bool, stealth bool) (*IPScanResult, error) {
-	var results []portResult
-
+func scanIPPorts(
+	ctx context.Context, hostname string, laddr string, proto string, fastscan bool, stealth bool, ports ...int,
+) (*IPScanResult, error) {
 	// checks if device is online
 	addr, err := net.LookupIP(hostname)
 	if err != nil {
@@ -39,67 +44,101 @@ func scanIPPorts(hostname string, laddr string, proto string, fastscan bool, ste
 	// This gets the device name. ('/etc/hostname')
 	// This is typically a good indication of if a host is 'up'
 	// but can cause false-negatives in certain situations.
-	// For this reason when in fastscan mode, devices without
-	// names are ignored but are fully scanned in slowmode.
 	hname, err := net.LookupAddr(hostname)
-	if fastscan {
-		if err != nil {
-			return nil, err
-		}
-	} else if err != nil {
+	if err != nil {
 		hname = append(hname, "Unknown")
 	}
 
+	depth := 4
+	list := make(map[int]string)
+
+	for _, p := range ports {
+		if svc, ok := detailedlist[p]; ok {
+			list[p] = svc
+		} else {
+			list[p] = "unknown"
+		}
+	}
+
+	if len(list) == 0 {
+		if fastscan {
+			list = commonlist
+			depth = 8
+		} else {
+			list = detailedlist
+			depth = 16
+		}
+	}
+
 	// Start prepping channels and vars for worker pool
 	in := make(chan int)
 	go func() {
-		for i := 0; i <= 65535; i++ {
-			in <- i
+		defer close(in)
+
+		for p := range list {
+			select {
+			case <-ctx.Done():
+				return
+			case in <- p:
+			}
 		}
-		close(in)
 	}()
 
-	var list map[int]string
-	var depth int
-
-	if fastscan {
-		list = commonlist
-		depth = 50
-	} else {
-		list = detailedlist
-		depth = 500
-	}
-
 	tasks := len(list)
 
 	// Create results channel and worker function
-	resultChannel := make(chan portResult, tasks)
+	resultChannel := make(chan portResult)
+
+	var wg sync.WaitGroup
 	worker := func() {
-		for port := range in {
-			if service, ok := list[port]; ok {
+		defer wg.Done()
+
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case port, ok := <-in:
+				if !ok {
+					return
+				}
+
+				service := list[port]
+
 				if stealth {
 					scanPortSyn(resultChannel, proto, hostname, service, port, laddr)
 				} else {
 					scanPort(resultChannel, proto, hostname, service, port)
 				}
 			}
+
 		}
 	}
 
 	// Deploy a pool of workers
 	for i := 0; i < depth; i++ {
+		wg.Add(1)
 		go worker()
 	}
 
-	// Combines all results from resultChannel and return a IPScanResult
-	for result := range resultChannel {
-		results = append(results, result)
-		fmt.Printf("\033[2K\rHost: %s | Ports Scanned %d/%d", hostname, len(results), tasks)
+	var results []portResult
+
+	var wgResult sync.WaitGroup
 
-		if len(results) == tasks {
-			close(resultChannel)
+	wgResult.Add(1)
+	go func() {
+		defer wgResult.Done()
+		// Combines all results from resultChannel and return a IPScanResult
+		for result := range resultChannel {
+			results = append(results, result)
+			fmt.Printf("\033[2K\rHost: %s | Ports Scanned %d/%d", hostname, len(results), tasks)
 		}
-	}
+	}()
+
+	wg.Wait()
+
+	close(resultChannel)
+
+	wgResult.Wait()
 
 	return &IPScanResult{
 		Hostname: hname[0],
@@ -114,6 +153,7 @@ func scanIPPorts(hostname string, laddr string, proto string, fastscan bool, ste
 func scanPort(resultChannel chan<- portResult, protocol, hostname, service string, port int) {
 	result := portResult{Port: port, Service: service}
 	address := hostname + ":" + strconv.Itoa(port)
+
 	conn, err := net.DialTimeout(protocol, address, 3*time.Second)
 	if err != nil {
 		result.State = false
@@ -122,6 +162,7 @@ func scanPort(resultChannel chan<- portResult, protocol, hostname, service strin
 	}
 
 	defer conn.Close()
+
 	result.State = true
 	resultChannel <- result
 }

gomap_test.go

@@ -1,6 +1,7 @@
 package gomap_test
 
 import (
+	"context"
 	"fmt"
 	"testing"
 
@@ -15,7 +16,7 @@ func TestMain(m *testing.M) {
 	)
 
 	// results, err := gomap.ScanIP("192.168.1.1", proto, fastscan, stealth)
-	results, err := gomap.ScanRange(proto, fastscan, stealth)
+	results, err := gomap.ScanLocal(context.Background(), proto, fastscan, stealth)
 	if err != nil {
 		panic(err)
 	} else {