diff --git a/http_modifier.go b/http_modifier.go index 5023008..d820df6 100644 --- a/http_modifier.go +++ b/http_modifier.go @@ -1,6 +1,7 @@ package main import ( + "bytes" "github.com/buger/gor/proto" "hash/fnv" ) @@ -17,6 +18,7 @@ func NewHTTPModifier(config *HTTPModifierConfig) *HTTPModifier { len(config.headerFilters) == 0 && len(config.headerHashFilters) == 0 && len(config.paramHashFilters) == 0 && + len(config.params) == 0 && len(config.headers) == 0 && len(config.methods) == 0 { return nil @@ -26,8 +28,33 @@ func NewHTTPModifier(config *HTTPModifierConfig) *HTTPModifier { } func (m *HTTPModifier) Rewrite(payload []byte) (response []byte) { - if len(m.config.methods) > 0 && !m.config.methods.Contains(proto.Method(payload)) { - return + if len(m.config.methods) > 0 { + method := proto.Method(payload) + + matched := false + + for _, m := range m.config.methods { + if bytes.Equal(method, m) { + matched = true + break + } + } + + if !matched { + return + } + } + + if len(m.config.headers) > 0 { + for _, header := range m.config.headers { + payload = proto.SetHeader(payload, []byte(header.Name), []byte(header.Value)) + } + } + + if len(m.config.params) > 0 { + for _, param := range m.config.params { + payload = proto.SetPathParam(payload, param.Name, param.Value) + } } if len(m.config.urlRegexp) > 0 { @@ -38,6 +65,7 @@ func (m *HTTPModifier) Rewrite(payload []byte) (response []byte) { for _, f := range m.config.urlRegexp { if f.regexp.Match(path) { matched = true + break } } @@ -109,11 +137,5 @@ func (m *HTTPModifier) Rewrite(payload []byte) (response []byte) { } } - if len(m.config.headers) > 0 { - for _, header := range m.config.headers { - payload = proto.SetHeader(payload, []byte(header.Name), []byte(header.Value)) - } - } - return payload } diff --git a/http_modifier_settings.go b/http_modifier_settings.go index cfbc82e..854b11a 100644 --- a/http_modifier_settings.go +++ b/http_modifier_settings.go @@ -1,7 +1,6 @@ package main import ( - "bytes" "errors" "fmt" "regexp" @@ -17,6 +16,7 @@ type HTTPModifierConfig struct { headerHashFilters HTTPHashFilters paramHashFilters HTTPHashFilters + params HTTPParams headers HTTPHeaders methods HTTPMethods } @@ -121,6 +121,34 @@ func (h *HTTPHeaders) Set(value string) error { return nil } +// +// Handling of --http-set-param option +// +type HTTPParams []HTTPParam +type HTTPParam struct { + Name []byte + Value []byte +} + +func (h *HTTPParams) String() string { + return fmt.Sprint(*h) +} + +func (h *HTTPParams) Set(value string) error { + v := strings.SplitN(value, "=", 2) + if len(v) != 2 { + return errors.New("Expected `Key=Value`") + } + + param := HTTPParam{ + []byte(strings.TrimSpace(v[0])), + []byte(strings.TrimSpace(v[1])), + } + + *h = append(*h, param) + return nil +} + // // Handling of --http-allow-method option // @@ -135,15 +163,6 @@ func (h *HTTPMethods) Set(value string) error { return nil } -func (h *HTTPMethods) Contains(value []byte) bool { - for _, method := range *h { - if bytes.Equal(value, method) { - return true - } - } - return false -} - // // Handling of --http-rewrite-url option // diff --git a/http_modifier_settings_test.go b/http_modifier_settings_test.go index 5e135b6..60f382a 100644 --- a/http_modifier_settings_test.go +++ b/http_modifier_settings_test.go @@ -51,21 +51,6 @@ func TestHTTPHashFilters(t *testing.T) { } } -func TestHTTPMethods(t *testing.T) { - methods := HTTPMethods{} - - methods.Set("GET") - methods.Set("POST") - - if !methods.Contains([]byte("GET")) { - t.Error("Does not contain GET") - } - - if !methods.Contains([]byte("POST")) { - t.Error("Does not contain POST") - } -} - func TestUrlRewriteMap(t *testing.T) { var err error rewrites := UrlRewriteMap{} diff --git a/http_modifier_test.go b/http_modifier_test.go index 55d9a19..40e5ec4 100644 --- a/http_modifier_test.go +++ b/http_modifier_test.go @@ -186,3 +186,35 @@ func TestHTTPModifierURLNegativeRegexp(t *testing.T) { t.Error("Should not pass url") } } + +func TestHTTPModifierSetHeader(t *testing.T) { + filters := HTTPHeaders{} + filters.Set("User-Agent:Gor") + + modifier := NewHTTPModifier(&HTTPModifierConfig{ + headers: filters, + }) + + payload := []byte("POST /post HTTP/1.1\r\nContent-Length: 7\r\nHost: www.w3.org\r\n\r\na=1&b=2") + payload_after := []byte("POST /post HTTP/1.1\r\nUser-Agent: Gor\r\nContent-Length: 7\r\nHost: www.w3.org\r\n\r\na=1&b=2") + + if payload = modifier.Rewrite(payload); !bytes.Equal(payload_after, payload) { + t.Error("Should add new header", string(payload)) + } +} + +func TestHTTPModifierSetParam(t *testing.T) { + filters := HTTPParams{} + filters.Set("api_key=1") + + modifier := NewHTTPModifier(&HTTPModifierConfig{ + params: filters, + }) + + payload := []byte("POST /post?api_key=1234 HTTP/1.1\r\nContent-Length: 7\r\nHost: www.w3.org\r\n\r\na=1&b=2") + payload_after := []byte("POST /post?api_key=1 HTTP/1.1\r\nContent-Length: 7\r\nHost: www.w3.org\r\n\r\na=1&b=2") + + if payload = modifier.Rewrite(payload); !bytes.Equal(payload_after, payload) { + t.Error("Should override param", string(payload)) + } +} diff --git a/proto/proto.go b/proto/proto.go index 9b0e5cf..626d289 100644 --- a/proto/proto.go +++ b/proto/proto.go @@ -104,6 +104,38 @@ func PathParam(payload, name []byte) (value []byte, valueStart, valueEnd int) { } } +func SetPathParam(payload, name, value []byte) []byte { + path := Path(payload) + _, vs, ve := PathParam(payload, name) + + if vs != -1 { + newPath := make([]byte, len(path)) + copy(newPath, path) + newPath = byteutils.Replace(newPath, vs, ve, value) + + return SetPath(payload, newPath) + } else { // if param not found append to end of url + // Adding 2 because of '?' or '&' at start, and '=' in middle + newParam := make([]byte, len(name) + len(value) + 2) + + if bytes.IndexByte(path, '?') == -1 { + newParam[0] = '?' + } else { + newParam[0] = '&' + } + + copy(newParam[1:], name) + newParam[1+len(name)] = '=' + copy(newParam[2+len(name):], value) + + newPath := make([]byte, len(path) + len(newParam)) + copy(newPath, path) + copy(newPath[len(path):], newParam) + + return SetPath(payload, newPath) + } +} + func SetHost(payload, url, host []byte) []byte { // If this is HTTP 1.0 traffic or proxy traffic it may include host right into path variable, so instead of setting Host header we rewrite Path // Fix for https://github.com/buger/gor/issues/156 diff --git a/proto/proto_test.go b/proto/proto_test.go index 12227e2..080f023 100644 --- a/proto/proto_test.go +++ b/proto/proto_test.go @@ -104,6 +104,42 @@ func TestPathParam(t *testing.T) { } +func TestSetPathParam(t *testing.T) { + var payload, payload_after []byte + + payload = []byte("POST /post?param=test&user_id=1 HTTP/1.1\r\nContent-Length: 7\r\nHost: www.w3.org\r\n\r\na=1&b=2") + payload_after = []byte("POST /post?param=new&user_id=1 HTTP/1.1\r\nContent-Length: 7\r\nHost: www.w3.org\r\n\r\na=1&b=2") + + if payload = SetPathParam(payload, []byte("param"), []byte("new")); !bytes.Equal(payload, payload_after) { + t.Error("Should replace existing value", string(payload)) + } + + payload = []byte("POST /post?param=test&user_id=1 HTTP/1.1\r\nContent-Length: 7\r\nHost: www.w3.org\r\n\r\na=1&b=2") + payload_after = []byte("POST /post?param=test&user_id=2 HTTP/1.1\r\nContent-Length: 7\r\nHost: www.w3.org\r\n\r\na=1&b=2") + + if payload = SetPathParam(payload, []byte("user_id"), []byte("2")); !bytes.Equal(payload, payload_after) { + t.Error("Should replace existing value", string(payload)) + } + + + payload = []byte("POST /post HTTP/1.1\r\nContent-Length: 7\r\nHost: www.w3.org\r\n\r\na=1&b=2") + payload_after = []byte("POST /post?param=test HTTP/1.1\r\nContent-Length: 7\r\nHost: www.w3.org\r\n\r\na=1&b=2") + + if payload = SetPathParam(payload, []byte("param"), []byte("test")); !bytes.Equal(payload, payload_after) { + t.Error("Should set param if url have no params", string(payload)) + } + + + payload = []byte("POST /post?param=test HTTP/1.1\r\nContent-Length: 7\r\nHost: www.w3.org\r\n\r\na=1&b=2") + payload_after = []byte("POST /post?param=test&user_id=1 HTTP/1.1\r\nContent-Length: 7\r\nHost: www.w3.org\r\n\r\na=1&b=2") + + if payload = SetPathParam(payload, []byte("user_id"), []byte("1")); !bytes.Equal(payload, payload_after) { + t.Error("Should set param at the end if url params", string(payload)) + } +} + + + func TestSetHostHTTP10(t *testing.T) { var payload, payload_after []byte diff --git a/settings.go b/settings.go index 8fd2127..68e8f3a 100644 --- a/settings.go +++ b/settings.go @@ -89,6 +89,8 @@ func init() { flag.Var(&Settings.modifierConfig.headers, "http-set-header", "Inject additional headers to http reqest:\n\tgor --input-raw :8080 --output-http staging.com --http-set-header 'User-Agent: Gor'") flag.Var(&Settings.modifierConfig.headers, "output-http-header", "WARNING: `--output-http-header` DEPRECATED, use `--http-set-header` instead") + flag.Var(&Settings.modifierConfig.params, "http-set-param", "Set request url param, if param already exists it will be overwritten:\n\tgor --input-raw :8080 --output-http staging.com --http-set-param api_key=1") + flag.Var(&Settings.modifierConfig.methods, "http-allow-method", "Whitelist of HTTP methods to replay. Anything else will be dropped:\n\tgor --input-raw :8080 --output-http staging.com --http-allow-method GET --http-allow-method OPTIONS") flag.Var(&Settings.modifierConfig.methods, "output-http-method", "WARNING: `--output-http-method` DEPRECATED, use `--http-allow-method` instead") @@ -97,7 +99,7 @@ func init() { flag.Var(&Settings.modifierConfig.urlNegativeRegexp, "http-diallow-url", "A regexp to match requests against. Filter get matched agains full url with domain. Anything else will be dropped:\n\t gor --input-raw :8080 --output-http staging.com --http-disallow-url ^www.") - flag.Var(&Settings.modifierConfig.urlRewrite, "http-rewrite-url", "Rewrite the requst url based on a mapping:\n\tgor --input-raw :8080 --output-http staging.com --http-rewrite-url /v1/user/([^\\/]+)/ping:/v2/user/$1/ping") + flag.Var(&Settings.modifierConfig.urlRewrite, "http-rewrite-url", "Rewrite the request url based on a mapping:\n\tgor --input-raw :8080 --output-http staging.com --http-rewrite-url /v1/user/([^\\/]+)/ping:/v2/user/$1/ping") flag.Var(&Settings.modifierConfig.urlRewrite, "output-http-rewrite-url", "WARNING: `--output-http-rewrite-url` DEPRECATED, use `--http-rewrite-url` instead") flag.Var(&Settings.modifierConfig.headerFilters, "http-allow-header", "A regexp to match a specific header against. Requests with non-matching headers will be dropped:\n\t gor --input-raw :8080 --output-http staging.com --http-allow-header api-version:^v1")