From e21274c6e9e826dfcf8e508888f8942be3846eb5 Mon Sep 17 00:00:00 2001
From: lqqyt2423 <974923609@qq.com>
Date: Mon, 13 Feb 2023 15:25:08 +0800
Subject: [PATCH] support ignore_hosts allow_hosts options

---
 Makefile                      |  6 ++++-
 cmd/go-mitmproxy/config.go    | 37 ++++++++++++++++++++++++++
 cmd/go-mitmproxy/main.go      | 50 +++++++++++++++++++++++------------
 cmd/go-mitmproxy/main_test.go | 43 ++++++++++++++++++++++++++++++
 4 files changed, 118 insertions(+), 18 deletions(-)
 create mode 100644 cmd/go-mitmproxy/config.go
 create mode 100644 cmd/go-mitmproxy/main_test.go

diff --git a/Makefile b/Makefile
index 58546dc..2745e33 100644
--- a/Makefile
+++ b/Makefile
@@ -2,7 +2,7 @@ all: mitmproxy
 
 .PHONY: mitmproxy
 mitmproxy:
-	go build -o go-mitmproxy cmd/go-mitmproxy/main.go
+	go build -o go-mitmproxy cmd/go-mitmproxy/*.go
 
 .PHONY: dummycert
 dummycert:
@@ -15,3 +15,7 @@ clean:
 .PHONY: test
 test:
 	go test ./... -v
+
+.PHONY: dev
+dev:
+	go run $(shell ls cmd/go-mitmproxy/*.go | grep -v _test.go)
diff --git a/cmd/go-mitmproxy/config.go b/cmd/go-mitmproxy/config.go
new file mode 100644
index 0000000..0f24ac0
--- /dev/null
+++ b/cmd/go-mitmproxy/config.go
@@ -0,0 +1,37 @@
+package main
+
+import (
+	"flag"
+	"fmt"
+)
+
+func loadConfig() *Config {
+	config := new(Config)
+
+	flag.IntVar(&config.debug, "debug", 0, "debug mode: 1 - print debug log, 2 - show debug from")
+	flag.BoolVar(&config.version, "version", false, "show version")
+	flag.StringVar(&config.addr, "addr", ":9080", "proxy listen addr")
+	flag.StringVar(&config.webAddr, "web_addr", ":9081", "web interface listen addr")
+	flag.BoolVar(&config.ssl_insecure, "ssl_insecure", false, "not verify upstream server SSL/TLS certificates.")
+	flag.StringVar(&config.dump, "dump", "", "dump filename")
+	flag.IntVar(&config.dumpLevel, "dump_level", 0, "dump level: 0 - header, 1 - header + body")
+	flag.StringVar(&config.mapperDir, "mapper_dir", "", "mapper files dirpath")
+	flag.StringVar(&config.certPath, "cert_path", "", "path of generate cert files")
+	flag.Var((*arrayValue)(&config.ignoreHosts), "ignore_hosts", "a list of ignore hosts")
+	flag.Var((*arrayValue)(&config.allowHosts), "allow_hosts", "a list of allow hosts")
+	flag.Parse()
+
+	return config
+}
+
+// arrayValue 实现了 flag.Value 接口
+type arrayValue []string
+
+func (a *arrayValue) String() string {
+	return fmt.Sprint(*a)
+}
+
+func (a *arrayValue) Set(value string) error {
+	*a = append(*a, value)
+	return nil
+}
diff --git a/cmd/go-mitmproxy/main.go b/cmd/go-mitmproxy/main.go
index 233fe65..16efe0c 100644
--- a/cmd/go-mitmproxy/main.go
+++ b/cmd/go-mitmproxy/main.go
@@ -1,10 +1,10 @@
 package main
 
 import (
-	"flag"
 	"fmt"
 	rawLog "log"
 	"os"
+	"strings"
 
 	"github.com/lqqyt2423/go-mitmproxy/addon"
 	"github.com/lqqyt2423/go-mitmproxy/proxy"
@@ -25,23 +25,9 @@ type Config struct {
 	dumpLevel int    // dump level
 
 	mapperDir string
-}
 
-func loadConfig() *Config {
-	config := new(Config)
-
-	flag.IntVar(&config.debug, "debug", 0, "debug mode: 1 - print debug log, 2 - show debug from")
-	flag.BoolVar(&config.version, "version", false, "show version")
-	flag.StringVar(&config.addr, "addr", ":9080", "proxy listen addr")
-	flag.StringVar(&config.webAddr, "web_addr", ":9081", "web interface listen addr")
-	flag.BoolVar(&config.ssl_insecure, "ssl_insecure", false, "not verify upstream server SSL/TLS certificates.")
-	flag.StringVar(&config.dump, "dump", "", "dump filename")
-	flag.IntVar(&config.dumpLevel, "dump_level", 0, "dump level: 0 - header, 1 - header + body")
-	flag.StringVar(&config.mapperDir, "mapper_dir", "", "mapper files dirpath")
-	flag.StringVar(&config.certPath, "cert_path", "", "path of generate cert files")
-	flag.Parse()
-
-	return config
+	ignoreHosts []string
+	allowHosts  []string
 }
 
 func main() {
@@ -81,6 +67,17 @@ func main() {
 
 	log.Infof("go-mitmproxy version %v\n", p.Version)
 
+	if len(config.ignoreHosts) > 0 {
+		p.SetShouldInterceptRule(func(address string) bool {
+			return !matchHost(address, config.ignoreHosts)
+		})
+	}
+	if len(config.allowHosts) > 0 {
+		p.SetShouldInterceptRule(func(address string) bool {
+			return matchHost(address, config.allowHosts)
+		})
+	}
+
 	p.AddAddon(&proxy.LogAddon{})
 	p.AddAddon(web.NewWebAddon(config.webAddr))
 
@@ -96,3 +93,22 @@ func main() {
 
 	log.Fatal(p.Start())
 }
+
+func matchHost(address string, hosts []string) bool {
+	hostname, port := splitHostPort(address)
+	for _, host := range hosts {
+		h, p := splitHostPort(host)
+		if h == hostname && (p == "" || p == port) {
+			return true
+		}
+	}
+	return false
+}
+
+func splitHostPort(address string) (string, string) {
+	index := strings.LastIndex(address, ":")
+	if index == -1 {
+		return address, ""
+	}
+	return address[:index], address[index+1:]
+}
diff --git a/cmd/go-mitmproxy/main_test.go b/cmd/go-mitmproxy/main_test.go
new file mode 100644
index 0000000..56963e2
--- /dev/null
+++ b/cmd/go-mitmproxy/main_test.go
@@ -0,0 +1,43 @@
+package main
+
+import (
+	"testing"
+)
+
+func TestMatchHost(t *testing.T) {
+	address := "www.baidu.com:443"
+	hosts := []string{
+		"www.baidu.com:443",
+		"www.baidu.com",
+		"www.google.com",
+	}
+	expected := true
+	result := matchHost(address, hosts)
+	if result != expected {
+		t.Errorf("Expected %t but got %t", expected, result)
+	}
+
+	address = "www.google.com:80"
+	hosts = []string{
+		"www.baidu.com:443",
+		"www.baidu.com",
+		"www.google.com",
+	}
+	expected = true
+	result = matchHost(address, hosts)
+	if result != expected {
+		t.Errorf("Expected %t but got %t", expected, result)
+	}
+
+	address = "www.test.com:80"
+	hosts = []string{
+		"www.baidu.com:443",
+		"www.baidu.com",
+		"www.google.com",
+	}
+	expected = false
+	result = matchHost(address, hosts)
+	if result != expected {
+		t.Errorf("Expected %t but got %t", expected, result)
+	}
+}