增加mux hook

Author: piexlMax(奇淼

core/mux/muxServerHTTP/install.go

@@ -0,0 +1,30 @@
+package muxServeHTTP
+
+import (
+	"fmt"
+	"github.com/HXSecurity/DongTai-agent-go/model"
+	"github.com/brahma-adshonor/gohook"
+	"github.com/go-mux/mux"
+)
+
+func init() {
+	model.HookMap["muxServeHTTP"] = new(MuxServeHTTP)
+}
+
+type MuxServeHTTP struct {
+}
+
+func (h *MuxServeHTTP) Hook() {
+	mx := &mux.Router{}
+	err := gohook.HookMethod(mx, "ServeHTTP", MyServer, MyServerTemp)
+	if err != nil {
+		fmt.Println(err, "HttpServeHTTP")
+	} else {
+		fmt.Println("HttpServeHTTP")
+	}
+}
+
+func (h *MuxServeHTTP) UnHook() {
+	mx := &mux.Router{}
+	gohook.UnHookMethod(mx, "ServeHTTP")
+}

core/mux/muxServerHTTP/replacement.go

@@ -0,0 +1,120 @@
+package muxServeHTTP
+
+import (
+	"bufio"
+	"bytes"
+	"encoding/base64"
+	"github.com/HXSecurity/DongTai-agent-go/api"
+	"github.com/HXSecurity/DongTai-agent-go/global"
+	"github.com/HXSecurity/DongTai-agent-go/model/request"
+	"github.com/HXSecurity/DongTai-agent-go/utils"
+	"github.com/go-mux/mux"
+	"net/http"
+	"reflect"
+	"strconv"
+	"strings"
+)
+
+func MyServer(server *mux.Router, w http.ResponseWriter, r *http.Request) {
+	worker, _ := utils.NewWorker(global.AgentId)
+
+	TraceId := global.TraceId + "-" + strconv.Itoa(int(worker.GetId()))
+	global.TargetTraceId = TraceId
+	MyServerTemp(server, w, r)
+	id := utils.CatGoroutineID()
+	go func() {
+		t := reflect.ValueOf(r.Body)
+		var headerBase string
+		body := ""
+		for k, v := range r.Header {
+			headerBase += k + ": " + strings.Join(v, ",") + "\n"
+		}
+		tranceID := TraceId + "." + strconv.Itoa(global.AgentId) + ".0.0.0"
+		headerBase += "dt-traceid:" + tranceID
+		if t.Kind() == reflect.Ptr {
+			buf := t.
+				Elem().
+				FieldByName("src").
+				Elem().Elem().
+				FieldByName("R").
+				Elem().Elem().
+				FieldByName("buf").Bytes()
+			buf = buf[:bytes.IndexByte(buf, 0)]
+			reader := bufio.NewReader(bytes.NewReader(buf))
+			var reqArr []string
+			for {
+				line, _, err := reader.ReadLine()
+				if err != nil {
+					break
+				}
+				reqArr = append(reqArr, string(line))
+			}
+			body = reqArr[len(reqArr)-1]
+		}
+		header := base64.StdEncoding.EncodeToString([]byte(headerBase))
+		scheme := "http"
+		if r.TLS != nil {
+			scheme = "https"
+		}
+		onlyKey := int(worker.GetId())
+
+		HookGroup := &request.UploadReq{
+			Type:     36,
+			InvokeId: onlyKey,
+			Detail: request.Detail{
+				AgentId: global.AgentId,
+				Function: request.Function{
+					Method:        r.Method,
+					Url:           scheme + "://" + r.Host + r.RequestURI,
+					Uri:           r.URL.Path,
+					Protocol:      r.Proto,
+					ClientIp:      r.RemoteAddr,
+					Language:      "GO",
+					ReplayRequest: false,
+					ReqHeader:     header,
+					ReqBody:       body,
+					QueryString:   r.URL.RawQuery,
+					Pool:          []request.Pool{},
+					TraceId:       tranceID,
+				},
+			},
+		}
+		var resBody string
+		var resH string
+		res, ok := global.ResponseMap.Load(id)
+		if ok {
+			global.ResponseMap.Delete(id)
+			resBody = res.(string)
+		}
+		value2, ok2 := global.ResponseHeaderMap.Load(id)
+		if ok2 {
+			global.ResponseHeaderMap.Delete(id)
+			resH = value2.(string)
+		}
+		for k, v := range w.Header() {
+			resH += k + ": " + strings.Join(v, ",") + "\n"
+		}
+		resHeader := base64.StdEncoding.EncodeToString([]byte(resH))
+		HookGroup.Detail.ResHeader = resHeader
+		HookGroup.Detail.ResBody = resBody
+		goroutineIDs := make(map[string]bool)
+		global.PoolTreeMap.Range(func(key, value interface{}) bool {
+			if value.(*request.PoolTree).IsThisBegin(id) {
+				global.PoolTreeMap.Delete(key)
+				value.(*request.PoolTree).FMT(&HookGroup.Detail.Function.Pool, worker, goroutineIDs, HookGroup.Detail.Function.TraceId)
+				return false
+			}
+			return true
+		})
+		api.ReportUpload(*HookGroup)
+		request.RunMapGCbYGoroutineID(goroutineIDs)
+	}()
+	return
+}
+
+func MyServerTemp(server *mux.Router, w http.ResponseWriter, r *http.Request) {
+	for i := 0; i < 100; i++ {
+
+	}
+	return
+}

go.mod

@@ -16,6 +16,7 @@ require (
 	github.com/StackExchange/wmi v1.2.1 // indirect
 	github.com/elazarl/goproxy v0.0.0-20210801061803-8e322dfb79c4 // indirect
 	github.com/go-chi/chi/v5 v5.0.7
+	github.com/go-mux/mux v1.0.0
 	github.com/gorilla/rpc v1.2.0
 	github.com/jinzhu/now v1.1.3 // indirect
 	github.com/json-iterator/go v1.1.11 // indirect
@@ -32,5 +33,3 @@ require (
 	gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b
 	moul.io/http2curl v1.0.0 // indirect
 )
-
-

go.sum

@@ -35,6 +35,8 @@ github.com/gin-gonic/gin v1.7.7 h1:3DoBmSbJbZAWqXJC3SLjAPfutPJJRN1U5pALB7EeTTs=
 github.com/gin-gonic/gin v1.7.7/go.mod h1:axIBovoeJpVj8S3BwE0uPMTeReE4+AfFtqpqaZ1qq1U=
 github.com/go-chi/chi/v5 v5.0.7 h1:rDTPXLDHGATaeHvVlLcR4Qe0zftYethFucbjVQ1PxU8=
 github.com/go-chi/chi/v5 v5.0.7/go.mod h1:DslCQbL2OYiznFReuXYUmQ2hGd1aDpCnlMNITLSKoi8=
+github.com/go-mux/mux v1.0.0 h1:cjme5KED7XU2D1XPqfepZpr/ouRoa8St5/Q1Wn5/K7s=
+github.com/go-mux/mux v1.0.0/go.mod h1:WycrREmM6C8dXy3aKkr8IQ1XwEqXSsQUm1Hyvs8X0Co=
 github.com/go-ole/go-ole v1.2.5 h1:t4MGB5xEDZvXI+0rMjjsfBsD7yAgp/s9ZDkL1JndXwY=
 github.com/go-ole/go-ole v1.2.5/go.mod h1:pprOEPIfldk/42T2oK7lQ4v4JSDwmV0As9GaiUsvbm0=
 github.com/go-playground/assert/v2 v2.0.1 h1:MsBgLAaY856+nPRTKrp3/OZK38U/wa0CcBYNjji3q3A=

hook/mux.go

@@ -0,0 +1,18 @@
+package hook
+
+type Mux struct {
+}
+
+func (h *Mux) GetHook() []string {
+	return []string{
+		"muxServeHTTP",
+	}
+}
+
+func (h *Mux) HookAll() {
+	Hook(h.GetHook())
+}
+
+func (h *Mux) UnHookAll() {
+	UnHook(h.GetHook())
+}

run/mux/hookMux.go

@@ -0,0 +1,13 @@
+package http
+
+import (
+	_ "github.com/HXSecurity/DongTai-agent-go/core/mux/muxServerHTTP"
+	"github.com/HXSecurity/DongTai-agent-go/global"
+	"github.com/HXSecurity/DongTai-agent-go/hook"
+)
+
+func init() {
+	h := new(hook.Mux)
+	global.AllHooks = append(global.AllHooks, h)
+	h.HookAll()
+}