Skip to content

Commit

Permalink
feat 使用协程池
Browse files Browse the repository at this point in the history
  • Loading branch information
@yhy0
yhy0 committed May 7, 2024
1 parent 8c9679d commit d975ba3
Show file tree
Hide file tree
Showing 7 changed files with 338 additions and 281 deletions.
10 changes: 6 additions & 4 deletions pkg/mitmproxy/go-mitmproxy.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,11 +7,11 @@ package mitmproxy
**/

import (
"github.com/panjf2000/ants/v2"
"github.com/yhy0/Jie/conf"
"github.com/yhy0/Jie/pkg/mitmproxy/go-mitmproxy/proxy"
"github.com/yhy0/Jie/pkg/task"
"github.com/yhy0/logging"
"github.com/yhy0/sizedwaitgroup"
)

var t *task.Task
Expand All @@ -33,10 +33,12 @@ func NewMitmproxy() {
ScanTask: make(map[string]*task.ScanTask),
}

t.Wg = sizedwaitgroup.New(t.Parallelism)
pool, _ := ants.NewPool(t.Parallelism)
t.Pool = pool
defer t.Pool.Release() // 释放协程池

// 先加一，这里会一直阻塞，这样就不会马上退出, 这里要的就是一直阻塞，所以不使用 wg.Done()
t.Wg.Add()
t.WG.Add(1)

var err error
PassiveProxy, err = proxy.NewProxy(opts)
Expand All @@ -53,5 +55,5 @@ func NewMitmproxy() {
}
}()

