/*
 * Copyright (C) 2016 The Android Open Source Project
 *
 * 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 Lice`nse 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 androidx.test.rule;

import static androidx.test.internal.util.Checks.checkArgument;
import static androidx.test.internal.util.Checks.checkNotNull;

import android.util.Log;
import androidx.annotation.NonNull;
import androidx.annotation.RestrictTo;
import androidx.annotation.RestrictTo.Scope;
import androidx.annotation.VisibleForTesting;
import java.util.Properties;
import org.junit.rules.TestRule;
import org.junit.runner.Description;
import org.junit.runners.model.Statement;

/**
 * A {@code TestRule} to forward network traffic to a specific port. By default all traffic is
 * forwarded to default address {@code #DEFAULT_PROXY_HOST}:{@code #DEFAULT_PROXY_PORT} unless
 * otherwise specified
 *
 * <p>Note: Traffic forwarding will only apply to the current process under test.
 *
 * <p><b>This API is currently in beta.</b>
 *
 * @hide
 */
@RestrictTo(Scope.LIBRARY_GROUP)
public class PortForwardingRule implements TestRule {

  private static final String TAG = "PortForwardingRule";

  public static final int MIN_PORT = 1024;
  public static final int MAX_PORT = 65535;

  @VisibleForTesting static final int DEFAULT_PROXY_PORT = 8080;
  @VisibleForTesting static final String DEFAULT_PROXY_HOST = "127.0.0.1";
  @VisibleForTesting static final String HTTP_HOST_PROPERTY = "http.proxyHost";
  @VisibleForTesting static final String HTTPS_HOST_PROPERTY = "https.proxyHost";
  @VisibleForTesting static final String HTTP_PORT_PROPERTY = "http.proxyPort";
  @VisibleForTesting static final String HTTPS_PORT_PROPERTY = "https.proxyPort";

  @VisibleForTesting final String proxyHost;
  @VisibleForTesting final int proxyPort;
  @VisibleForTesting Properties prop;

  private Properties backUpProp;

  /** @hide */
  public static class Builder {

    private String proxyHost = DEFAULT_PROXY_HOST;
    private int proxyPort = DEFAULT_PROXY_PORT;
    private Properties prop = System.getProperties();

    /**
     * Builder to set a specific host address to forward the network traffic to.
     *
     * @param proxyHost The host address to which the network traffic is forwarded during tests.
     */
    public Builder withProxyHost(@NonNull String proxyHost) {
      this.proxyHost = checkNotNull(proxyHost);
      return this;
    }

    /**
     * Builder to set a specific port number to forward the network traffic to.
     *
     * @param proxyPort The port number to which the network traffic is forwarded during tests.
     */
    public Builder withProxyPort(int proxyPort) {
      checkArgument(
          proxyPort >= MIN_PORT && proxyPort <= MAX_PORT,
          "%d is used as a proxy port, must in range [%d, %d]",
          proxyPort,
          MIN_PORT,
          MAX_PORT);
      this.proxyPort = proxyPort;
      return this;
    }

    /**
     * Builder which allows to pass a {@link Properties} object for testing. This will help to avoid
     * the system properties being affected by tests.
     *
     * @param properties A pre-constructed properties object for testing.
     */
    public Builder withProperties(@NonNull Properties properties) {
      prop = checkNotNull(properties);
      return this;
    }

    public PortForwardingRule build() {
      return new PortForwardingRule(this);
    }
  }

  private PortForwardingRule(Builder builder) {
    this(builder.proxyHost, builder.proxyPort, builder.prop);
  }

  protected PortForwardingRule(int proxyPort) {
    this(DEFAULT_PROXY_HOST, proxyPort, System.getProperties());
  }

  @VisibleForTesting
  PortForwardingRule(String proxyHost, int proxyPort, @NonNull Properties properties) {
    this.proxyHost = proxyHost;
    this.proxyPort = proxyPort;
    prop = checkNotNull(properties);
    backUpProp = new Properties();
    backUpProperties();
  }

  protected static int getDefaultPort() {
    return DEFAULT_PROXY_PORT;
  }

