/*
 * Copyright 2020 Signal Messenger, LLC
 * SPDX-License-Identifier: AGPL-3.0-only
 */

package org.whispersystems.textsecuregcm.filters;

import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyByte;
import static org.mockito.ArgumentMatchers.anyInt;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

import com.google.common.net.HttpHeaders;
import com.google.common.net.InetAddresses;
import com.google.protobuf.ByteString;
import com.google.rpc.ErrorInfo;
import com.google.rpc.Status;
import com.vdurmont.semver4j.Semver;
import io.dropwizard.auth.basic.BasicCredentials;
import io.grpc.ManagedChannel;
import io.grpc.Server;
import io.grpc.StatusRuntimeException;
import io.grpc.inprocess.InProcessChannelBuilder;
import io.grpc.inprocess.InProcessServerBuilder;
import io.grpc.protobuf.StatusProto;
import jakarta.servlet.FilterChain;
import jakarta.servlet.ServletException;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.time.Instant;
import java.util.EnumMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.UUID;
import java.util.stream.Stream;
import javax.annotation.Nullable;
import org.apache.commons.lang3.RandomStringUtils;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.Arguments;
import org.junit.jupiter.params.provider.MethodSource;
import org.signal.chat.rpc.EchoRequest;
import org.signal.chat.rpc.EchoServiceGrpc;
import org.whispersystems.textsecuregcm.auth.AccountAuthenticator;
import org.whispersystems.textsecuregcm.auth.AuthenticatedDevice;
import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration;
import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicRemoteDeprecationConfiguration;
import org.whispersystems.textsecuregcm.grpc.EchoServiceImpl;
import org.whispersystems.textsecuregcm.grpc.GrpcExceptions;
import org.whispersystems.textsecuregcm.grpc.MockRequestAttributesInterceptor;
import org.whispersystems.textsecuregcm.grpc.RequestAttributes;
import org.whispersystems.textsecuregcm.storage.Account;
import org.whispersystems.textsecuregcm.storage.AccountsManager;
import org.whispersystems.textsecuregcm.storage.Device;
import org.whispersystems.textsecuregcm.storage.DeviceCapability;
import org.whispersystems.textsecuregcm.storage.DynamicConfigurationManager;
import org.whispersystems.textsecuregcm.tests.util.AuthHelper;
import org.whispersystems.textsecuregcm.util.ua.ClientPlatform;

class RemoteDeprecationFilterTest {

  private DynamicConfiguration dynamicConfiguration;

  private RemoteDeprecationFilter remoteDeprecationFilter;

  private static final UUID ACCOUNT_IDENTIFIER_WITH_SPQR = UUID.randomUUID();
  private static final UUID ACCOUNT_IDENTIFIER_WITHOUT_SPQR = UUID.randomUUID();

  private static final String PASSWORD = RandomStringUtils.insecure().nextAlphanumeric(16);

  @BeforeEach
  void setUp() {
    final AccountsManager accountsManager = mock(AccountsManager.class);
    when(accountsManager.getByAccountIdentifier(any())).thenReturn(Optional.empty());

    final Account accountWithSpqr = buildMockAccount(true);
    final Account accountWithoutSpqr = buildMockAccount(false);

    when(accountsManager.getByAccountIdentifier(ACCOUNT_IDENTIFIER_WITH_SPQR))
        .thenReturn(Optional.of(accountWithSpqr));

    when(accountsManager.getByAccountIdentifier(ACCOUNT_IDENTIFIER_WITHOUT_SPQR))
        .thenReturn(Optional.of(accountWithoutSpqr));

    final AccountAuthenticator accountAuthenticator = mock(AccountAuthenticator.class);
    when(accountAuthenticator.authenticate(any())).thenReturn(Optional.empty());

    when(accountAuthenticator.authenticate(new BasicCredentials(ACCOUNT_IDENTIFIER_WITH_SPQR.toString(), PASSWORD)))
        .thenReturn(Optional.of(new AuthenticatedDevice(ACCOUNT_IDENTIFIER_WITH_SPQR, Device.PRIMARY_ID, Instant.now())));

    when(accountAuthenticator.authenticate(new BasicCredentials(ACCOUNT_IDENTIFIER_WITHOUT_SPQR.toString(), PASSWORD)))
        .thenReturn(Optional.of(new AuthenticatedDevice(ACCOUNT_IDENTIFIER_WITHOUT_SPQR, Device.PRIMARY_ID, Instant.now())));

    @SuppressWarnings("unchecked") final DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager =
        mock(DynamicConfigurationManager.class);

    dynamicConfiguration = mock(DynamicConfiguration.class);

    when(dynamicConfigurationManager.getConfiguration()).thenReturn(dynamicConfiguration);
    when(dynamicConfiguration.getRemoteDeprecationConfiguration())
        .thenReturn(DynamicRemoteDeprecationConfiguration.DEFAULT);

    remoteDeprecationFilter =
        new RemoteDeprecationFilter(accountsManager, accountAuthenticator, dynamicConfigurationManager);
  }

