kyoketsu/pkg/scanner_test.go

package kyoketsu

import (
	"fmt"
	"log"
	"net"
	"sync"
	"testing"
)

func startTestSever(port int, t *testing.T) {
	listen, err := net.Listen("tcp", fmt.Sprintf("localhost:%v", port))
	if err != nil {
		log.Fatal(err)
	}
	defer listen.Close()
	for {
		conn, err := listen.Accept()
		if err != nil {
			log.Fatal(err)
		}
		t.Logf("Recieved dial from PortWalk of port: %v from host: %s\n", port, conn.RemoteAddr().String())
		conn.Close()
	}
}

func TestPortWalk(t *testing.T) {
	type TestCase struct {
		Name       string
		ScanPort   map[int]string
		ListenPort map[int]string
		ShouldFail bool
	}
	tc := []TestCase{
		TestCase{
			Name:       "Passing test, listening on 8080/27017 and scanned on 8080/27017.",
			ScanPort:   map[int]string{8080: "unknown", 27017: "mongo"},
			ListenPort: map[int]string{8080: "unknown", 27017: "mongo"},
			ShouldFail: false,
		},
		TestCase{
			Name:       "Failing test, listening on 8081 scanned on 6439",
			ScanPort:   map[int]string{6439: "random_service"},
			ListenPort: map[int]string{8081: "unknown"},
			ShouldFail: true,
		},
	}
	wg := &sync.WaitGroup{}
	for i := range tc {
		for port, _ := range tc[i].ListenPort {
			wg.Add(1)
			go startTestSever(port, t)
		}
		got := PortWalk("localhost", tc[i].ScanPort)
		for k, _ := range tc[i].ListenPort {
			wg.Done()

			_, ok := got.ListeningPorts[k]
			if !ok {
				if !tc[i].ShouldFail {
					t.Errorf("Test '%s' failed! PortWalk didnt detect the test server was listening on: %+v\n", tc[i].Name, tc[i].ListenPort)
				}
			}
			t.Logf("Test '%s' passed! Scanned port: '%v'", tc[i].Name, k)

		}

	}
}