From 3e25111a77925393c8d3488b80a84f7ddc64be8b Mon Sep 17 00:00:00 2001 From: shadowabi <50265741+shadowabi@users.noreply.github.com> Date: Tue, 30 Apr 2024 13:40:48 +0800 Subject: [PATCH] optimize --- define/var.go | 13 ------------- pkg/data.go | 22 +++------------------- pkg/req.go | 36 +++++++++++++++++++++++------------- 3 files changed, 26 insertions(+), 45 deletions(-) diff --git a/define/var.go b/define/var.go index 8a1ca8e..4003532 100644 --- a/define/var.go +++ b/define/var.go @@ -1,7 +1,5 @@ package define -import "net" - type Configure struct { Mode string `mapstructure:"Mode" json:"Mode" yaml:"Mode"` CloudEndpoint string `mapstructure:"CloudEndpoint" json:"CloudEndpoint" yaml:"CloudEndpoint"` @@ -16,14 +14,3 @@ var ( TimeOut int Port string ) - -type IPRecord struct { - Domain string - Ip []net.IP -} - -type Record struct { - Ip []net.IP - SvcDomain string - SrvRecords []*net.SRV -} diff --git a/pkg/data.go b/pkg/data.go index df168b5..02ce8a7 100644 --- a/pkg/data.go +++ b/pkg/data.go @@ -2,13 +2,12 @@ package pkg import ( "bufio" - "fmt" - "github.com/wgpsec/EndpointSearch/define" - "github.com/wgpsec/EndpointSearch/internal/config" - "github.com/wgpsec/EndpointSearch/utils/Error" "os" "regexp" "strings" + + "github.com/wgpsec/EndpointSearch/internal/config" + "github.com/wgpsec/EndpointSearch/utils/Error" ) func ParseFileParameter(fileName string) (fileHostList []string) { @@ -52,21 +51,6 @@ func ConvertToReqList(endpoint string, param ...string) (reqList []string) { return reqList } -func ParseRecordResult(recordList ...define.Record) (resultList []string) { - if len(recordList) != 0 { - for _, record := range recordList { - if len(record.SrvRecords) != 0 { - for _, srv := range record.SrvRecords { - result := strings.Join([]string{srv.Target, ":", fmt.Sprintf("%v", srv.Port)}, "") - resultList = append(resultList, result) - } - } - resultList = append(resultList, record.SvcDomain) - } - } - return resultList -} - func JudgeEndpoint(respList ...ResponseData) (resultList []string) { if len(respList) != 0 { for _, resp := range respList { diff --git a/pkg/req.go b/pkg/req.go index 7915f6d..acdfd8e 100644 --- a/pkg/req.go +++ b/pkg/req.go @@ -2,7 +2,6 @@ package pkg import ( "fmt" - "github.com/wgpsec/EndpointSearch/define" "net" "net/http" "strings" @@ -14,43 +13,54 @@ type ResponseData struct { contentType string } -func SearchDomain(domainList ...string) (ipRecordList []define.IPRecord) { +type IPRecord struct { + domain string + ip []net.IP +} + +type Record struct { + ip []net.IP + svcDomain string + srvRecords []*net.SRV +} + +func SearchDomain(domainList ...string) (ipRecordList []IPRecord) { if len(domainList) != 0 { var answers []net.IP for _, domain := range domainList { answers, _ = net.LookupIP(domain) if len(answers) != 0 { - ipRecordList = append(ipRecordList, define.IPRecord{Domain: domain, Ip: answers}) + ipRecordList = append(ipRecordList, IPRecord{domain: domain, ip: answers}) } } } return ipRecordList } -func SearchSRVRecord(ipRecordList ...define.IPRecord) (recordList []define.Record) { +func SearchSRVRecord(ipRecordList ...IPRecord) (recordList []Record) { if len(ipRecordList) != 0 { for _, ipRecord := range ipRecordList { - _, srv, err := net.LookupSRV(ipRecord.Domain, "tcp", "") + _, srv, err := net.LookupSRV(ipRecord.domain, "tcp", "") if err != nil { - recordList = append(recordList, define.Record{Ip: ipRecord.Ip, SvcDomain: ipRecord.Domain}) + recordList = append(recordList, Record{ip: ipRecord.ip, svcDomain: ipRecord.domain}) } else { - recordList = append(recordList, define.Record{Ip: ipRecord.Ip, SvcDomain: ipRecord.Domain, SrvRecords: srv}) + recordList = append(recordList, Record{ip: ipRecord.ip, svcDomain: ipRecord.domain, srvRecords: srv}) } } } return recordList } -func SearchEndpoint(client *http.Client, portList []string, recordList ...define.Record) (respList []ResponseData) { +func SearchEndpoint(client *http.Client, portList []string, recordList ...Record) (respList []ResponseData) { if len(recordList) != 0 { resultsChan := make(chan ResponseData, cap(recordList)) var wg1 sync.WaitGroup var wg2 sync.WaitGroup for _, record := range recordList { for _, port := range portList { - if len(record.SrvRecords) != 0 { - resultsChan2 := make(chan ResponseData, cap(record.SrvRecords)) - for _, srv := range record.SrvRecords { + if len(record.srvRecords) != 0 { + resultsChan2 := make(chan ResponseData, cap(record.srvRecords)) + for _, srv := range record.srvRecords { go func(srv *net.SRV, port string, wg *sync.WaitGroup) { defer wg.Done() //logrus.Debug(srv.Target, ":", port, " Working") @@ -66,9 +76,9 @@ func SearchEndpoint(client *http.Client, portList []string, recordList ...define close(resultsChan2) } wg1.Add(1) - go func(record define.Record, port string, wg *sync.WaitGroup) { + go func(record Record, port string, wg *sync.WaitGroup) { //logrus.Debug(record.SvcDomain, ":", port, " Working") - requestStr := strings.Join([]string{record.SvcDomain, ":", port}, "") + requestStr := strings.Join([]string{record.svcDomain, ":", port}, "") SendHttpRequest(client, requestStr, resultsChan, wg) //logrus.Debug(record.SvcDomain, ":", port, " Done") }(record, port, &wg1)