/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You 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 org.apache.activemq.artemis.tests.integration.routing;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.fail;
import static org.junit.jupiter.api.Assumptions.assumeTrue;

import org.apache.activemq.artemis.api.core.BroadcastGroupConfiguration;
import org.apache.activemq.artemis.api.core.DiscoveryGroupConfiguration;
import org.apache.activemq.artemis.api.core.UDPBroadcastEndpointFactory;
import org.apache.activemq.artemis.core.remoting.impl.netty.TransportConstants;
import org.apache.activemq.artemis.core.security.Role;
import org.apache.activemq.artemis.core.server.ActiveMQServers;
import org.apache.activemq.artemis.core.server.cluster.impl.MessageLoadBalancingType;
import org.apache.activemq.artemis.core.server.routing.policies.FirstElementPolicy;
import org.apache.activemq.artemis.core.server.routing.policies.PolicyFactoryResolver;
import org.apache.activemq.artemis.core.server.routing.targets.Target;
import org.apache.activemq.artemis.core.server.routing.KeyType;
import org.apache.activemq.artemis.core.settings.HierarchicalRepository;
import org.apache.activemq.artemis.spi.core.security.ActiveMQJAASSecurityManager;
import org.apache.activemq.artemis.tests.extensions.parameterized.ParameterizedTestExtension;
import org.apache.activemq.artemis.tests.extensions.parameterized.Parameters;
import org.apache.activemq.artemis.tests.integration.security.SecurityTest;
import org.eclipse.paho.client.mqttv3.MqttClient;
import org.eclipse.paho.client.mqttv3.MqttConnectOptions;
import org.eclipse.paho.client.mqttv3.persist.MemoryPersistence;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.TestTemplate;
import org.junit.jupiter.api.extension.ExtendWith;

import javax.jms.Connection;
import javax.jms.ConnectionFactory;
import javax.net.ssl.KeyManager;
import javax.net.ssl.KeyManagerFactory;
import javax.net.ssl.SSLContext;
import javax.net.ssl.SSLSocketFactory;
import javax.net.ssl.TrustManagerFactory;
import java.lang.management.ManagementFactory;
import java.net.InetAddress;
import java.net.URL;
import java.net.UnknownHostException;
import java.security.KeyStore;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.UUID;
import java.util.stream.Collectors;

@ExtendWith(ParameterizedTestExtension.class)
public class KeyTypeTest extends RoutingTestBase {

   private static final String MOCK_POLICY_NAME = "MOCK_POLICY";

   @Parameters(name = "protocol: {0}")
   public static Collection<Object[]> data() {
      Collection<Object[]> data = new ArrayList<>();

      for (String protocol : Arrays.asList(AMQP_PROTOCOL, CORE_PROTOCOL, MQTT_PROTOCOL, OPENWIRE_PROTOCOL)) {
         data.add(new Object[] {protocol});
      }

      return data;
   }


   static {
      String path = System.getProperty("java.security.auth.login.config");
      if (path == null) {
         URL resource = SecurityTest.class.getClassLoader().getResource("login.config");
         if (resource != null) {
            path = resource.getFile();
            System.setProperty("java.security.auth.login.config", path);
         }
      }
   }

   private final String protocol;

   private final List<String> keys = new ArrayList<>();

   public KeyTypeTest(String protocol) {
      this.protocol = protocol;
   }

   @BeforeEach
   public void setup() throws Exception {
      PolicyFactoryResolver.getInstance().registerPolicyFactory(MOCK_POLICY_NAME, () -> new FirstElementPolicy(MOCK_POLICY_NAME) {
         @Override
         public Target selectTarget(List<Target> targets, String key) {
            keys.add(key);
            return super.selectTarget(targets, key);
         }
      });
   }

