From 0b56b19134122f7cad1a6b21a004c20a0c2a5fc8 Mon Sep 17 00:00:00 2001 From: yungtravla Date: Wed, 21 Mar 2018 19:47:15 +1000 Subject: [PATCH] Update http_proxy_js_response.go --- modules/http_proxy_js_response.go | 51 ++++++++++++++++++------------- 1 file changed, 30 insertions(+), 21 deletions(-) diff --git a/modules/http_proxy_js_response.go b/modules/http_proxy_js_response.go index 43034158..d195acb7 100644 --- a/modules/http_proxy_js_response.go +++ b/modules/http_proxy_js_response.go @@ -8,14 +8,12 @@ import ( "strings" "github.com/elazarl/goproxy" - - "github.com/bettercap/bettercap/core" ) type JSResponse struct { Status int ContentType string - Headers string + Headers []JSHeader Body string refHash string @@ -25,17 +23,18 @@ type JSResponse struct { func NewJSResponse(res *http.Response) *JSResponse { cType := "" - headers := "" + headers := make([]JSHeader, 0) code := 200 if res != nil { code = res.StatusCode for name, values := range res.Header { for _, value := range values { + headers = append(headers, JSHeader{name, value}) + if name == "Content-Type" { cType = value } - headers += name + ": " + value + "\r\n" } } } @@ -53,7 +52,11 @@ func NewJSResponse(res *http.Response) *JSResponse { } func (j *JSResponse) NewHash() string { - return fmt.Sprintf("%d.%s.%s", j.Status, j.ContentType, j.Headers) + hash := fmt.Sprintf("%d.%s", j.Status, j.ContentType) + for _, h := range j.Headers { + hash += fmt.Sprintf(".%s-%s", h.Name, h.Value) + } + return hash } func (j *JSResponse) UpdateHash() { @@ -73,29 +76,35 @@ func (j *JSResponse) WasModified() bool { return false } -func (j *JSResponse) Header(name, deflt string) string { +func (j *JSResponse) GetHeader(name, deflt string) string { name = strings.ToLower(name) - for _, header := range strings.Split(j.Headers, "\n") { - parts := strings.SplitN(core.Trim(header), ":", 2) - if len(parts) == 2 && strings.ToLower(parts[0]) == name { - return parts[1] + for _, h := range j.Headers { + if name == strings.ToLower(h.Name) { + return h.Value } } return deflt } +func (j *JSResponse) SetHeader(name, value string) { + name = strings.ToLower(name) + found := false + for _, h := range j.Headers { + if name == strings.ToLower(h.Name) { + found = true + h.Value = value + } + } + if found == false { + j.Headers = append(j.Headers, JSHeader{name, value}) + } +} + func (j *JSResponse) ToResponse(req *http.Request) (resp *http.Response) { resp = goproxy.NewResponse(req, j.ContentType, j.Status, j.Body) - if j.Headers != "" { - for _, header := range strings.Split(j.Headers, "\n") { - header = core.Trim(header) - if header == "" { - continue - } - parts := strings.SplitN(header, ":", 2) - if len(parts) == 2 { - resp.Header.Add(parts[0], parts[1]) - } + if len(j.Headers) > 0 { + for _, h := range j.Headers { + resp.Header.Add(h.Name, h.Value) } } return