From 8150254b763cbf8f75847eb5df4793c7becd8eea Mon Sep 17 00:00:00 2001 From: lqqyt2423 <974923609@qq.com> Date: Fri, 3 Dec 2021 16:07:32 +0800 Subject: [PATCH] add flowmapper addon --- addon/flowmapper/mapper.go | 96 +++++++++++++++ addon/flowmapper/parser.go | 206 ++++++++++++++++++++++++++++++++ addon/flowmapper/parser_test.go | 41 +++++++ cmd/go-mitmproxy/main.go | 23 +++- 4 files changed, 360 insertions(+), 6 deletions(-) create mode 100644 addon/flowmapper/mapper.go create mode 100644 addon/flowmapper/parser.go create mode 100644 addon/flowmapper/parser_test.go diff --git a/addon/flowmapper/mapper.go b/addon/flowmapper/mapper.go new file mode 100644 index 0000000..46146d9 --- /dev/null +++ b/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 +} diff --git a/addon/flowmapper/parser.go b/addon/flowmapper/parser.go new file mode 100644 index 0000000..d8a9939 --- /dev/null +++ b/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 +} diff --git a/addon/flowmapper/parser_test.go b/addon/flowmapper/parser_test.go new file mode 100644 index 0000000..623971f --- /dev/null +++ b/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") + } +} diff --git a/cmd/go-mitmproxy/main.go b/cmd/go-mitmproxy/main.go index 541072d..b1ca80e 100644 --- a/cmd/go-mitmproxy/main.go +++ b/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 + 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()) }