   @TestTemplate
   public void testClientIDKey() throws Exception {
      setupPrimaryServerWithDiscovery(0, GROUP_ADDRESS, GROUP_PORT, true, true, false);
      setupRouterServerWithDiscovery(0, KeyType.CLIENT_ID, MOCK_POLICY_NAME, null, true, null, 1);
      startServers(0);

      testConnection(TransportConstants.DEFAULT_HOST, TransportConstants.DEFAULT_PORT + 0,
              "test", null, null, false, false, -1);

      assertEquals(1, keys.size());
      assertEquals("test", keys.get(0));
   }

   @Override
   protected boolean isForceUniqueStorageManagerIds() {
      return false;
   }

   @TestTemplate
   public void testClientIDKeyOnBackup() throws Exception {
      setupPrimaryServerWithDiscovery(0, GROUP_ADDRESS, GROUP_PORT, true, true, false);
      setupDiscoveryClusterConnection("cluster0", 0, "dg1", "queues", MessageLoadBalancingType.OFF, 1, true);
      setupRouterServerWithCluster(0, KeyType.CLIENT_ID, FirstElementPolicy.NAME, null, true, null, 1, "cluster0");
      setupBackupServer(1, 0, false, HAType.SharedNothingReplication, true);
      UDPBroadcastEndpointFactory endpoint = new UDPBroadcastEndpointFactory().setGroupAddress(GROUP_ADDRESS).setGroupPort(GROUP_PORT);
      List<String> connectorInfos = getServer(1).getConfiguration().getConnectorConfigurations().keySet().stream().collect(Collectors.toList());
      BroadcastGroupConfiguration bcConfig = new BroadcastGroupConfiguration().setName("bg1").setBroadcastPeriod(1000).setConnectorInfos(connectorInfos).setEndpointFactory(endpoint);
      DiscoveryGroupConfiguration dcConfig = new DiscoveryGroupConfiguration().setName("dg1").setRefreshTimeout(5000).setDiscoveryInitialWaitTimeout(5000).setBroadcastEndpointFactory(endpoint);
      getServer(1).getConfiguration().addBroadcastGroupConfiguration(bcConfig).addDiscoveryGroupConfiguration(dcConfig.getName(), dcConfig);
      setupDiscoveryClusterConnection("cluster0", 1, "dg1", "queues", MessageLoadBalancingType.OFF, 1, true);
      setupRouterServerWithCluster(1, KeyType.CLIENT_ID, MOCK_POLICY_NAME, null, true, null, 1, "cluster0");
      startServers(0, 1);

      waitForTopology(getServer(0), 1, 1);

      getServer(0).fail(true);

      waitForFailoverTopology(1);

      testConnection(TransportConstants.DEFAULT_HOST, TransportConstants.DEFAULT_PORT + 1,
              "test", null, null, false, false, -1);

      assertEquals(1, keys.size());
      assertEquals("test", keys.get(0));
   }

   @TestTemplate
   public void testSNIHostKey() throws Exception {
      String localHostname = "localhost.localdomain";

      if (!checkLocalHostname(localHostname)) {
         localHostname = "artemis.localtest.me";

         if (!checkLocalHostname(localHostname)) {
            localHostname = "localhost";

            assumeTrue(CORE_PROTOCOL.equals(protocol) && checkLocalHostname(localHostname));
         }
      }

      setupPrimaryServerWithDiscovery(0, GROUP_ADDRESS, GROUP_PORT, true, true, false);
      getDefaultServerAcceptor(0).getParams().put(TransportConstants.SSL_ENABLED_PROP_NAME, true);
      getDefaultServerAcceptor(0).getParams().put(TransportConstants.KEYSTORE_PATH_PROP_NAME, "server-keystore.jks");
      getDefaultServerAcceptor(0).getParams().put(TransportConstants.KEYSTORE_PASSWORD_PROP_NAME, "securepass");

      setupRouterServerWithDiscovery(0, KeyType.SNI_HOST, MOCK_POLICY_NAME, null, true, null, 1);
      startServers(0);

      testConnection(localHostname, TransportConstants.DEFAULT_PORT + 0,
              null, null, null, true, false, -1);

      assertEquals(1, keys.size());
      assertEquals(localHostname, keys.get(0));
   }

