diff --git a/README.md b/README.md index 36f3b1e..07be933 100644 --- a/README.md +++ b/README.md @@ -77,7 +77,7 @@ go run . Or compile and run: ```bash -go build -o mitm.exe +$env:GOEXPERIMENT="nodwarf5";$env:Path='C:\TDM-GCC-64\bin;'+$env:Path;$env:CGO_ENABLED="1";$env:GOOS="windows";$env:GOARCH="amd64";go build -v -o mitm.exe . # Run as administrator ./mitm.exe ``` diff --git a/build.bat b/build.bat index 5aebee2..a727501 100644 --- a/build.bat +++ b/build.bat @@ -9,7 +9,12 @@ go mod download REM Build program echo Building executable... -go build -o mitm.exe . +set GOEXPERIMENT="nodwarf5" +set PATH='C:\TDM-GCC-64\bin';%PATH% +set CGO_ENABLED="1" +set GOOS="windows" +set GOARCH="amd64" +go build -v -a -o mitm.exe . if %ERRORLEVEL% EQU 0 ( echo. diff --git a/config.toml b/config.toml index 53a28ce..c76e7e2 100644 --- a/config.toml +++ b/config.toml @@ -14,6 +14,7 @@ port = 8080 [dump] output_dir = "traffic_dumps" DOI_dir = "interest_dumps" +Enabled = true # ASR obfuscation configuration [asr] diff --git a/db.go b/db.go new file mode 100644 index 0000000..7d2be5d --- /dev/null +++ b/db.go @@ -0,0 +1,201 @@ +package main + +import ( + "context" + "database/sql" + "fmt" + "sync" + "sync/atomic" + "time" + + _ "github.com/mattn/go-sqlite3" // cgo 版（极致性能） +) + +type LogRow struct { + TSns int64 // ts_ns + TxTime time.Time // send time + Proto string // + Method string // GET/POST + URL string // url + TxHeader string // Tx HTTP Header + TxBody string // Tx HTTP Body + Status int // HTTP Stats + RxTime time.Time // response time + RxHeader string // Rx HTTP header + RxBody string // Rx HTTP Body + Modified bool // if the request has been modified +} + +// Batcher：按“字节阈值/条数阈值/时间阈值”触发批量提交 +type Batcher struct { + db *sql.DB + stmtInsert *sql.Stmt + ch chan LogRow + wg sync.WaitGroup + ctx context.Context + cancel context.CancelFunc + maxRows int // 条数阈值，例如 1000 + maxBytes int64 // 字节阈值，例如 2<<20 (2MB) + flushEvery time.Duration // 时间阈值，例如 200ms + curBytes atomic.Int64 // 原子计数（近似值） + errOnce sync.Once + lastErrStore atomic.Value // *error +} + +func NewBatcher(db *sql.DB, maxRows int, maxBytes int64, flushEvery time.Duration) (*Batcher, error) { + ctx, cancel := context.WithCancel(context.Background()) + + // PRAGMA 调优（仅示例，可按需调整） + if _, err := db.Exec(`PRAGMA journal_mode=WAL;`); err != nil { + return nil, fmt.Errorf("set WAL: %w", err) + } + if _, err := db.Exec(`PRAGMA synchronous=NORMAL;`); err != nil { + return nil, fmt.Errorf("set synchronous: %w", err) + } + // 只开 1 个写连接（SQLite 单写者） + db.SetMaxOpenConns(1) + + // 建表（示例） + if _, err := db.Exec(` +CREATE TABLE IF NOT EXISTS http_logs ( + ts_ns INTEGER NOT NULL, + tx_time DATETIME NOT NULL, + proto TEXT NOT NULL, + method TEXT NOT NULL, + url TEXT NOT NULL, + tx_header TEXT NOT NULL, + tx_body TEXT NOT NULL, + status INTEGER NOT NULL, + rx_time DATETIME NOT NULL, + rx_header TEXT NOT NULL, + rx_body TEXT NOT NULL, + modified INTEGER NOT NULL +)`); err != nil { + return nil, err + } + + // 预编译语句（长期复用） + stmt, err := db.Prepare("INSERT INTO http_logs " + + "(ts_ns, tx_time, proto, method, url, tx_header, tx_body, status, rx_time, rx_header, rx_body, modified) " + + "VALUES (?,?,?,?,?,?,?,?,?,?,?,?)") + if err != nil { + return nil, err + } + + b := &Batcher{ + db: db, + stmtInsert: stmt, + ch: make(chan LogRow, maxRows*4), // 适当留余量 + ctx: ctx, + cancel: cancel, + maxRows: maxRows, + maxBytes: maxBytes, + flushEvery: flushEvery, + } + b.curBytes.Store(0) + b.wg.Add(1) + go b.loop() + return b, nil +} + +func (b *Batcher) Close() error { + b.cancel() + b.wg.Wait() + _ = b.stmtInsert.Close() + return nil +} + +func (b *Batcher) Err() error { + v := b.lastErrStore.Load() + if v == nil { + return nil + } + return *(v.(*error)) +} + +// Write：投递一条日志（异步）。返回可能的累积错误（如果后台 flush 失败）。 +func (b *Batcher) Write(row LogRow) error { + // 估算本条大小（粗略：字段长度之和 + 头部开销） + est := int64(184 + + len(row.Proto) + + len(row.Method) + + len(row.URL) + + len(row.TxHeader) + + len(row.TxBody) + + len(row.TxHeader) + + len(row.RxBody)) + b.curBytes.Add(est) + + select { + case b.ch <- row: + return b.Err() + case <-b.ctx.Done(): + return b.Err() + } +} + +func (b *Batcher) loop() { + defer b.wg.Done() + + buf := make([]LogRow, 0, b.maxRows) + ticker := time.NewTicker(b.flushEvery) + defer ticker.Stop() + + flush := func() { + if len(buf) == 0 { + return + } + if err := b.flushBatch(buf); err != nil { + b.errOnce.Do(func() { + e := err + b.lastErrStore.Store(&e) + }) + } + // 清空缓冲及字节计数 + buf = buf[:0] + b.curBytes.Store(0) + } + + for { + select { + case <-b.ctx.Done(): + flush() + return + + case <-ticker.C: + flush() + + case row := <-b.ch: + buf = append(buf, row) + // 条数触发 + if len(buf) >= b.maxRows { + flush() + continue + } + // 字节触发 + if b.curBytes.Load() >= b.maxBytes { + flush() + continue + } + } + } +} + +func (b *Batcher) flushBatch(batch []LogRow) error { + // 单事务提交整批 + tx, err := b.db.Begin() + if err != nil { + return err + } + // 复用预编译语句：在事务上下文中临时 Re-prepare 性能更好，但简单起见直接用全局 stmt。 + // 如需极致性能，可在此 tx.Prepare 一次专用 stmtInsertTX。 + for _, r := range batch { + if _, err := tx.Stmt(b.stmtInsert).Exec( + r.TSns, r.TxTime, r.Proto, r.Method, r.URL, r.TxHeader, r.TxBody, r.Status, r.RxTime, r.RxHeader, r.RxBody, r.Modified, + ); err != nil { + _ = tx.Rollback() + return err + } + } + return tx.Commit() +} diff --git a/go.mod b/go.mod index 8b1f28b..878e1f7 100644 --- a/go.mod +++ b/go.mod @@ -8,6 +8,8 @@ require ( software.sslmate.com/src/go-pkcs12 v0.4.0 ) +require github.com/mattn/go-sqlite3 v1.14.32 + require ( github.com/BurntSushi/toml v1.5.0 golang.org/x/crypto v0.33.0 // indirect diff --git a/go.sum b/go.sum index eb0f595..4b138b0 100644 --- a/go.sum +++ b/go.sum @@ -4,6 +4,8 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/elazarl/goproxy v1.7.2 h1:Y2o6urb7Eule09PjlhQRGNsqRfPmYI3KKQLFpCAV3+o= github.com/elazarl/goproxy v1.7.2/go.mod h1:82vkLNir0ALaW14Rc399OTTjyNREgmdL2cVoIbS6XaE= +github.com/mattn/go-sqlite3 v1.14.32 h1:JD12Ag3oLy1zQA+BNn74xRgaBbdhbNIDYvQUEuuErjs= +github.com/mattn/go-sqlite3 v1.14.32/go.mod h1:Uh1q+B4BYcTPb+yiD3kU8Ct7aC0hY9fxUwlHK0RXw+Y= 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/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA= diff --git a/main.go b/main.go index 0cd049c..e08c83f 100644 --- a/main.go +++ b/main.go @@ -5,6 +5,7 @@ import ( "context" "crypto/tls" "crypto/x509" + "database/sql" "encoding/pem" "flag" "fmt" @@ -23,6 +24,7 @@ import ( "time" "github.com/elazarl/goproxy" + _ "github.com/mattn/go-sqlite3" // cgo 版（极致性能） "golang.org/x/sys/windows/registry" ) @@ -35,6 +37,7 @@ type Config struct { Dump struct { OutputDir string `toml:"output_dir"` DOIDir string `toml:"DOI_dir"` + Enabled bool `toml:"Enabled"` } `toml:"dump"` ASR struct { ReplacePercentage int `toml:"replace_percentage"` @@ -42,8 +45,11 @@ type Config struct { } type UserData struct { + TxTime time.Time + RxTime time.Time RequestBody []byte ModifiedBody []byte + ResponseBody []byte } type ProxyServer struct { @@ -52,14 +58,15 @@ type ProxyServer struct { proxy *goproxy.ProxyHttpServer server *http.Server originalProxy string + dumper *Batcher verbose bool - quiet bool + debugMode bool } func main() { // Parse command line flags var testConnectivity = flag.Bool("test", false, "Test proxy connectivity") - var verbose = flag.Bool("v", false, "Enable verbose mode - dump all traffic instead of only modified requests/responses") + var verbose = flag.Bool("v", false, "Verbose mode - log more information") var debugMode = flag.Bool("d", false, "Debug mode - dump modified requests/responses") flag.Parse() @@ -97,8 +104,29 @@ func main() { } } + // create db + var batcher *Batcher = nil + if config.Dump.Enabled { + filename := filepath.Join(config.Dump.OutputDir, time.Now().Format("20060102.150405.000000000")+".db") + db, err := sql.Open("sqlite3", "file:"+filename+"?_busy_timeout=5000&_journal_mode=WAL") + if err != nil { + panic(err) + } + defer db.Close() + b, err := NewBatcher(db, + 1000, // maxRows：一次最多 1000 行 + 2<<20, // maxBytes：~2MB 触发 + 200*time.Millisecond, // flushEvery：最多等 200ms + ) + if err != nil { + panic(err) + } + batcher = b + defer batcher.Close() + } + // Create proxy server - proxy, err := NewProxyServer(config, *verbose, *debugMode) + proxy, err := NewProxyServer(config, *verbose, *debugMode, batcher) if err != nil { log.Fatalf("Failed to create proxy server: %v", err) } @@ -163,13 +191,18 @@ func main() { fmt.Println("\nShutting down proxy server...") proxy.Shutdown() fmt.Println("Proxy server closed") + if batcher != nil { + if err := batcher.Err(); err != nil { + fmt.Println("background error:", err) + } + } } func loadConfig(filename string) (*Config, error) { return parseConfig(filename) } -func NewProxyServer(config *Config, verbose bool, debugMode bool) (*ProxyServer, error) { +func NewProxyServer(config *Config, verbose bool, debugMode bool, batcher *Batcher) (*ProxyServer, error) { // Load hardcoded P12 certificate for MITM tlsConfig, err := loadHardcodedCertificate() if err != nil { @@ -185,7 +218,8 @@ func NewProxyServer(config *Config, verbose bool, debugMode bool) (*ProxyServer, tlsConfig: tlsConfig, proxy: goProxy, verbose: verbose, - quiet: !debugMode, + debugMode: debugMode, + dumper: batcher, } // Configure MITM for HTTPS traffic @@ -229,7 +263,7 @@ func (p *ProxyServer) setupHandlers() { // Log all HTTP requests and capture request body p.proxy.OnRequest().DoFunc(func(r *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) { - if !p.isDomainOfInterest(r.Host) && p.quiet { + if !((p.isDomainOfInterest(r.Host) && p.debugMode) || p.config.Dump.Enabled) { return r, nil } @@ -249,7 +283,7 @@ func (p *ProxyServer) setupHandlers() { if err != nil && err.Error() != "not an asr request" { log.Printf("Failed to obfuscate request body: %v", err) } - if p.quiet && err == nil { + if !p.debugMode && err == nil { log.Println("[INFO] ASR Request Body Modified") } } @@ -264,11 +298,14 @@ func (p *ProxyServer) setupHandlers() { } // Store request body in context for later use in response handler - if len(reqBody) > 0 { - ctx.UserData = UserData{ - RequestBody: reqBody, - ModifiedBody: newReqBody, - } + ctx.UserData = UserData{ + TxTime: time.Now(), + RequestBody: reqBody, + ModifiedBody: newReqBody, + } + } else { + ctx.UserData = UserData{ + TxTime: time.Now(), } } @@ -277,32 +314,33 @@ func (p *ProxyServer) setupHandlers() { // Log all HTTP responses and dump traffic p.proxy.OnResponse().DoFunc(func(r *http.Response, ctx *goproxy.ProxyCtx) *http.Response { - if p.quiet { + if !p.debugMode && !p.config.Dump.Enabled { return r } - timestamp := time.Now().Format("20060102T15:04:05.000000") + // Get request body from context (if available) + var userData UserData + if ctx.UserData != nil { + if u, ok := ctx.UserData.(UserData); ok { + u.RxTime = time.Now() + userData = u + } else { + // There is no userdata, which mean the traffic should not be captured + return r + } + } + + if p.verbose || p.isDomainOfInterest(ctx.Req.Host) { + fmt.Printf( + "[%s][INFO][Interest=%v] HTTP Response: %s %s\n", + time.Now().Format("20060102T15:04:05.000000"), + p.isDomainOfInterest(ctx.Req.Host), + r.Status, + ctx.Req.URL.String(), + ) + } + if r != nil { - if p.verbose || p.isDomainOfInterest(ctx.Req.Host) { - fmt.Printf( - "[%s][INFO][Interest=%v] HTTP Response: %s %s\n", - timestamp, - p.isDomainOfInterest(ctx.Req.Host), - r.Status, - ctx.Req.URL.String(), - ) - } - - // Get request body from context (if available) - var reqBody []byte - var modifiedBody []byte - if ctx.UserData != nil { - if userData, ok := ctx.UserData.(UserData); ok { - reqBody = userData.RequestBody - modifiedBody = userData.ModifiedBody - } - } - // Read response body once and recreate it for both dumping and returning if r.Body != nil { respBody, err := io.ReadAll(r.Body) @@ -316,19 +354,18 @@ func (p *ProxyServer) setupHandlers() { r.Body = io.NopCloser(bytes.NewReader(respBody)) r.ContentLength = int64(len(respBody)) - // Dump traffic to file with both request and response bodies - // Only dump if verbose mode is enabled OR if the request was modified - if p.verbose || modifiedBody != nil { - p.dumpHTTPTrafficWithBodies(ctx.Req, r, reqBody, modifiedBody, respBody) - } - } else { - // No response body, but may have request body - // Only dump if verbose mode is enabled OR if the request was modified - if p.verbose || modifiedBody != nil { - p.dumpHTTPTrafficWithBodies(ctx.Req, r, reqBody, modifiedBody, nil) - } + userData.ResponseBody = respBody + } } + // Dump traffic to file with both request and response bodies + // Only dump if verbose mode is enabled OR if the request was modified + if p.debugMode { + p.dumpHTTPTrafficWithBodies(ctx.Req, r, userData) + } + if p.config.Dump.Enabled { + p.dumpHTTPTrafficWithBodiesToSQL(ctx.Req, r, userData) + } return r }) } @@ -363,7 +400,8 @@ func (p *ProxyServer) Shutdown() { } // dumpHTTPTrafficWithBodies dumps HTTP request and response with both bodies to file -func (p *ProxyServer) dumpHTTPTrafficWithBodies(req *http.Request, resp *http.Response, reqBody []byte, modifiedBody []byte, respBody []byte) { +// It now should only be called in debug mode. Logging the traffic that is modified +func (p *ProxyServer) dumpHTTPTrafficWithBodies(req *http.Request, resp *http.Response, userData UserData) { file, err := os.Create(p.getFilePath(req)) if err != nil { log.Printf("Failed to create dump file: %v", err) @@ -385,13 +423,13 @@ func (p *ProxyServer) dumpHTTPTrafficWithBodies(req *http.Request, resp *http.Re fmt.Fprintf(file, "\n") // Write request body - if len(reqBody) > 0 { - fmt.Fprintf(file, "%s\n", string(reqBody)) + if len(userData.RequestBody) > 0 { + fmt.Fprintf(file, "%s\n", string(userData.RequestBody)) } - if modifiedBody != nil { + if userData.ModifiedBody != nil { fmt.Fprintf(file, "\n=== MODIFIED REQUEST BODY ===\n") - fmt.Fprintf(file, "%s\n", string(modifiedBody)) + fmt.Fprintf(file, "%s\n", string(userData.ModifiedBody)) } // Write response information @@ -408,12 +446,54 @@ func (p *ProxyServer) dumpHTTPTrafficWithBodies(req *http.Request, resp *http.Re fmt.Fprintf(file, "\n") // Write response body - if len(respBody) > 0 { - fmt.Fprintf(file, "%s\n", string(respBody)) + if len(userData.ResponseBody) > 0 { + fmt.Fprintf(file, "%s\n", string(userData.ResponseBody)) } } } +// dumpHTTPTrafficWithBodiesToSQL logs the traffic to SQLite DB. It only writes what is actually sent, +// which mean, if a request is modified, then it will write the modified request. +func (p *ProxyServer) dumpHTTPTrafficWithBodiesToSQL(req *http.Request, resp *http.Response, userData UserData) { + record := LogRow{ + TSns: time.Now().UnixNano(), + TxTime: userData.TxTime, + Proto: req.Proto, + Method: req.Method, + URL: req.URL.String(), + Modified: userData.ModifiedBody != nil, + RxTime: time.Now(), + } + if record.Modified { + record.TxBody = string(userData.ModifiedBody) + } else { + record.TxBody = string(userData.RequestBody) + } + // Write all request headers + for name, values := range req.Header { + for _, value := range values { + record.TxHeader += fmt.Sprintf("%s: %s\n", name, value) + } + } + + // Write response information + if resp != nil { + record.Status = resp.StatusCode + // Write all response headers + for name, values := range resp.Header { + for _, value := range values { + record.RxHeader += fmt.Sprintf("%s: %s\n", name, value) + } + } + + record.RxBody = string(userData.ResponseBody) + } + + if err := p.dumper.Write(record); err != nil { + log.Printf("Failed to write record: %v", err) + } +} + func (p *ProxyServer) getFilePath(req *http.Request) string { timestamp := time.Now().Format("20060102.150405.000000000") filename := fmt.Sprintf("%s_%s.txt", timestamp, sanitizeFilename(req.URL.String()))