  private static Account buildMockAccount(final boolean hasSpqr) {
    final Device device = mock(Device.class);
    when(device.hasCapability(any())).thenReturn(false);
    when(device.hasCapability(DeviceCapability.SPARSE_POST_QUANTUM_RATCHET)).thenReturn(hasSpqr);

    final Account account = mock(Account.class);
    when(account.getDevice(anyByte())).thenReturn(Optional.empty());
    when(account.getDevice(Device.PRIMARY_ID)).thenReturn(Optional.of(device));

    return account;
  }

  @Test
  void testEmptyMap() throws IOException, ServletException {
    final HttpServletRequest servletRequest = mock(HttpServletRequest.class);
    final HttpServletResponse servletResponse = mock(HttpServletResponse.class);
    final FilterChain filterChain = mock(FilterChain.class);

    when(servletRequest.getHeader("UserAgent")).thenReturn("Signal-Android/4.68.3");

    assertDoesNotThrow(() -> remoteDeprecationFilter.doFilter(servletRequest, servletResponse, filterChain));

    verify(filterChain).doFilter(servletRequest, servletResponse);
    verify(servletResponse, never()).sendError(anyInt());
  }

  private static DynamicRemoteDeprecationConfiguration configuration() {
    final EnumMap<ClientPlatform, Semver> minimumVersionsByPlatform = new EnumMap<>(Map.of(
        ClientPlatform.ANDROID, new Semver("1.0.0"),
        ClientPlatform.IOS, new Semver("1.0.0"),
        ClientPlatform.DESKTOP, new Semver("1.0.0")
    ));

    final EnumMap<ClientPlatform, Semver> versionsPendingDeprecationByPlatform = new EnumMap<>(Map.of(
        ClientPlatform.ANDROID, new Semver("1.1.0"),
        ClientPlatform.IOS, new Semver("1.1.0"),
        ClientPlatform.DESKTOP, new Semver("1.1.0")
    ));

    final EnumMap<ClientPlatform, Set<Semver>> blockedVersionsByPlatform = new EnumMap<>(Map.of(
        ClientPlatform.DESKTOP, Set.of(new Semver("8.0.0-beta.2"))
    ));

    final EnumMap<ClientPlatform, Set<Semver>> versionsPendingBlockByPlatform = new EnumMap<>(Map.of(
        ClientPlatform.DESKTOP, Set.of(new Semver("8.0.0-beta.3"))
    ));

    return new DynamicRemoteDeprecationConfiguration(
        minimumVersionsByPlatform,
        versionsPendingDeprecationByPlatform,
        blockedVersionsByPlatform,
        versionsPendingBlockByPlatform,
        false,
        false);
  }

  @ParameterizedTest
  @MethodSource
  void testFilter(final String userAgent, final boolean expectDeprecation) throws IOException, ServletException {
    when(dynamicConfiguration.getRemoteDeprecationConfiguration()).thenReturn(configuration());

    final HttpServletRequest servletRequest = mock(HttpServletRequest.class);
    final HttpServletResponse servletResponse = mock(HttpServletResponse.class);
    final FilterChain filterChain = mock(FilterChain.class);

    when(servletRequest.getHeader(HttpHeaders.USER_AGENT)).thenReturn(userAgent);

    remoteDeprecationFilter.doFilter(servletRequest, servletResponse, filterChain);

    if (expectDeprecation) {
      verify(filterChain, never()).doFilter(any(), any());
      verify(servletResponse).sendError(499);
    } else {
      verify(filterChain).doFilter(servletRequest, servletResponse);
      verify(servletResponse, never()).sendError(anyInt());
    }
  }

  @SuppressWarnings("DataFlowIssue")
  @ParameterizedTest
  @MethodSource(value="testFilter")
  void testGrpcFilter(final String userAgentString, final boolean expectDeprecation) throws IOException, InterruptedException {
    when(dynamicConfiguration.getRemoteDeprecationConfiguration()).thenReturn(configuration());

    final MockRequestAttributesInterceptor mockRequestAttributesInterceptor = new MockRequestAttributesInterceptor();
    mockRequestAttributesInterceptor.setRequestAttributes(new RequestAttributes(InetAddresses.forString("127.0.0.1"), userAgentString, null));

    final Server testServer = InProcessServerBuilder.forName("RemoteDeprecationFilterTest")
        .directExecutor()
        .addService(new EchoServiceImpl())
        .intercept(remoteDeprecationFilter)
        .intercept(mockRequestAttributesInterceptor)
        .build()
        .start();

    final ManagedChannel channel = InProcessChannelBuilder.forName("RemoteDeprecationFilterTest")
        .directExecutor()
        .userAgent(userAgentString)
        .build();

    try {
      final String echoString = "Test";

      final EchoServiceGrpc.EchoServiceBlockingStub client = EchoServiceGrpc.newBlockingStub(channel);
      final EchoRequest req = EchoRequest.newBuilder()
          .setPayload(ByteString.copyFromUtf8(echoString))
          .build();

      if (expectDeprecation) {
        final StatusRuntimeException e = assertThrows(StatusRuntimeException.class, () -> client.echo(req));
        final Status status = StatusProto.fromThrowable(e);
        final ErrorInfo errorInfo = assertDoesNotThrow(() -> status.getDetailsList().stream()
            .filter(any -> any.is(ErrorInfo.class)).findFirst()
            .orElseThrow(() -> new AssertionError("No error info found"))
            .unpack(ErrorInfo.class));

        assertEquals(GrpcExceptions.DOMAIN, errorInfo.getDomain());
        assertEquals(io.grpc.Status.Code.INVALID_ARGUMENT.value(), status.getCode());
        assertEquals("UPGRADE_REQUIRED", errorInfo.getReason());
      } else {
        assertEquals(echoString, client.echo(req).getPayload().toStringUtf8());
      }
    } finally {
      testServer.shutdownNow();
      testServer.awaitTermination();
    }
  }