   @TestTemplate
   public void testSourceIPKey() throws Exception {
      setupPrimaryServerWithDiscovery(0, GROUP_ADDRESS, GROUP_PORT, true, true, false);
      setupRouterServerWithDiscovery(0, KeyType.SOURCE_IP, MOCK_POLICY_NAME, null, true, null, 1);
      startServers(0);

      testConnection(TransportConstants.DEFAULT_HOST, TransportConstants.DEFAULT_PORT + 0,
              null, null, null, false, false, -1);

      assertEquals(1, keys.size());
      assertEquals(InetAddress.getLoopbackAddress().getHostAddress(), keys.get(0));
   }

   @TestTemplate
   public void testUserNameKey() throws Exception {
      setupPrimaryServerWithDiscovery(0, GROUP_ADDRESS, GROUP_PORT, true, true, false);
      setupRouterServerWithDiscovery(0, KeyType.USER_NAME, MOCK_POLICY_NAME, null, true, null, 1);
      startServers(0);

      testConnection(TransportConstants.DEFAULT_HOST, TransportConstants.DEFAULT_PORT + 0,
              null, "admin", "admin", false, false, -1);

      assertEquals(1, keys.size());
      assertEquals("admin", keys.get(0));
   }

   @TestTemplate
   public void testUserNameKeyFromCertificate() throws Exception {
      setupPrimaryServerWithDiscovery(0, GROUP_ADDRESS, GROUP_PORT, true, true, false);
      getDefaultServerAcceptor(0).getParams().put("securityDomain", "CertLogin");
      getDefaultServerAcceptor(0).getParams().put(TransportConstants.SSL_ENABLED_PROP_NAME, true);
      getDefaultServerAcceptor(0).getParams().put(TransportConstants.KEYSTORE_PATH_PROP_NAME, "server-keystore.jks");
      getDefaultServerAcceptor(0).getParams().put(TransportConstants.KEYSTORE_PASSWORD_PROP_NAME, "securepass");
      getDefaultServerAcceptor(0).getParams().put(TransportConstants.WANT_CLIENT_AUTH_PROP_NAME, "true");
      getDefaultServerAcceptor(0).getParams().put(TransportConstants.NEED_CLIENT_AUTH_PROP_NAME, "true");
      getDefaultServerAcceptor(0).getParams().put(TransportConstants.TRUSTSTORE_PATH_PROP_NAME, "client-ca-truststore.jks");
      getDefaultServerAcceptor(0).getParams().put(TransportConstants.TRUSTSTORE_PASSWORD_PROP_NAME, "securepass");

      getServer(0).getConfiguration().setSecurityEnabled(true);
      getServer(0).getSecurityRepository().addMatch("#", Collections.singleton(
              new Role("programmers", true, true, true, true, true, true, true, true, true, true, false, false)));

      setupRouterServerWithDiscovery(0, KeyType.USER_NAME, MOCK_POLICY_NAME, null, true, null, 1);

      startServers(0);

      testConnection(TransportConstants.DEFAULT_HOST, TransportConstants.DEFAULT_PORT + 0,
              null, null, null, true, true, -1);

      assertEquals(1, keys.size());
      assertEquals("first", keys.get(0));
   }

