add flowmapper addon

addon/flowmapper/mapper.go

@@ -0,0 +1,96 @@
+package flowmapper
+
+import (
+	"io/ioutil"
+	"path/filepath"
+	"regexp"
+	"strings"
+
+	"github.com/lqqyt2423/go-mitmproxy/addon"
+	"github.com/lqqyt2423/go-mitmproxy/flow"
+	_log "github.com/sirupsen/logrus"
+)
+
+var log = _log.WithField("at", "changeflow addon")
+var httpsRegexp = regexp.MustCompile(`^https://`)
+
+type Mapper struct {
+	addon.Base
+	reqResMap map[string]*flow.Response
+}
+
+func NewMapper(dirname string) *Mapper {
+	infos, err := ioutil.ReadDir(dirname)
+	if err != nil {
+		panic(err)
+	}
+
+	filenames := make([]string, 0)
+
+	for _, info := range infos {
+		if info.IsDir() {
+			continue
+		}
+		if !strings.HasSuffix(info.Name(), ".map.txt") {
+			continue
+		}
+
+		filenames = append(filenames, filepath.Join(dirname, info.Name()))
+	}
+
+	if len(filenames) == 0 {
+		return &Mapper{
+			reqResMap: make(map[string]*flow.Response),
+		}
+	}
+
+	ch := make(chan interface{}, len(filenames))
+	for _, filename := range filenames {
+		go func(filename string, ch chan<- interface{}) {
+			f, err := ParseFlowFromFile(filename)
+			if err != nil {
+				log.Error(err)
+				ch <- err
+				return
+			}
+			ch <- f
+		}(filename, ch)
+	}
+
+	reqResMap := make(map[string]*flow.Response)
+
+	for i := 0; i < len(filenames); i++ {
+		flowOrErr := <-ch
+		if f, ok := flowOrErr.(*flow.Flow); ok {
+			key := buildReqKey(f.Request)
+			log.Infof("add request mapper: %v", key)
+			reqResMap[key] = f.Response
+		}
+	}
+
+	return &Mapper{
+		reqResMap: reqResMap,
+	}
+}
+
+func ParseFlowFromFile(filename string) (*flow.Flow, error) {
+	p, err := NewParserFromFile(filename)
+	if err != nil {
+		return nil, err
+	}
+	return p.Parse()
+}
+
+func (c *Mapper) Request(f *flow.Flow) {
+	key := buildReqKey(f.Request)
+	if resp, ok := c.reqResMap[key]; ok {
+		f.Response = resp
+	}
+}
+
+func buildReqKey(req *flow.Request) string {
+	url := req.URL.String()
+	url = httpsRegexp.ReplaceAllString(url, "http://")
+	key := req.Method + " " + url
+	return key
+}

addon/flowmapper/parser.go

