/*
 * 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.kafka.common.security.authenticator;

import org.apache.kafka.common.config.SaslConfigs;
import org.apache.kafka.common.config.types.Password;
import org.apache.kafka.common.network.ListenerName;
import org.apache.kafka.common.security.JaasContext;
import org.apache.kafka.common.security.auth.AuthenticateCallbackHandler;
import org.apache.kafka.common.security.auth.Login;
import org.apache.kafka.common.security.plain.PlainLoginModule;
import org.apache.kafka.common.utils.Utils;

import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.mockito.MockedStatic;
import org.mockito.Mockito;

import java.util.Collections;
import java.util.HashMap;
import java.util.Map;

import javax.security.auth.login.LoginException;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNotSame;
import static org.junit.jupiter.api.Assertions.assertSame;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.mockStatic;
import static org.mockito.Mockito.verify;

public class LoginManagerTest {

    private Password dynamicPlainContext;
    private Password dynamicDigestContext;

    @BeforeEach
    public void setUp() {
        dynamicPlainContext = new Password(PlainLoginModule.class.getName() +
                " required user=\"plainuser\" password=\"plain-secret\";");
        dynamicDigestContext = new Password(TestDigestLoginModule.class.getName() +
                " required user=\"digestuser\" password=\"digest-secret\";");
        TestJaasConfig.createConfiguration("SCRAM-SHA-256",
                Collections.singletonList("SCRAM-SHA-256"));
    }

    @AfterEach
    public void tearDown() {
        LoginManager.closeAll();
    }

    @Test
    public void testClientLoginManager() throws Exception {
        Map<String, ?> configs = Collections.singletonMap("sasl.jaas.config", dynamicPlainContext);
        JaasContext dynamicContext = JaasContext.loadClientContext(configs);
        JaasContext staticContext = JaasContext.loadClientContext(Collections.emptyMap());

        LoginManager dynamicLogin = LoginManager.acquireLoginManager(dynamicContext, "PLAIN",
                DefaultLogin.class, configs);
        assertEquals(dynamicPlainContext, dynamicLogin.cacheKey());
        LoginManager staticLogin = LoginManager.acquireLoginManager(staticContext, "SCRAM-SHA-256",
                DefaultLogin.class, configs);
        assertNotSame(dynamicLogin, staticLogin);
        assertEquals("KafkaClient", staticLogin.cacheKey());

        assertSame(dynamicLogin, LoginManager.acquireLoginManager(dynamicContext, "PLAIN",
                DefaultLogin.class, configs));
        assertSame(staticLogin, LoginManager.acquireLoginManager(staticContext, "SCRAM-SHA-256",
                DefaultLogin.class, configs));

        verifyLoginManagerRelease(dynamicLogin, 2, dynamicContext, configs);
        verifyLoginManagerRelease(staticLogin, 2, staticContext, configs);
    }

    @Test
    public void testServerLoginManager() throws Exception {
        Map<String, Object> configs = new HashMap<>();
        configs.put("plain.sasl.jaas.config", dynamicPlainContext);
        configs.put("digest-md5.sasl.jaas.config", dynamicDigestContext);
        ListenerName listenerName = new ListenerName("listener1");
        JaasContext plainJaasContext = JaasContext.loadServerContext(listenerName, "PLAIN", configs);
        JaasContext digestJaasContext = JaasContext.loadServerContext(listenerName, "DIGEST-MD5", configs);
        JaasContext scramJaasContext = JaasContext.loadServerContext(listenerName, "SCRAM-SHA-256", configs);

        LoginManager dynamicPlainLogin = LoginManager.acquireLoginManager(plainJaasContext, "PLAIN",
                DefaultLogin.class, configs);
        assertEquals(dynamicPlainContext, dynamicPlainLogin.cacheKey());
        LoginManager dynamicDigestLogin = LoginManager.acquireLoginManager(digestJaasContext, "DIGEST-MD5",
                DefaultLogin.class, configs);
        assertNotSame(dynamicPlainLogin, dynamicDigestLogin);
        assertEquals(dynamicDigestContext, dynamicDigestLogin.cacheKey());
        LoginManager staticScramLogin = LoginManager.acquireLoginManager(scramJaasContext, "SCRAM-SHA-256",
                DefaultLogin.class, configs);
        assertNotSame(dynamicPlainLogin, staticScramLogin);
        assertEquals("KafkaServer", staticScramLogin.cacheKey());

        assertSame(dynamicPlainLogin, LoginManager.acquireLoginManager(plainJaasContext, "PLAIN",
                DefaultLogin.class, configs));
        assertSame(dynamicDigestLogin, LoginManager.acquireLoginManager(digestJaasContext, "DIGEST-MD5",
                DefaultLogin.class, configs));
        assertSame(staticScramLogin, LoginManager.acquireLoginManager(scramJaasContext, "SCRAM-SHA-256",
                DefaultLogin.class, configs));

        verifyLoginManagerRelease(dynamicPlainLogin, 2, plainJaasContext, configs);
        verifyLoginManagerRelease(dynamicDigestLogin, 2, digestJaasContext, configs);
        verifyLoginManagerRelease(staticScramLogin, 2, scramJaasContext, configs);
    }

    @Test
    public void testLoginManagerWithDifferentConfigs() throws Exception {
        Map<String, Object> configs1 = new HashMap<>();
        configs1.put("sasl.jaas.config", dynamicPlainContext);
        configs1.put("client.id", "client");
        configs1.put(SaslConfigs.SASL_OAUTHBEARER_TOKEN_ENDPOINT_URL, "http://host1:1234");
        Map<String, Object> configs2 = new HashMap<>(configs1);
        configs2.put(SaslConfigs.SASL_OAUTHBEARER_TOKEN_ENDPOINT_URL, "http://host2:1234");
        JaasContext dynamicContext = JaasContext.loadClientContext(configs1);
        Map<String, Object> configs3 = new HashMap<>(configs1);
        configs3.put("client.id", "client3");

        LoginManager dynamicLogin1 = LoginManager.acquireLoginManager(dynamicContext, "PLAIN", DefaultLogin.class, configs1);
        LoginManager dynamicLogin2 = LoginManager.acquireLoginManager(dynamicContext, "PLAIN", DefaultLogin.class, configs2);

        assertEquals(dynamicPlainContext, dynamicLogin1.cacheKey());
        assertEquals(dynamicPlainContext, dynamicLogin2.cacheKey());
        assertNotSame(dynamicLogin1, dynamicLogin2);

        assertSame(dynamicLogin1, LoginManager.acquireLoginManager(dynamicContext, "PLAIN", DefaultLogin.class, configs1));
        assertSame(dynamicLogin2, LoginManager.acquireLoginManager(dynamicContext, "PLAIN", DefaultLogin.class, configs2));
        assertSame(dynamicLogin1, LoginManager.acquireLoginManager(dynamicContext, "PLAIN", DefaultLogin.class, new HashMap<>(configs1)));
        assertSame(dynamicLogin1, LoginManager.acquireLoginManager(dynamicContext, "PLAIN", DefaultLogin.class, configs3));


        JaasContext staticContext = JaasContext.loadClientContext(Collections.emptyMap());
        LoginManager staticLogin1 = LoginManager.acquireLoginManager(staticContext, "SCRAM-SHA-256", DefaultLogin.class, configs1);
        LoginManager staticLogin2 = LoginManager.acquireLoginManager(staticContext, "SCRAM-SHA-256", DefaultLogin.class, configs2);
        assertNotSame(staticLogin1, dynamicLogin1);
        assertNotSame(staticLogin2, dynamicLogin2);
        assertNotSame(staticLogin1, staticLogin2);
        assertSame(staticLogin1, LoginManager.acquireLoginManager(staticContext, "SCRAM-SHA-256", DefaultLogin.class, configs1));
        assertSame(staticLogin2, LoginManager.acquireLoginManager(staticContext, "SCRAM-SHA-256", DefaultLogin.class, configs2));
        assertSame(staticLogin1, LoginManager.acquireLoginManager(staticContext, "SCRAM-SHA-256", DefaultLogin.class, new HashMap<>(configs1)));
        assertSame(staticLogin1, LoginManager.acquireLoginManager(staticContext, "SCRAM-SHA-256", DefaultLogin.class, configs3));

        verifyLoginManagerRelease(dynamicLogin1, 4, dynamicContext, configs1);
        verifyLoginManagerRelease(dynamicLogin2, 2, dynamicContext, configs2);
        verifyLoginManagerRelease(staticLogin1, 4, staticContext, configs1);
        verifyLoginManagerRelease(staticLogin2, 2, staticContext, configs2);
    }

    @Test
    public void testShouldReThrowExceptionOnErrorLoginAttempt() throws Exception {
        Map<String, Object> config = new HashMap<>();
        config.put(SaslConfigs.SASL_JAAS_CONFIG, dynamicPlainContext);
        config.put(SaslConfigs.SASL_LOGIN_CLASS, Login.class);
        config.put(SaslConfigs.SASL_LOGIN_CALLBACK_HANDLER_CLASS, AuthenticateCallbackHandler.class);
        JaasContext dynamicContext = JaasContext.loadClientContext(config);

        Login mockLogin = mock(Login.class);
        AuthenticateCallbackHandler mockHandler = mock(AuthenticateCallbackHandler.class);

        doThrow(new LoginException("Expecting LoginException")).when(mockLogin).login();

        try (MockedStatic<Utils> mockedUtils = mockStatic(Utils.class, Mockito.CALLS_REAL_METHODS)) {
            mockedUtils.when(() -> Utils.newInstance(Login.class)).thenReturn(mockLogin);
            mockedUtils.when(() -> Utils.newInstance(AuthenticateCallbackHandler.class)).thenReturn(mockHandler);

            assertThrows(LoginException.class, () ->
                    LoginManager.acquireLoginManager(dynamicContext, "PLAIN", DefaultLogin.class, config)
            );

            verify(mockLogin).close();
            verify(mockHandler).close();
        }
    }

    private void verifyLoginManagerRelease(LoginManager loginManager, int acquireCount, JaasContext jaasContext,
                                           Map<String, ?> configs) throws Exception {

        // Release all except one reference and verify that the loginManager is still cached
        for (int i = 0; i < acquireCount - 1; i++)
            loginManager.release();
        assertSame(loginManager, LoginManager.acquireLoginManager(jaasContext, "PLAIN",
                DefaultLogin.class, configs));

        // Release all references and verify that new LoginManager is created on next acquire
        for (int i = 0; i < 2; i++) // release all references
            loginManager.release();
        LoginManager newLoginManager = LoginManager.acquireLoginManager(jaasContext, "PLAIN",
                DefaultLogin.class, configs);
        assertNotSame(loginManager, newLoginManager);
        newLoginManager.release();
    }
}