  /**
   * Override this method to execute any code that should run before port forwarding. This method is
   * called before each test method, including any method annotated with <a
   * href="http://junit.sourceforge.net/javadoc/org/junit/Before.html"><code>Before</code></a>.
   */
  protected void beforePortForwarding() {
    // empty by default
  }

  /**
   * Override this method to execute any code that should run after port forwarding is set up, but
   * before any test code is run including any method annotated with <a
   * href="http://junit.sourceforge.net/javadoc/org/junit/Before.html"><code>Before</code></a>.
   */
  protected void afterPortForwarding() {
    // empty by default
  }

  /**
   * Override this method to execute any code that should run before port forwarding is restored.
   * This method is called after each test method, including any method annotated with <a
   * href="http://junit.sourceforge.net/javadoc/org/junit/After.html"><code>After</code></a>.
   */
  protected void beforeRestoreForwarding() {
    // empty by default
  }

  /**
   * Override this method to execute any code that should run after port forwarding is restored.
   * This method is called after each test method, including any method annotated with <a
   * href="http://junit.sourceforge.net/javadoc/org/junit/After.html"><code>After</code></a>.
   */
  protected void afterRestoreForwarding() {
    // empty by default
  }

  /**
   * Set the port forwarding system properties (Note: for the process under test only). This method
   * will also back up the existing system properties for later to restore.
   */
  private void setPortForwarding() {
    beforePortForwarding();
    prop.setProperty(HTTP_HOST_PROPERTY, proxyHost);
    prop.setProperty(HTTPS_HOST_PROPERTY, proxyHost);
    prop.setProperty(HTTP_PORT_PROPERTY, String.valueOf(proxyPort));
    prop.setProperty(HTTPS_PORT_PROPERTY, String.valueOf(proxyPort));
    afterPortForwarding();
  }

  /** Restore the system properties backed up (Note: for the process under test only) */
  private void restorePortForwarding() {
    try {
      beforeRestoreForwarding();
    } finally {
      restoreOneProperty(prop, backUpProp, HTTP_HOST_PROPERTY);
      restoreOneProperty(prop, backUpProp, HTTPS_HOST_PROPERTY);
      restoreOneProperty(prop, backUpProp, HTTP_PORT_PROPERTY);
      restoreOneProperty(prop, backUpProp, HTTPS_PORT_PROPERTY);
      afterRestoreForwarding();
    }
  }

  private void backUpProperties() {
    if (prop.getProperty(HTTP_HOST_PROPERTY) != null) {
      backUpProp.setProperty(HTTP_HOST_PROPERTY, prop.getProperty(HTTP_HOST_PROPERTY));
    }
    if (prop.getProperty(HTTPS_HOST_PROPERTY) != null) {
      backUpProp.setProperty(HTTPS_HOST_PROPERTY, prop.getProperty(HTTPS_HOST_PROPERTY));
    }
    if (prop.getProperty(HTTP_PORT_PROPERTY) != null) {
      backUpProp.setProperty(HTTP_PORT_PROPERTY, prop.getProperty(HTTP_PORT_PROPERTY));
    }
    if (prop.getProperty(HTTPS_PORT_PROPERTY) != null) {
      backUpProp.setProperty(HTTPS_PORT_PROPERTY, prop.getProperty(HTTPS_PORT_PROPERTY));
    }
  }

  private void restoreOneProperty(Properties prop, Properties backUpProp, String key) {
    if (backUpProp.getProperty(key) != null) {
      prop.setProperty(key, backUpProp.getProperty(key));
    } else {
      prop.remove(key);
    }
  }

  @Override
  public Statement apply(final Statement base, Description description) {
    return new PortForwardingStatement(base);
  }

  /** {@link Statement} that set/restore port forwarding before/after the execution of the test. */
  private class PortForwardingStatement extends Statement {

    private final Statement base;

    public PortForwardingStatement(Statement base) {
      this.base = base;
    }

    @Override
    public void evaluate() throws Throwable {
      try {
        setPortForwarding();
        Log.i(
            TAG,
            String.format(
                "The current process traffic is forwarded to %s:%d", proxyHost, proxyPort));
        base.evaluate();
      } finally {
        restorePortForwarding();
        Log.i(TAG, "Current process traffic forwarding is cancelled");
      }
    }
  }
}