コード例 #1
0
ファイル: token_modifier.go プロジェクト: juan15Chen/gor
func process(buf []byte) {
	// First byte indicate payload type, possible values:
	//  1 - Request
	//  2 - Response
	//  3 - ReplayedResponse
	payloadType := buf[0]
	headerSize := bytes.IndexByte(buf, '\n') + 1
	header := buf[:headerSize-1]

	// Header contains space separated values of: request type, request id, and request start time (or round-trip time for responses)
	meta := bytes.Split(header, []byte(" "))
	// For each request you should receive 3 payloads (request, response, replayed response) with same request id
	reqID := string(meta[1])
	payload := buf[headerSize:]

	Debug("Received payload:", string(buf))

	switch payloadType {
	case '1': // Request
		if bytes.Equal(proto.Path(payload), []byte("/token")) {
			originalTokens[reqID] = []byte{}
			Debug("Found token request:", reqID)
		} else {
			token, vs, _ := proto.PathParam(payload, []byte("token"))

			if vs != -1 { // If there is GET token param
				if alias, ok := tokenAliases[string(token)]; ok {
					// Rewrite original token to alias
					payload = proto.SetPathParam(payload, []byte("token"), alias)

					// Copy modified payload to our buffer
					buf = append(buf[:headerSize], payload...)
				}
			}
		}

		// Emitting data back
		os.Stdout.Write(encode(buf))
	case '2': // Original response
		if _, ok := originalTokens[reqID]; ok {
			// Token is inside response body
			secureToken := proto.Body(payload)
			originalTokens[reqID] = secureToken
			Debug("Remember origial token:", string(secureToken))
		}
	case '3': // Replayed response
		if originalToken, ok := originalTokens[reqID]; ok {
			delete(originalTokens, reqID)
			secureToken := proto.Body(payload)
			tokenAliases[string(originalToken)] = secureToken

			Debug("Create alias for new token token, was:", string(originalToken), "now:", string(secureToken))
		}
	}
}
コード例 #2
0
ファイル: http_modifier.go プロジェクト: robot0x/gor
func (m *HTTPModifier) Rewrite(payload []byte) (response []byte) {
	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 {
		path := proto.Path(payload)

		matched := false

		for _, f := range m.config.urlRegexp {
			if f.regexp.Match(path) {
				matched = true
				break
			}
		}

		if !matched {
			return
		}
	}

	if len(m.config.urlNegativeRegexp) > 0 {
		path := proto.Path(payload)

		for _, f := range m.config.urlNegativeRegexp {
			if f.regexp.Match(path) {
				return
			}
		}
	}

	if len(m.config.headerFilters) > 0 {
		for _, f := range m.config.headerFilters {
			value := proto.Header(payload, f.name)

			if len(value) > 0 && !f.regexp.Match(value) {
				return
			}
		}
	}

	if len(m.config.headerNegativeFilters) > 0 {
		for _, f := range m.config.headerNegativeFilters {
			value := proto.Header(payload, f.name)

			if len(value) > 0 && f.regexp.Match(value) {
				return
			}
		}
	}

	if len(m.config.headerHashFilters) > 0 {
		for _, f := range m.config.headerHashFilters {
			value := proto.Header(payload, f.name)

			if len(value) > 0 {
				hasher := fnv.New32a()
				hasher.Write(value)

				if (hasher.Sum32() % 100) >= f.percent {
					return
				}
			}
		}
	}

	if len(m.config.paramHashFilters) > 0 {
		for _, f := range m.config.paramHashFilters {
			value, s, _ := proto.PathParam(payload, f.name)

			if s != -1 {
				hasher := fnv.New32a()
				hasher.Write(value)

				if (hasher.Sum32() % 100) >= f.percent {
					return
				}
			}
		}
	}

	if len(m.config.urlRewrite) > 0 {
		path := proto.Path(payload)

		for _, f := range m.config.urlRewrite {
			if f.src.Match(path) {
				path = f.src.ReplaceAll(path, f.target)
				payload = proto.SetPath(payload, path)

				break
			}
		}
	}

	return payload
}