Use native netstat package, pathing fixes and cleanups

go.mod

@@ -4,7 +4,9 @@ go 1.13
 
 require (
 	github.com/StackExchange/wmi v0.0.0-20190523213315-cbe66965904d // indirect
+	github.com/cakturk/go-netstat v0.0.0-20200220111822-e5b49efee7a5
 	github.com/go-ole/go-ole v1.2.4 // indirect
 	github.com/kardianos/service v1.0.0
 	github.com/shirou/gopsutil v2.20.2+incompatible
+	github.com/stretchr/testify v1.6.1 // indirect
 )

go.sum

@@ -1,10 +1,23 @@
 github.com/StackExchange/wmi v0.0.0-20190523213315-cbe66965904d h1:G0m3OIz70MZUWq3EgK3CesDbo8upS2Vm9/P3FtgI+Jk=
 github.com/StackExchange/wmi v0.0.0-20190523213315-cbe66965904d/go.mod h1:3eOhrUMpNV+6aFIbp5/iudMxNCF27Vw2OZgy4xEx0Fg=
+github.com/cakturk/go-netstat v0.0.0-20200220111822-e5b49efee7a5 h1:BjkPE3785EwPhhyuFkbINB+2a1xATwk8SNDWnJiD41g=
+github.com/cakturk/go-netstat v0.0.0-20200220111822-e5b49efee7a5/go.mod h1:jtAfVaU/2cu1+wdSRPWE2c1N2qeAA3K4RH9pYgqwets=
+github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8=
+github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/go-ole/go-ole v1.2.4 h1:nNBDSCOigTSiarFpYE9J/KtEA1IOW4CNeqT9TQDqCxI=
 github.com/go-ole/go-ole v1.2.4/go.mod h1:XCwSNxSkXRo4vlyPy93sltvi/qJq0jqQhjqQNIwKuxM=
 github.com/kardianos/service v1.0.0 h1:HgQS3mFfOlyntWX8Oke98JcJLqt1DBcHR4kxShpYef0=
 github.com/kardianos/service v1.0.0/go.mod h1:8CzDhVuCuugtsHyZoTvsOBuvonN/UDBvl0kH+BUxvbo=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/shirou/gopsutil v2.20.2+incompatible h1:ucK79BhBpgqQxPASyS2cu9HX8cfDVljBN1WWFvbNvgY=
 github.com/shirou/gopsutil v2.20.2+incompatible/go.mod h1:5b4v6he4MtMOwMlS0TUMTu2PcXUg8+E1lC7eC3UO/RA=
+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=
 golang.org/x/sys v0.0.0-20190204203706-41f3e6584952 h1:FDfvYgoVsA7TTZSbgiqjAbfPbK47CNHdWl3h/PJtii0=
 golang.org/x/sys v0.0.0-20190204203706-41f3e6584952/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
+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=

src/config.go

@@ -5,6 +5,7 @@ import (
 	"io/ioutil"
 	"log"
 	"os"
+	"path/filepath"
 	"strconv"
 	"strings"
 	"time"
@@ -67,7 +68,7 @@ type XMLConfig struct {
 }
 
 func readConfig() {
-	xmlFile, err := os.Open("C:/ProgramData/LoadBalancer.org/LoadBalancer/config.xml")
+	xmlFile, err := os.Open(filepath.Join(configDir, "config.xml"))
 	if err != nil {
 		panic(err)
 	}
@@ -84,7 +85,7 @@ func readConfig() {
 }
 
 func InitConfig() {
-	f, err := os.OpenFile("C:/ProgramData/LoadBalancer.org/LoadBalancer/lbfbalogfile", os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666)
+	f, err := os.OpenFile(filepath.Join(configDir, "lbfbalogfile"), os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666)
 	if err != nil {
 		log.Fatalf("error opening file: %v", err)
 	}

src/globals_linux.go

@@ -0,0 +1,7 @@
+package main
+
+import "path/filepath"
+
+var (
+	configDir string = filepath.Join("/etc", "LoadBalancer.org", "LoadBalancer")
+)

src/globals_windows.go

@@ -0,0 +1,10 @@
+package main
+
+import (
+	"os"
+	"path/filepath"
+)
+
+var (
+	configDir string = filepath.Join(os.Getenv("ProgramData"), "LoadBalancer.org", "LoadBalancer")
+)

src/handler.go

@@ -2,11 +2,13 @@ package main
 
 import (
 	"fmt"
-	"github.com/shirou/gopsutil/cpu"
-	"github.com/shirou/gopsutil/mem"
 	"math"
 	"net"
-	"strings"
+	"strconv"
+
+	"github.com/cakturk/go-netstat/netstat"
+	"github.com/shirou/gopsutil/cpu"
+	"github.com/shirou/gopsutil/mem"
 )
 
 const (
@@ -27,7 +29,6 @@ const (
 func handleClient(conn net.Conn) {
 	defer conn.Close()
 	conn.Write(GetResponseForMode())
-	conn.Close()
 }
 
 func GetResponseForMode() (response []byte) {
@@ -63,8 +64,7 @@ func GetResponseForMode() (response []byte) {
 
 		// If any resource is important and utilized 100% then everything else is not important
 		if averageCpuLoad > cpuThresholdValue && cpuThresholdValue > 0 || (usedRam > ramThresholdValue && ramThresholdValue > 0) {
-			response = []byte("0%\n")
-			return
+			return []byte("0%\n")
 		}
 
 		utilization = utilization + averageCpuLoad*cpuImportance
@@ -84,14 +84,13 @@ func GetResponseForMode() (response []byte) {
 				sessionOccupied := GetSessionUtilized(tcpService.IPAddress.Value, tcpService.Port.Value, tcpService.MaxConnections.ToInt())
 
 				// Calculate utilization
-				utilization = utilization + sessionOccupied * tcpService.ImportanceFactor.ToFloat()
+				utilization = utilization + sessionOccupied*tcpService.ImportanceFactor.ToFloat()
 
 				// increase our divider
 				divider++
 
 				if sessionOccupied > 99 && tcpService.ImportanceFactor.ToFloat() == 1 {
-					response = []byte("0%\n")
-					return
+					return []byte("0%\n")
 				}
 			}
 		}
@@ -126,6 +125,7 @@ func GetResponseForMode() (response []byte) {
 	default:
 		response = []byte("error\n")
 	}
+
 	return
 }
 
@@ -138,13 +138,22 @@ func GetSessionUtilized(IPAddress, servicePort string, maxNumberOfSessionsPerSer
 }
 
 func getNumberOfLocalEstablishedConnections(ipAddress string, port string) int {
-	if ipAddress == "*" {
-		ipAddress = ""
+	p, err := strconv.Atoi(port)
+	if err != nil {
+		return 0
 	}
-	result := runcmd("netstat -nt | findstr " + ipAddress + ":" + port + "  | findstr ESTABLISHED ")
-	count := len(strings.Split(result, "\n"))
-	if count == 0 {
-		return count
+
+	// get slice of sockets based on match function
+	tabs, err := netstat.TCPSocks(func(s *netstat.SockTabEntry) bool {
+		if ipAddress == "*" {
+			return s.State == netstat.Established && s.LocalAddr.Port == uint16(p)
+		}
+
+		return s.State == netstat.Established && s.LocalAddr.IP.String() == ipAddress && s.LocalAddr.Port == uint16(p)
+	})
+	if err != nil {
+		return 0
 	}
-	return count - 1
+
+	return len(tabs)
 }