  private static Stream<Arguments> testFilter() {
    return Stream.of(
        Arguments.of("Unrecognized UA", false),
        Arguments.of("Signal-Android/4.68.3", false),
        Arguments.of("Signal-iOS/3.9.0", false),
        Arguments.of("Signal-Desktop/1.2.3", false),
        Arguments.of("Signal-Android/0.68.3", true),
        Arguments.of("Signal-iOS/0.9.0", true),
        Arguments.of("Signal-Desktop/0.2.3", true),
        Arguments.of("Signal-Desktop/8.0.0-beta.2", true),
        Arguments.of("Signal-Desktop/8.0.0-beta.1", false),
        Arguments.of("Signal-iOS/8.0.0-beta.2", false));
  }

  @ParameterizedTest
  @MethodSource
  void isMissingSpqrCapability(@Nullable final String authHeader, final boolean expectMissingCapability) {
    assertEquals(expectMissingCapability, remoteDeprecationFilter.isMissingSpqrCapability(authHeader));
  }

  private static List<Arguments> isMissingSpqrCapability() {
    final String password = RandomStringUtils.insecure().nextAlphanumeric(16);
    AuthHelper.getAuthHeader(ACCOUNT_IDENTIFIER_WITH_SPQR, password);

    return List.of(
        Arguments.argumentSet("No authentication header", null, false),

        Arguments.argumentSet("Authentication header for device with SPQR",
            AuthHelper.getAuthHeader(ACCOUNT_IDENTIFIER_WITH_SPQR, PASSWORD), false),

        Arguments.argumentSet("Authentication header for device without SPQR",
            AuthHelper.getAuthHeader(ACCOUNT_IDENTIFIER_WITHOUT_SPQR, PASSWORD), true)
    );
  }

  @ParameterizedTest
  @MethodSource
  void shouldBlockSpqr(final boolean enforcementPending, final boolean spqrRequired, final String authHeader, final boolean expectBlock) {
    when(dynamicConfiguration.getRemoteDeprecationConfiguration())
        .thenReturn(new DynamicRemoteDeprecationConfiguration(
            DynamicRemoteDeprecationConfiguration.DEFAULT.minimumVersions(),
            DynamicRemoteDeprecationConfiguration.DEFAULT.versionsPendingDeprecation(),
            DynamicRemoteDeprecationConfiguration.DEFAULT.blockedVersions(),
            DynamicRemoteDeprecationConfiguration.DEFAULT.versionsPendingBlock(),
            enforcementPending,
            spqrRequired));

    assertEquals(expectBlock, remoteDeprecationFilter.shouldBlock(null, authHeader));
  }

  private static List<Arguments> shouldBlockSpqr() {
    return List.of(
        Arguments.argumentSet("Has capability, no enforcement",
            false, false, AuthHelper.getAuthHeader(ACCOUNT_IDENTIFIER_WITH_SPQR, PASSWORD), false),

        Arguments.argumentSet("Has capability, enforcement pending",
            true, false, AuthHelper.getAuthHeader(ACCOUNT_IDENTIFIER_WITH_SPQR, PASSWORD), false),

        Arguments.argumentSet("Has capability, enforcement active",
            false, true, AuthHelper.getAuthHeader(ACCOUNT_IDENTIFIER_WITH_SPQR, PASSWORD), false),

        Arguments.argumentSet("Missing capability, no enforcement",
            false, false, AuthHelper.getAuthHeader(ACCOUNT_IDENTIFIER_WITHOUT_SPQR, PASSWORD), false),

        Arguments.argumentSet("Missing capability, enforcement pending",
            true, false, AuthHelper.getAuthHeader(ACCOUNT_IDENTIFIER_WITHOUT_SPQR, PASSWORD), false),

        Arguments.argumentSet("Missing capability, enforcement active",
            false, true, AuthHelper.getAuthHeader(ACCOUNT_IDENTIFIER_WITHOUT_SPQR, PASSWORD), true)
    );
  }
}