@@ -0,0 +1,206 @@
+package flowmapper
+
+import (
+	"errors"
+	"io/ioutil"
+	"net/http"
+	"net/url"
+	"regexp"
+	"strconv"
+	"strings"
+
+	"github.com/lqqyt2423/go-mitmproxy/flow"
+)
+
+type Parser struct {
+	lines    []string
+	url      string
+	request  *flow.Request
+	response *flow.Response
+}
+
+func NewParserFromFile(filename string) (*Parser, error) {
+	bytes, err := ioutil.ReadFile(filename)
+	if err != nil {
+		return nil, err
+	}
+
+	return NewParserFromString(string(bytes))
+}
+
+func NewParserFromString(content string) (*Parser, error) {
+	content = strings.TrimSpace(content)
+	lines := strings.Split(content, "\n")
+	if len(lines) == 0 {
+		return nil, errors.New("no lines")
+	}
+
+	return &Parser{
+		lines: lines,
+	}, nil
+}
+
+func (p *Parser) Parse() (*flow.Flow, error) {
+	if err := p.parseRequest(); err != nil {
+		return nil, err
+	}
+
+	if err := p.parseResponse(); err != nil {
+		return nil, err
+	}
+
+	return &flow.Flow{
+		Request:  p.request,
+		Response: p.response,
+	}, nil
+}
+
+func (p *Parser) parseRequest() error {
+	if err := p.parseReqHead(); err != nil {
+		return err
+	}
+
+	if header, err := p.parseHeader(); err != nil {
+		return err
+	} else {
+		p.request.Header = header
+	}
+
+	// parse url
+	if !strings.HasPrefix(p.url, "http") {
+		host := p.request.Header.Get("host")
+		if host == "" {
+			return errors.New("no request host")
+		}
+		p.url = "http://" + host + p.url
+	}
+	url, err := url.Parse(p.url)
+	if err != nil {
+		return err
+	}
+	p.request.URL = url
+
+	p.parseReqBody()
+
+	return nil
+}
+
+func (p *Parser) parseReqHead() error {
+	line, _ := p.getLine()
+	re := regexp.MustCompile(`^(GET|POST|PUT|DELETE)\s+?(.+)`)
+	matches := re.FindStringSubmatch(line)
+	if len(matches) == 0 {
+		return errors.New("request head parse error")
+	}
+
+	p.request = &flow.Request{
+		Method: matches[1],
+	}
+	p.url = matches[2]
+
+	return nil
+}
+
+func (p *Parser) parseHeader() (http.Header, error) {
+	header := make(http.Header)
+	re := regexp.MustCompile(`^([\w-]+?):\s*(.+)$`)
+
+	for {
+		line, ok := p.getLine()
+		if !ok {
+			break
+		}
+		line = strings.TrimSpace(line)
+		if line == "" {
+			break
+		}
+		matches := re.FindStringSubmatch(line)
+		if len(matches) == 0 {
+			return nil, errors.New("request header parse error")
+		}
+
+		key := matches[1]
+		val := matches[2]
+		header.Add(key, val)
+	}
+
+	return header, nil
+}
+
+func (p *Parser) parseReqBody() {
+	bodyLines := make([]string, 0)
+
+	for {
+		line, ok := p.getLine()
+		if !ok {
+			break
+		}
+
+		if len(bodyLines) == 0 {
+			line = strings.TrimSpace(line)
+			if line == "" {
+				continue
+			}
+		}
+
+		if strings.HasPrefix(line, "HTTP/1.1 ") {
+			p.lines = append([]string{line}, p.lines...)
+			break
+		}
+		bodyLines = append(bodyLines, line)
+	}
+
+	body := strings.Join(bodyLines, "\n")
+	body = strings.TrimSpace(body)
+	p.request.Body = []byte(body)
+}
+
+func (p *Parser) parseResponse() error {
+	if err := p.parseResHead(); err != nil {
+		return err
+	}
+
+	if header, err := p.parseHeader(); err != nil {
+		return err
+	} else {
+		p.response.Header = header
+	}
+
+	// all left content
+	body := strings.Join(p.lines, "\n")
+	body = strings.TrimSpace(body)
+	p.response.Body = []byte(body)
+	p.response.Header.Set("Content-Length", strconv.Itoa(len(p.response.Body)))
+
+	return nil
+}
+
+func (p *Parser) parseResHead() error {
+	line, ok := p.getLine()
+	if !ok {
+		return errors.New("response no head line")
+	}
+
+	re := regexp.MustCompile(`^HTTP/1\.1\s+?(\d+)`)
+	matches := re.FindStringSubmatch(line)
+	if len(matches) == 0 {
+		return errors.New("response head parse error")
+	}
+
+	code, _ := strconv.Atoi(matches[1])
+	p.response = &flow.Response{
+		StatusCode: code,
+	}
+
+	return nil
+}
+
+func (p *Parser) getLine() (string, bool) {
+	if len(p.lines) == 0 {
+		return "", false
+	}
+
+	line := p.lines[0]
+	p.lines = p.lines[1:]
+	return line, true
+}

addon/flowmapper/parser_test.go

@@ -0,0 +1,41 @@
+package flowmapper
+
+import "testing"
+
+func TestParser(t *testing.T) {
+	content := `
+GET /index.html
+Host: www.baidu.com
+Accept: */*
+
+hello world
+
+HTTP/1.1 200
+
+ok
+`
+	p, err := NewParserFromString(content)
+	if err != nil {
+		t.Fatal(err)
+	}
+	f, err := p.Parse()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if f.Request.Method != "GET" {
+		t.Fatal("request method error")
+	}
+	if f.Request.URL.String() != "http://www.baidu.com/index.html" {
+		t.Fatal("request url error")
+	}
+	if f.Response.StatusCode != 200 {
+		t.Fatal("response status code error")
+	}
+	if string(f.Response.Body) != "ok" {
+		t.Fatal("response body error")
+	}
+	if f.Response.Header.Get("Content-Length") != "2" {
+		t.Fatal("response header content-length error")
+	}
+}

cmd/go-mitmproxy/main.go

@@ -6,19 +6,24 @@ import (
 	"os"
 
 	"github.com/lqqyt2423/go-mitmproxy/addon"
+	"github.com/lqqyt2423/go-mitmproxy/addon/flowmapper"
 	"github.com/lqqyt2423/go-mitmproxy/addon/web"
 	"github.com/lqqyt2423/go-mitmproxy/proxy"
 	log "github.com/sirupsen/logrus"
 )
 
-const version = "0.1.0"
+const version = "0.1.1"
 
 type Config struct {
 	version bool
+
 	addr    string
 	webAddr string
+
 	dump      string // dump filename
 	dumpLevel int    // dump level
+
+	mapperDir string
 }
 
 func loadConfig() *Config {
@@ -29,6 +34,7 @@ func loadConfig() *Config {
 	flag.StringVar(&config.webAddr, "web_addr", ":9081", "web interface listen addr")
 	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.Parse()
 
 	return config
@@ -61,13 +67,18 @@ func main() {
 		log.Fatal(err)
 	}
 
+	p.AddAddon(&addon.Log{})
+	p.AddAddon(web.NewWebAddon(config.webAddr))
+
 	if config.dump != "" {
 		dumper := addon.NewDumper(config.dump, config.dumpLevel)
 		p.AddAddon(dumper)
 	}
 
-	p.AddAddon(&addon.Log{})
-	p.AddAddon(web.NewWebAddon(config.webAddr))
+	if config.mapperDir != "" {
+		mapper := flowmapper.NewMapper(config.mapperDir)
+		p.AddAddon(mapper)
+	}
 
 	log.Fatal(p.Start())
 }