   @TestTemplate
   public void testRoleNameKeyLocalTarget() throws Exception {
      ActiveMQJAASSecurityManager securityManager = new ActiveMQJAASSecurityManager("PropertiesLogin");
      servers[0] = addServer(ActiveMQServers.newActiveMQServer(createDefaultConfig(true).setSecurityEnabled(true), ManagementFactory.getPlatformMBeanServer(), securityManager, false));
      setupRouterServerWithLocalTarget(0, KeyType.ROLE_NAME, "b", "b");

      // ensure advisory permission is present for openwire connection creation by 'b'
      HierarchicalRepository<Set<Role>> securityRepository = servers[0].getSecurityRepository();
      Role role = new Role("b", true, true, true, true, true, true, false, false, true, true, false, false);
      Set<Role> roles = new HashSet<>();
      roles.add(role);
      securityRepository.addMatch("ActiveMQ.Advisory.#", roles);

      startServers(0);

      // expect disconnect/reject as not role b
      try {
         final int noRetriesSuchThatWeGetAnErrorOnRejection = 0;
         testConnection(TransportConstants.DEFAULT_HOST, TransportConstants.DEFAULT_PORT + 0,
                 null, "a", "a", false, false, noRetriesSuchThatWeGetAnErrorOnRejection);
         fail("Expect to be rejected as not in role b");
      } catch (Exception expectedButNotSpecificDueToDifferentProtocolsInPlay) {
      }

      testConnection(TransportConstants.DEFAULT_HOST, TransportConstants.DEFAULT_PORT + 0,
              null, "b", "b", false, false, -1);
   }

   private void testConnection(String host, int port, String clientID, String user, String password, boolean sslEnabled, boolean needClientAuth, int retries) throws Exception {
      if (MQTT_PROTOCOL.equals(protocol)) {
         for (int i = 0; retries == -1 || i <= retries; i++) {
            try {
               testMQTTConnection(host, port, clientID, user, password, sslEnabled, needClientAuth);
               break;
            } catch (Throwable t) {
               if (i == retries)  {
                  throw t;
               }
            }
         }
      } else {
         ConnectionFactory connectionFactory = createFactory(protocol, sslEnabled, host,
                 port, clientID, user, password, needClientAuth, retries);

         try (Connection connection = connectionFactory.createConnection()) {
            connection.start();
         }
      }
   }

   private void testMQTTConnection(String host, int port, String clientID, String user, String password, boolean sslEnabled, boolean needClientAuth) throws Exception {
      MqttConnectOptions connOpts = new MqttConnectOptions();
      connOpts.setCleanSession(true);
      connOpts.setUserName(user);
      if (password != null) {
         connOpts.setPassword(password.toCharArray());
      }

      String serverURIScheme = "tcp";

      if (sslEnabled) {
         serverURIScheme = "ssl";

         SSLContext sslContext = SSLContext.getInstance("TLS");

         KeyManager[] kms = null;
         if (needClientAuth) {
            KeyStore keyStore = KeyStore.getInstance("JKS");
            keyStore.load(getClass().getClassLoader().getResourceAsStream("client-keystore.jks"), "securepass".toCharArray());
            KeyManagerFactory kmf = KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm());
            kmf.init(keyStore, "securepass".toCharArray());
            kms = kmf.getKeyManagers();
         }

         KeyStore trustStore = KeyStore.getInstance("JKS");
         trustStore.load(getClass().getClassLoader().getResourceAsStream("server-ca-truststore.jks"), "securepass".toCharArray());
         TrustManagerFactory tmf = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm());
         tmf.init(trustStore);

         sslContext.init(kms, tmf.getTrustManagers(), null); // null KeyManagers, null TrustManagers, null SecureRandom

         SSLSocketFactory socketFactory = sslContext.getSocketFactory();
         connOpts.setSocketFactory(socketFactory);
      }

      if (clientID == null) {
         clientID = UUID.randomUUID().toString();
      }

      try (MqttClient mqttClient = new MqttClient(serverURIScheme + "://" + host + ":" + port, clientID, new MemoryPersistence())) {
         mqttClient.connect(connOpts);
         mqttClient.disconnect();
      }
   }

   private boolean checkLocalHostname(String host) {
      try {
         return InetAddress.getByName(host).isLoopbackAddress();
      } catch (UnknownHostException ignore) {
         return false;
      }
   }
}