t.Wg.Wait()
t.WG.Wait()
}
4 changes: 2 additions & 2 deletions pkg/mitmproxy/passive.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@ func judge(f *proxy.Flow) {
flag = funk.Contains(conf.GlobalConfig.Mitmproxy.FilterSuffix, ext)
}
if !flag {
go distribution(f)
distribution(f)
}
}
} else {
Expand All @@ -55,7 +55,7 @@ func judge(f *proxy.Flow) {
flag = funk.Contains(conf.GlobalConfig.Mitmproxy.FilterSuffix, ext)
}
if !flag {
go distribution(f)
distribution(f)
}
}
}
24 changes: 16 additions & 8 deletions pkg/mitmproxy/task.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ import (
"github.com/yhy0/logging"
"net/url"
"strings"

"github.com/yhy0/Jie/pkg/input"
"github.com/yhy0/Jie/pkg/protocols/httpx"
"strconv"
Expand All @@ -25,7 +25,7 @@ func distribution(f *proxy.Flow) {
logging.Logger.Errorln(err)
return
}

var host string
// 有的会带80、443端口号，导致 example.com 和 example.com:80、example.com:443被认为是不同的网站
port := strings.Split(parseUrl.Host, ":")
Expand All @@ -34,14 +34,14 @@ func distribution(f *proxy.Flow) {
} else {
host = parseUrl.Host
}

// 使用解码后的，不然有的 js f.Response.Body 直接乱码
var body []byte
body, err = f.Response.DecodedBody()
if err != nil {
body = f.Response.Body
}

// TODO 将 http.Header 转换为 map[string]string 有的重复请求头，这里后面遇到了再优化吧
headerMap := make(map[string]string)
for key, values := range f.Request.Header {
Expand All @@ -50,11 +50,11 @@ func distribution(f *proxy.Flow) {
if key == "Set-Cookie" {
separator = ";"
}

// 将多个值连接成一个字符串，用逗号分隔
headerMap[key] = strings.Join(values, separator)
}

in := &input.CrawlResult{
Target: f.Request.URL.Host,
Url: f.Request.URL.String(),
Expand All @@ -73,6 +73,14 @@ func distribution(f *proxy.Flow) {
RawRequest: requestDump(f.Request),
RawResponse: responseDump(f),
}

t.Distribution(in)

t.WG.Add(1)
go func() {
err := t.Pool.Submit(t.Distribution(in))
if err != nil {
t.WG.Done()
logging.Logger.Errorf("add distribution err:%v, crawlResult:%v", err, in)
}
}()
// t.Distribution(in)
}
37 changes: 25 additions & 12 deletions pkg/mode/active.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,9 +3,9 @@ package mode
import (
"encoding/json"
"fmt"
"github.com/panjf2000/ants/v2"
"github.com/projectdiscovery/katana/pkg/output"
"github.com/thoas/go-funk"
regexp "github.com/wasilibs/go-re2"
"github.com/yhy0/Jie/conf"
"github.com/yhy0/Jie/crawler"
"github.com/yhy0/Jie/crawler/crawlergo"
Expand All @@ -22,6 +22,7 @@ import (
"github.com/yhy0/sizedwaitgroup"
"net/url"
"path"
"regexp"
"strings"
"time"
)
Expand Down Expand Up @@ -84,8 +85,9 @@ func Active(target string, fingerprint []string) ([]string, []string) {
Wg: sizedwaitgroup.New(3 + 3),
}

t.Wg = sizedwaitgroup.New(t.Parallelism)

pool, _ := ants.NewPool(t.Parallelism)
t.Pool = pool
defer t.Pool.Release() // 释放协程池
// 爬虫前，进行连接性、指纹识别、 waf 探测
resp, err := client.Request(target, "GET", "", nil)
if err != nil {
Expand All @@ -107,7 +109,7 @@ func Active(target string, fingerprint []string) ([]string, []string) {
subdomains = Katana(target, wafs, t, fingerprint)
}

t.Wg.Wait()
t.WG.Wait()

logging.Logger.Debugln("Fingerprints: ", t.Fingerprints)

Expand All @@ -117,8 +119,6 @@ func Active(target string, fingerprint []string) ([]string, []string) {
}

func Katana(target string, waf []string, t *task.Task, fingerprint []string) []string {
t.Wg.Add()
defer t.Wg.Done()
parseUrl, err := url.Parse(target)
if err != nil {
logging.Logger.Errorln(err)
Expand All @@ -132,7 +132,7 @@ func Katana(target string, waf []string, t *task.Task, fingerprint []string) []s
curl := strings.ReplaceAll(result.Request.URL, "\\n", "")
curl = strings.ReplaceAll(curl, "\\t", "")
curl = strings.ReplaceAll(curl, "\\n", "")
parseUrl, err := url.Parse(curl)
parseUrl, err = url.Parse(curl)
if err != nil {
logging.Logger.Errorln(err)
return
Expand All @@ -158,7 +158,8 @@ func Katana(target string, waf []string, t *task.Task, fingerprint []string) []s
} else {
host = parseUrl.Host
}
resp, err := t.ScanTask[host].Client.Request(result.Request.URL, result.Request.Method, result.Request.Body, result.Request.Headers)
var resp *httpx.Response
resp, err = t.ScanTask[host].Client.Request(result.Request.URL, result.Request.Method, result.Request.Body, result.Request.Headers)
if err != nil {
logging.Logger.Errorln(err)
return
Expand Down Expand Up @@ -200,7 +201,14 @@ func Katana(target string, waf []string, t *task.Task, fingerprint []string) []s
}

// 分发扫描任务
go t.Distribution(crawlResult)
t.WG.Add(1)
go func() {
err := t.Pool.Submit(t.Distribution(crawlResult))
if err != nil {
t.WG.Done()
logging.Logger.Errorf("add distribution err:%v, crawlResult:%v", err, crawlResult)
}
}()
}

if conf.GlobalConfig.WebScan.Craw == "k" {
Expand All @@ -216,8 +224,6 @@ func Katana(target string, waf []string, t *task.Task, fingerprint []string) []s

// Crawlergo 运行爬虫, 对爬虫结果进行处理
func Crawlergo(target string, waf []string, t *task.Task, fingerprint []string) []string {
t.Wg.Add()
defer t.Wg.Done()
var targets []*model.Request

var req model.Request
Expand Down Expand Up @@ -296,7 +302,14 @@ func Crawlergo(target string, waf []string, t *task.Task, fingerprint []string)
}

// 分发扫描任务
go t.Distribution(crawlResult)
t.WG.Add(1)
go func() {
err := t.Pool.Submit(t.Distribution(crawlResult))
if err != nil {
t.WG.Done()
logging.Logger.Errorf("add distribution err:%v, crawlResult:%v", err, crawlResult)
}
}()
}

// 开始爬虫任务
Expand Down
44 changes: 34 additions & 10 deletions pkg/output/SCopilot.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,13 @@ var DataUpdated = make(chan struct{})

// SCopilot 将数据存储到 SCopilotMessage 中
func SCopilot(host string, data SCopilotData) {
_host := strings.Split(host, ":")
if len(_host) > 1 {
if _host[1] == "80" {
host = _host[0]
}
}

lock.Lock()
defer lock.Unlock()
// 判断 map 中是否存在，存在的话就 append，不存在的话就创建
Expand All @@ -34,38 +41,55 @@ func SCopilot(host string, data SCopilotData) {
sort.SliceStable(SCopilotMessage[host].SiteMap, func(i, j int) bool {
return compareLinks(SCopilotMessage[host].SiteMap[i], SCopilotMessage[host].SiteMap[j])
})

SCopilotMessage[host].Fingerprints = funk.UniqString(append(SCopilotMessage[host].Fingerprints, data.Fingerprints...))

for _, v := range data.VulMessage {
if funk.Contains(SCopilotMessage[host].VulMessage, v) {
continue
}
SCopilotMessage[host].VulMessage = append(SCopilotMessage[host].VulMessage, v)

if SCopilotMessage[host].VulPlugin == nil {
SCopilotMessage[host].VulPlugin = make(map[string]int)
}
if SCopilotMessage[host].VulPlugin[v.Plugin] > 0 {
SCopilotMessage[host].VulPlugin[v.Plugin] = SCopilotMessage[host].VulPlugin[v.Plugin] + 1
} else {
SCopilotMessage[host].VulPlugin[v.Plugin] = 1
}
}

for _, v := range data.InfoMsg {
if funk.Contains(SCopilotMessage[host].InfoMsg, v) {
continue
}
SCopilotMessage[host].InfoMsg = append(SCopilotMessage[host].InfoMsg, v)
if SCopilotMessage[host].InfoPlugin == nil {
SCopilotMessage[host].InfoPlugin = make(map[string]int)
}
if SCopilotMessage[host].InfoPlugin[v.Plugin] > 0 {
SCopilotMessage[host].InfoPlugin[v.Plugin] = SCopilotMessage[host].InfoPlugin[v.Plugin] + 1
} else {
SCopilotMessage[host].InfoPlugin[v.Plugin] = 1
}
}

for _, v := range data.PluginMsg {
if funk.Contains(SCopilotMessage[host].PluginMsg, v) {
continue
}
SCopilotMessage[host].PluginMsg = append(SCopilotMessage[host].PluginMsg, v)
}

for _, v := range SCopilotLists {
if v.Host == host {
v.InfoCount = len(SCopilotMessage[host].InfoMsg)
v.ApiCount = len(SCopilotMessage[host].SiteMap)
v.VulnCount = len(SCopilotMessage[host].VulMessage)
}
}

SCopilotMessage[host].CollectionMsg.Subdomain = funk.UniqString(append(SCopilotMessage[host].CollectionMsg.Subdomain, data.CollectionMsg.Subdomain...))
SCopilotMessage[host].CollectionMsg.OtherDomain = funk.UniqString(append(SCopilotMessage[host].CollectionMsg.OtherDomain, data.CollectionMsg.OtherDomain...))
SCopilotMessage[host].CollectionMsg.PublicIp = funk.UniqString(append(SCopilotMessage[host].CollectionMsg.PublicIp, data.CollectionMsg.PublicIp...))
Expand All @@ -75,14 +99,14 @@ func SCopilot(host string, data SCopilotData) {
SCopilotMessage[host].CollectionMsg.Others = funk.UniqString(append(SCopilotMessage[host].CollectionMsg.Others, data.CollectionMsg.Others...))
SCopilotMessage[host].CollectionMsg.Urls = funk.UniqString(append(SCopilotMessage[host].CollectionMsg.Urls, data.CollectionMsg.Urls...))
SCopilotMessage[host].CollectionMsg.Api = funk.UniqString(append(SCopilotMessage[host].CollectionMsg.Api, data.CollectionMsg.Api...))

} else {
SCopilotMessage[host] = &data
SCopilotLists = append(SCopilotLists, &SCopilotList{
Host: host,
})
}

// 通知数据已更新 这样防止没有启动前端界面时，造成阻塞
select {
case DataUpdated <- struct{}{}:
Expand All @@ -102,12 +126,12 @@ func compareLinks(a, b string) bool {
}
aPathComponents := strings.Split(aURL.Path, "/")
bPathComponents := strings.Split(bURL.Path, "/")

for i := 0; i < len(aPathComponents) && i < len(bPathComponents); i++ {
if aPathComponents[i] != bPathComponents[i] {
return aPathComponents[i] < bPathComponents[i]
}
}

return len(aPathComponents) < len(bPathComponents)
}
20 changes: 11 additions & 9 deletions pkg/output/type.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,15 +37,17 @@ type VulnData struct {
}

type SCopilotData struct {
Target string `json:"target"`
Ip string `json:"ip"`
HostNoPort string `json:"host_no_port"`
SiteMap []string `json:"site_map"`
Fingerprints []string `json:"fingerprints"`
VulMessage []VulMessage `json:"vul_message"`
InfoMsg []PluginMsg `json:"info_msg"`
PluginMsg []PluginMsg `json:"plugin_msg"`
CollectionMsg Collection `json:"collection_msg"`
Target string `json:"target"`
Ip string `json:"ip"`
HostNoPort string `json:"host_no_port"`
SiteMap []string `json:"site_map"`
Fingerprints []string `json:"fingerprints"`
VulMessage []VulMessage `json:"vul_message"`
VulPlugin map[string]int `json:"vul_plugin"`
InfoMsg []PluginMsg `json:"info_msg"`
InfoPlugin map[string]int `json:"info_plugin"`
PluginMsg []PluginMsg `json:"plugin_msg"`
CollectionMsg Collection `json:"collection_msg"`
}

type Collection struct {
Expand Down
Loading

0 comments on commit d975ba3

Please sign in to comment.