All Downloads are FREE. Search and download functionalities are using the official Maven repository.

com.github.tomakehurst.wiremock.extension.requestfilter.RequestWrapper Maven / Gradle / Ivy

There is a newer version: 3.9.1
Show newest version
/*
 * Copyright (C) 2018-2023 Thomas Akehurst
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.github.tomakehurst.wiremock.extension.requestfilter;

import static com.github.tomakehurst.wiremock.common.Encoding.encodeBase64;
import static com.github.tomakehurst.wiremock.common.ParameterUtils.getFirstNonNull;
import static org.apache.commons.lang3.StringUtils.countMatches;
import static org.apache.commons.lang3.StringUtils.ordinalIndexOf;

import com.github.tomakehurst.wiremock.http.*;
import java.util.*;
import java.util.stream.Collectors;

public class RequestWrapper implements Request {

  private final Request delegate;

  private final RequestMethod method;
  private final FieldTransformer absoluteUrlTransformer;
  private final List addedHeaders;
  private final List removedHeaders;
  private final Map>> headerTransformers;
  private final Map additionalCookies;
  private final List cookiesToRemove;
  private final Map> cookieTransformers;
  private final FieldTransformer bodyTransformer;
  private final FieldTransformer multipartTransformer;

  public RequestWrapper(Request delegate) {
    this(
        delegate,
        null,
        null,
        Collections.emptyList(),
        Collections.emptyList(),
        Collections.emptyMap(),
        Collections.emptyMap(),
        Collections.emptyList(),
        Collections.emptyMap(),
        null,
        null);
  }

  public RequestWrapper(
      Request delegate,
      RequestMethod method,
      FieldTransformer absoluteUrlTransformer,
      List addedHeaders,
      List removedHeaders,
      Map>> headerTransformers,
      Map additionalCookies,
      List cookiesToRemove,
      Map> cookieTransformers,
      FieldTransformer bodyTransformer,
      FieldTransformer multipartTransformer) {
    this.delegate = delegate;

    this.method = method;
    this.absoluteUrlTransformer = absoluteUrlTransformer;
    this.addedHeaders = addedHeaders;
    this.removedHeaders = removedHeaders;
    this.headerTransformers = headerTransformers;
    this.additionalCookies = additionalCookies;
    this.cookiesToRemove = cookiesToRemove;
    this.cookieTransformers = cookieTransformers;
    this.bodyTransformer = bodyTransformer;
    this.multipartTransformer = multipartTransformer;
  }

  public static Builder create() {
    return new Builder();
  }

  @Override
  public String getUrl() {
    String absoluteUrl = getAbsoluteUrl();
    int relativeStartIndex =
        countMatches(absoluteUrl, "/") >= 3
            ? ordinalIndexOf(absoluteUrl, "/", 3)
            : absoluteUrl.length();
    return absoluteUrl.substring(relativeStartIndex);
  }

  @Override
  public String getAbsoluteUrl() {
    if (absoluteUrlTransformer != null) {
      return absoluteUrlTransformer.transform(delegate.getAbsoluteUrl());
    }

    return delegate.getAbsoluteUrl();
  }

  @Override
  public RequestMethod getMethod() {
    return getFirstNonNull(method, delegate.getMethod());
  }

  @Override
  public String getScheme() {
    return delegate.getScheme();
  }

  @Override
  public String getHost() {
    return delegate.getHost();
  }

  @Override
  public int getPort() {
    return delegate.getPort();
  }

  @Override
  public String getClientIp() {
    return delegate.getClientIp();
  }

  @Override
  public String getHeader(String key) {
    return getHeaders().getHeader(key).firstValue();
  }

  @Override
  public HttpHeader header(String key) {
    return getHeaders().getHeader(key);
  }

  @Override
  public ContentTypeHeader contentTypeHeader() {
    return delegate.contentTypeHeader();
  }

  @Override
  public HttpHeaders getHeaders() {
    List existingHeaders = new ArrayList<>(delegate.getHeaders().all());
    existingHeaders.addAll(addedHeaders);

    List combinedHeaders =
        existingHeaders.stream()
            .filter(httpHeader -> !removedHeaders.contains(httpHeader.key()))
            .map(
                httpHeader -> {
                  if (headerTransformers.containsKey(httpHeader.caseInsensitiveKey())) {
                    FieldTransformer> transformer =
                        headerTransformers.get(httpHeader.caseInsensitiveKey());
                    List newValues = transformer.transform(httpHeader.values());
                    return new HttpHeader(httpHeader.key(), newValues);
                  }

                  return httpHeader;
                })
            .collect(Collectors.toList());

    return new HttpHeaders(combinedHeaders);
  }

  @Override
  public boolean containsHeader(String key) {
    return getHeaders().keys().contains(key);
  }

  @Override
  public Set getAllHeaderKeys() {
    return getHeaders().keys();
  }

  @Override
  public Map getCookies() {
    Map cookieMap = new HashMap<>();
    for (Map.Entry entry : delegate.getCookies().entrySet()) {
      Cookie newCookie =
          cookieTransformers.containsKey(entry.getKey())
              ? cookieTransformers.get(entry.getKey()).transform(entry.getValue())
              : entry.getValue();

      if (!cookiesToRemove.contains(entry.getKey())) {
        cookieMap.put(entry.getKey(), newCookie);
      }
    }

    cookieMap.putAll(additionalCookies);

    return Collections.unmodifiableMap(cookieMap);
  }

  @Override
  public QueryParameter queryParameter(String key) {
    return delegate.queryParameter(key);
  }

  @Override
  public FormParameter formParameter(String key) {
    return delegate.formParameter(key);
  }

  @Override
  public Map formParameters() {
    return delegate.formParameters();
  }

  @Override
  public byte[] getBody() {
    if (bodyTransformer != null) {
      return bodyTransformer.transform(new Body(delegate.getBody())).asBytes();
    }

    return delegate.getBody();
  }

  @Override
  public String getBodyAsString() {
    if (bodyTransformer != null) {
      return bodyTransformer.transform(new Body(delegate.getBodyAsString())).asString();
    }

    return delegate.getBodyAsString();
  }

  @Override
  public String getBodyAsBase64() {
    return encodeBase64(getBody());
  }

  @Override
  public boolean isMultipart() {
    return delegate.isMultipart();
  }

  @Override
  public Collection getParts() {
    if (delegate.getParts() == null || multipartTransformer == null) {
      return delegate.getParts();
    }

    return delegate.getParts().stream()
        .map(multipartTransformer::transform)
        .collect(Collectors.toList());
  }

  @Override
  public Part getPart(String name) {
    if (multipartTransformer != null) {
      return multipartTransformer.transform(delegate.getPart(name));
    }

    return delegate.getPart(name);
  }

  @Override
  public boolean isBrowserProxyRequest() {
    return delegate.isBrowserProxyRequest();
  }

  @Override
  public Optional getOriginalRequest() {
    return delegate.getOriginalRequest();
  }

  @Override
  public String getProtocol() {
    return delegate.getProtocol();
  }

  public static class Builder {

    private RequestMethod requestMethod;
    private FieldTransformer absoluteUrlTransformer;

    private final List additionalHeaders = new ArrayList<>();
    private final List headersToRemove = new ArrayList<>();
    private final Map>> headerTransformers =
        new HashMap<>();

    private final Map additionalCookies = new HashMap<>();
    private final List cookiesToRemove = new ArrayList<>();
    private final Map> cookieTransformers = new HashMap<>();

    private FieldTransformer bodyTransformer;
    private FieldTransformer mutlipartTransformer;

    public Builder addHeader(String key, String... values) {
      additionalHeaders.add(new HttpHeader(key, values));
      return this;
    }

    public Builder removeHeader(String key) {
      headersToRemove.add(key);
      return this;
    }

    public Builder transformHeader(String key, FieldTransformer> transformer) {
      headerTransformers.put(CaseInsensitiveKey.from(key), transformer);
      return this;
    }

    public Builder setMethod(RequestMethod method) {
      requestMethod = method;
      return this;
    }

    public Builder transformAbsoluteUrl(FieldTransformer transformer) {
      absoluteUrlTransformer = transformer;
      return this;
    }

    public Request wrap(Request request) {
      return new RequestWrapper(
          request,
          requestMethod,
          absoluteUrlTransformer,
          additionalHeaders,
          headersToRemove,
          headerTransformers,
          additionalCookies,
          cookiesToRemove,
          cookieTransformers,
          bodyTransformer,
          mutlipartTransformer);
    }

    public Builder transformBody(FieldTransformer transformer) {
      bodyTransformer = transformer;
      return this;
    }

    public Builder transformCookie(String name, FieldTransformer transformer) {
      cookieTransformers.put(name, transformer);
      return this;
    }

    public Builder transformParts(FieldTransformer transformer) {
      mutlipartTransformer = transformer;
      return this;
    }

    public Builder addCookie(String name, Cookie value) {
      additionalCookies.put(name, value);
      return this;
    }

    public Builder removeCookie(String name) {
      cookiesToRemove.add(name);
      return this;
    }
  }
}




© 2015 - 2024 Weber Informatics LLC | Privacy Policy