Skip to content

Commit

Permalink
[IMP]optimize the logic of multiple functions
Browse files Browse the repository at this point in the history
  • Loading branch information
@randolphcyg
randolphcyg committed Mar 17, 2023
1 parent 1847d6f commit c8f2bb5
Show file tree
Hide file tree
Showing 5 changed files with 89 additions and 199 deletions.
33 changes: 2 additions & 31 deletions README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -88,23 +88,6 @@ var (
ErrRecRsp = errors.New("Error receiving response")
)

func GetLocalIP() (string, error) {
addrs, err := net.InterfaceAddrs()
if err != nil {
return "", err
}
for _, addr := range addrs {
ipNet, ok := addr.(*net.IPNet)
if ok && !ipNet.IP.IsLoopback() {
if ipNet.IP.To4() != nil {
return ipNet.IP.String(), nil
}
}
}

return "", errors.New("no IPv4 address found")
}

func ServiceDetect(host string, port int, probe parser.Probe) (serviceName string, info parser.VInfo, err error) {
addr := net.JoinHostPort(host, strconv.Itoa(port))
conn, err := net.Dial(strings.ToLower(probe.Protocol), addr)
Expand Down Expand Up @@ -172,19 +155,7 @@ func main() {
panic(err)
}

localIP, err := GetLocalIP()
if err != nil {
panic(err)
}

isLocal := false
var host string
if isLocal {
host = localIP
} else {
host = "127.0.0.1"
}

host := "127.0.0.1"
port := 6379

serviceName := ""
Expand All @@ -203,7 +174,7 @@ func main() {
}

if serviceName != "" && !info.IsVInfoEmpty() {
fmt.Println(serviceName, info) // redis {Redis key-value store 7.0.8 [{ redislabs redis 7.0.8 }]}
fmt.Println(serviceName, info)
} else {
fmt.Println("no match service!")
}
Expand Down
10 changes: 7 additions & 3 deletions helper.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -49,11 +49,15 @@ func HelperI(sign string, b []byte) (val uint32) {
}

// FillHelperFuncOrVariable replace versionInfo helper functions and Variable
func FillHelperFuncOrVariable(str string, src [][]byte) (string, error) {
func FillHelperFuncOrVariable(str string, src [][]byte) string {
if str == "" {
return str
}

for _, p := range patternFlags {
re, err := regexp.Compile(p)
if err != nil {
return str, err
return str
}

matches := re.FindAllString(str, -1)
Expand Down Expand Up @@ -108,5 +112,5 @@ func FillHelperFuncOrVariable(str string, src [][]byte) (string, error) {

}

return str, nil
return str
}
204 changes: 74 additions & 130 deletions parser.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,12 +2,12 @@ package parser

import (
"bufio"
"github.com/pkg/errors"
"os"
"reflect"
"strconv"
"strings"

"github.com/pkg/errors"
"github.com/randolphcyg/cpe"
)

Expand Down Expand Up @@ -50,77 +50,50 @@ func (v VInfo) IsVInfoEmpty() bool {
return reflect.DeepEqual(v, VInfo{})
}

func handleVInfoField(src, flagStr string) (ret string, srcRet string, err error) {
func handleVInfoField(src, flagStr string) (string, string, error) {
src = strings.TrimSpace(src)
srcRet = src
isFlagInSrc := strings.Index(src, flagStr)
if isFlagInSrc != -1 {
end := strings.Index(src[isFlagInSrc+len(flagStr):], "/")
if end == -1 {
return "", src, errors.New("vInfo field end is wrong")
}
ret = src[isFlagInSrc+len(flagStr) : isFlagInSrc+len(flagStr)+end]
if isFlagInSrc+len(flagStr)+end+1 <= len(src) {
srcRet = src[:isFlagInSrc] + src[isFlagInSrc+len(flagStr)+end+1:]
}
if isFlagInSrc == -1 {
return "", src, nil
}

return
}

func HandleVInfo(src string) (vInfo VInfo, err error) {
fieldP, src, err := handleVInfoField(src, "p/")
if err != nil {
return
}
vInfo.VendorProductName = fieldP
if src == "" {
return
end := strings.IndexByte(src[isFlagInSrc+len(flagStr):], '/')
if end == -1 {
return "", src, errors.New("vInfo field end is wrong")
}

fieldV, src, err := handleVInfoField(src, "v/")
if err != nil {
return
}
vInfo.Version = fieldV
if src == "" {
return
}
ret := src[isFlagInSrc+len(flagStr) : isFlagInSrc+len(flagStr)+end]
srcRet := src[:isFlagInSrc+len(flagStr)] + src[isFlagInSrc+len(flagStr)+end+1:]

fieldI, src, err := handleVInfoField(src, "i/")
if err != nil {
return
}
vInfo.Info = fieldI
if src == "" {
return
}
return ret, srcRet, nil
}

fieldH, src, err := handleVInfoField(src, "h/")
if err != nil {
return
}
vInfo.Hostname = fieldH
if src == "" {
return
}
type fieldInfo struct {
field string
set func(string)
}

fieldO, src, err := handleVInfoField(src, "o/")
if err != nil {
return
}
vInfo.OperatingSystem = fieldO
if src == "" {
return
func HandleVInfo(src string) (vInfo VInfo, err error) {
fields := []fieldInfo{
{"p/", func(v string) { vInfo.VendorProductName = v }},
{"v/", func(v string) { vInfo.Version = v }},
{"i/", func(v string) { vInfo.Info = v }},
{"h/", func(v string) { vInfo.Hostname = v }},
{"o/", func(v string) { vInfo.OperatingSystem = v }},
{"d/", func(v string) { vInfo.DeviceType = v }},
}

fieldD, src, err := handleVInfoField(src, "d/")
if err != nil {
return
}
vInfo.DeviceType = fieldD
if src == "" {
return
for _, field := range fields {
fieldValue := ""
fieldValue, src, err = handleVInfoField(src, field.field)
if err != nil {
return
}

field.set(fieldValue)
if src == "" {
return
}
}

// CPE handle logic
Expand All @@ -137,7 +110,7 @@ func HandleVInfo(src string) (vInfo VInfo, err error) {
}
}

cpeSlice := strings.Split(cpeSrcStr, " ")
cpeSlice := strings.SplitN(cpeSrcStr, " ", -1)
if len(cpeSlice) > 0 {
for _, c := range cpeSlice {
cRet, err := cpe.ParseCPE(c)
Expand Down Expand Up @@ -201,28 +174,28 @@ func ParseNmap(srcFilePath string) (probes []Probe, err error) {
}
defer file.Close()

// Create a scanner to read the file line by line
scanner := bufio.NewScanner(file)
probes = make([]Probe, 0, 200)

// Create an empty probe to hold current probe being parsed
var currentProbe Probe

// Loop through each line of the file
for scanner.Scan() {
// Create a scanner to read the file line by line; Loop through each line of the file
for scanner := bufio.NewScanner(file); scanner.Scan(); {
line := scanner.Text()

// Ignore comments and empty lines
if strings.HasPrefix(line, "#") || len(line) == 0 || strings.HasPrefix(line, "Exclude ") {
continue
}

// If the line starts with "Probe", start a new probe
if strings.HasPrefix(line, "Probe ") {
switch {
case strings.HasPrefix(line, "#"), len(line) == 0, strings.HasPrefix(line, "Exclude "):
continue
case strings.HasPrefix(line, "Probe "): // If the line starts with "Probe", start a new probe
// If we have an existing probe, append it to the slice of probes
if currentProbe.ProbeName != "" {
probes = append(probes, currentProbe)
}

// Create a new probe with the name and default values
currentProbe = Probe{
ProbeName: "",
Expand All @@ -243,31 +216,28 @@ func ParseNmap(srcFilePath string) (probes []Probe, err error) {
}
currentProbe.Protocol = lineSeg[1]
currentProbe.ProbeName = lineSeg[2]

probeStringSrc := strings.TrimLeft(lineSeg[3], "q|")
probeString := strings.TrimRight(probeStringSrc, "|")
probeStringSrc := strings.TrimPrefix(lineSeg[3], "q|")
probeString := strings.TrimSuffix(probeStringSrc, "|")
currentProbe.ProbeString = probeString

} else if strings.HasPrefix(line, "match ") || strings.HasPrefix(line, "softmatch ") {
case strings.HasPrefix(line, "match "), strings.HasPrefix(line, "softmatch "):
m, err := ParseMatch(line)
if err != nil {
continue
}
currentProbe.Matches = append(currentProbe.Matches, m)
} else if strings.HasPrefix(line, "ports ") {
case strings.HasPrefix(line, "ports "):
currentProbe.Ports = strings.Split(line[len("ports "):], ",")
} else if strings.HasPrefix(line, "sslports ") {
case strings.HasPrefix(line, "sslports "):
currentProbe.SslPorts = strings.Split(line[len("sslports "):], ",")
} else if strings.HasPrefix(line, "totalwaitms ") {
case strings.HasPrefix(line, "totalwaitms "):
currentProbe.TotalWaitMs = line[len("totalwaitms "):]
} else if strings.HasPrefix(line, "tcpwrappedms ") {
case strings.HasPrefix(line, "tcpwrappedms "):
currentProbe.TcpWrappedMs = line[len("tcpwrappedms "):]
} else if strings.HasPrefix(line, "rarity ") {
case strings.HasPrefix(line, "rarity "):
currentProbe.Rarity = line[len("rarity "):]
} else if strings.HasPrefix(line, "fallback ") {
case strings.HasPrefix(line, "fallback "):
currentProbe.Fallback = line[len("fallback "):]
}

}

// Append the last probe to the slice of probes
Expand All @@ -293,59 +263,33 @@ func UnquoteRawString(rawStr string) (string, error) {

// FillVersionInfoFields Replace the versionInfo and CPE placeholder elements with the matched real values
func FillVersionInfoFields(src [][]byte, match Match) VInfo {
tmpVerInfo := VInfo{}

if match.VersionInfo.Version != "" {
tmpVerInfo.Version, _ = FillHelperFuncOrVariable(match.VersionInfo.Version, src)
}
if match.VersionInfo.Info != "" {
tmpVerInfo.Info, _ = FillHelperFuncOrVariable(match.VersionInfo.Info, src)
}
if match.VersionInfo.Hostname != "" {
tmpVerInfo.Hostname, _ = FillHelperFuncOrVariable(match.VersionInfo.Hostname, src)
}
if match.VersionInfo.DeviceType != "" {
tmpVerInfo.DeviceType, _ = FillHelperFuncOrVariable(match.VersionInfo.DeviceType, src)
}
if match.VersionInfo.VendorProductName != "" {
tmpVerInfo.VendorProductName, _ = FillHelperFuncOrVariable(match.VersionInfo.VendorProductName, src)
}
if match.VersionInfo.OperatingSystem != "" {
tmpVerInfo.OperatingSystem, _ = FillHelperFuncOrVariable(match.VersionInfo.OperatingSystem, src)
versionInfo := match.VersionInfo
tmpVerInfo := VInfo{
VendorProductName: FillHelperFuncOrVariable(versionInfo.VendorProductName, src),
Version: FillHelperFuncOrVariable(versionInfo.Version, src),
Info: FillHelperFuncOrVariable(versionInfo.Info, src),
Hostname: FillHelperFuncOrVariable(versionInfo.Hostname, src),
OperatingSystem: FillHelperFuncOrVariable(versionInfo.OperatingSystem, src),
DeviceType: FillHelperFuncOrVariable(versionInfo.DeviceType, src),
Cpe: nil,
}

for _, c := range match.VersionInfo.Cpe {
tmpCPE := cpe.CPE{}
if len(versionInfo.Cpe) > 0 {
for _, c := range versionInfo.Cpe {
tmpCPE := cpe.CPE{
Version: FillHelperFuncOrVariable(c.Version, src),
Language: FillHelperFuncOrVariable(c.Language, src),
Vendor: FillHelperFuncOrVariable(c.Vendor, src),
Update: FillHelperFuncOrVariable(c.Update, src),
Other: FillHelperFuncOrVariable(c.Other, src),
TargetSw: FillHelperFuncOrVariable(c.TargetSw, src),
SwEdition: FillHelperFuncOrVariable(c.SwEdition, src),
TargetHw: FillHelperFuncOrVariable(c.TargetHw, src),
Product: FillHelperFuncOrVariable(c.Product, src),
}

if c.Version != "" {
tmpCPE.Version, _ = FillHelperFuncOrVariable(c.Version, src)
}
if c.Language != "" {
tmpCPE.Language, _ = FillHelperFuncOrVariable(c.Language, src)
}
if c.Vendor != "" {
tmpCPE.Vendor, _ = FillHelperFuncOrVariable(c.Vendor, src)
tmpVerInfo.Cpe = append(tmpVerInfo.Cpe, tmpCPE)
}
if c.Update != "" {
tmpCPE.Update, _ = FillHelperFuncOrVariable(c.Update, src)
}
if c.Other != "" {
tmpCPE.Other, _ = FillHelperFuncOrVariable(c.Other, src)
}
if c.TargetSw != "" {
tmpCPE.TargetSw, _ = FillHelperFuncOrVariable(c.TargetSw, src)
}
if c.SwEdition != "" {
tmpCPE.SwEdition, _ = FillHelperFuncOrVariable(c.SwEdition, src)
}
if c.TargetHw != "" {
tmpCPE.TargetHw, _ = FillHelperFuncOrVariable(c.TargetHw, src)
}
if c.Product != "" {
tmpCPE.Product, _ = FillHelperFuncOrVariable(c.Product, src)
}

tmpVerInfo.Cpe = append(tmpVerInfo.Cpe, tmpCPE)
}

return tmpVerInfo
Expand Down
Loading

0 comments on commit c8f2bb5

Please sign in to comment.