package com.atlassian.bitbucket.mesh.auth;

import com.atlassian.bitbucket.mesh.config.ApplicationConfig;
import com.atlassian.bitbucket.mesh.config.AuthenticationConfig;
import com.atlassian.bitbucket.mesh.config.NodeConfig;
import com.atlassian.bitbucket.mesh.junit.LogListener;
import com.atlassian.bitbucket.mesh.management.SidebandManager;
import com.atlassian.bitbucket.mesh.network.MeshManager;
import com.atlassian.bitbucket.mesh.network.MeshNode;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcNodeKeyRequest;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcSidebandFollowerFragment;
import com.atlassian.bitbucket.mesh.util.KeyUtils;
import io.jsonwebtoken.SignatureAlgorithm;
import io.jsonwebtoken.security.Keys;
import org.junit.Before;
import org.junit.BeforeClass;
import org.junit.Rule;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.Mock;
import org.mockito.junit.MockitoJUnitRunner;

import java.security.PublicKey;
import java.time.Duration;
import java.time.temporal.ChronoUnit;

import static com.atlassian.bitbucket.mesh.rpc.RpcConstants.Token.NID_CONTROL_PLANE;
import static org.junit.Assert.*;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.*;

@RunWith(MockitoJUnitRunner.class)
public class JwtAuthenticatorTest {

    private final String NODE_ID = "1";
    private static final String RSA_DER = "MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAmrPKDLXmM/Iyt8JI4BwF56AAut" +
            "g+rWHI4iOI/8G2bzw4YdODwpCzlmmZQXbak6Tq10EamAJVARIN30i7CvyO2WIrDTYZcyLodQA9pMXv/Y7AyZlsAhuojUlkaIIRHK" +
            "IxGPy9bZ7ZzKVso9JfMozehiz50fyu4pGv0k3QqSmmG866Eq0jUDsHgFeNP8awK5PozDgT7HWHaR7J45X8HdI08NlKtA6fjwjRPm" +
            "sOLvDJdKJFgNAvQz0LwqLvaZ6Lhrh7+30D070ECjUTi3wjlC3bQKthdOTa/TZVYvv2k/tYAJMYa/Rt7EdzV8J4oKbzUg27v6XE/l" +
            "FvcLkmQeTgF8gpnwIDAQAB";
    private static final String RSA_FINGERPRINT = "hJAhBrb+Pf+MUpcPm5i3R/izb2eb6lyEDLySeCkxCWo";

    private static SigningKey signingKey;

    @Rule
    public final LogListener logListener = new LogListener(JwtAuthenticator.class);

    @Mock
    private AuthenticationConfig authenticationConfig;
    @Mock
    private AuthenticationManager authenticationManager;
    private JwtAuthenticator authenticator;
    @Mock
    private ApplicationConfig config;
    @Mock
    private PublicKeyRegistry keyRegistry;
    @Mock
    private MeshManager meshManager;
    @Mock
    private NodeConfig nodeConfig;
    @Mock
    private SidebandManager sidebandManager;

    @BeforeClass
    public static void setUpKeys() {
        signingKey = new SimpleSigningKey(Keys.keyPairFor(SignatureAlgorithm.RS256));
    }

    @Before
    public void setUp() throws Exception {
        when(authenticationManager.getSigningKey()).thenReturn(signingKey);

        when(config.getAuthentication()).thenReturn(authenticationConfig);
        when(authenticationConfig.getKeyRequestInterval()).thenReturn(Duration.of(2, ChronoUnit.SECONDS));
        when(config.getNode()).thenReturn(nodeConfig);

        when(nodeConfig.getId()).thenReturn(NODE_ID);

        when(keyRegistry.findByFingerprint(eq(NODE_ID), eq(signingKey.getFingerprint())))
                .thenReturn(signingKey.getPublic());

        authenticator = new JwtAuthenticator(authenticationManager, config, keyRegistry, meshManager, sidebandManager);
    }

    @Test
    public void testAuthenticateAllowsTokenFromControlPlane() throws Exception {
        registerPublicKey(NID_CONTROL_PLANE);

        //This token has no "nid" header, which should mean it defaults to "control-plane"
        String token = "eyJmcCI6ImhKQWhCcmIrUGYrTVVwY1BtNWkzUi9pemIyZWI2bHlFREx5U2VDa3hDV28iLCJhbGciOiJSUzI1NiJ9." +
                "eyJpYXQiOjE2NDUzMjQ0MTksInN1YiI6Ikp3dEF1dGhlbnRpY2F0b3JUZXN0In0.gKiXjWVzZmOvsf528AXjWmmcXMwcip55" +
                "ew-pH20YwNKWetEhwXKNDOyXLJnW3dUZdGcyCzE27zgWftuYgpW_4Yizw_jT6mgMAGZVs5l5nFb9c1BV1q78-BqyAJlkiAbQ" +
                "hCz3vMnxNLQWIYLsC6ZGL1TYXoH7Rt9VuS45ygxE1I52JQTCWmiDtU0Tj9s7Vk-hDCdM6YpXz6yss9oEGDKD6TTE8GoFHCBA" +
                "40dSQYeAX75tpb4zIIjGjYdXvNuLAdZnRi2K8If9TbAKxZFcIwUBBp6JI6s--uICM8c7QhQOKhOXtQhWBvMRbJs_ZS7NwCNF" +
                "xClstoCPQvks9JAemFczgQ";

        Principal principal = authenticator.authenticate(getClass().getSimpleName(), token);
        assertNotNull(principal);
        assertNull(principal.getId());
        assertSame(ApplicationRole.CONTROL_PLANE, principal.getRole());

        verify(keyRegistry).findByFingerprint(eq(NID_CONTROL_PLANE), eq(RSA_FINGERPRINT));
    }

    @Test
    public void testAuthenticateRejectsNonsenseToken() {
        assertNull(authenticator.authenticate("Fizz", "Buzz"));

        logListener.assertWarn("Rejecting invalid token");
    }

    @Test
    public void testAuthenticateRejectsTokenAfterExpiry() throws Exception {
        registerPublicKey(NODE_ID);

        String token = "eyJmcCI6ImhKQWhCcmIrUGYrTVVwY1BtNWkzUi9pemIyZWI2bHlFREx5U2VDa3hDV28iLCJuaWQiOiIxIiwiYWxnI" +
                "joiUlMyNTYifQ.eyJpYXQiOjE2NDUzMjMyNzEsInN1YiI6Ikp3dEF1dGhlbnRpY2F0b3JUZXN0IiwiZXhwIjoxNjQ1MzIzMz" +
                "AxfQ.LK2q60SysH-uRe-gXREGWnSBVmpVVY4wKms9yi3r0c4iNhe4GknXfuFAd6MYQ5hz6SZRjBNmLcG50yJJnWUW0GcAZS3" +
                "FynYQCkwbgHfHLfL1vFO5JRWi6BgMaK6a1rdWzT05esEFW65Fwj_7fNRfm3W_C5Fd6N1KAOat1UrOuU9Q2KePMgYz6xR3Pys" +
                "8GNM6WAvjdS5mMwTMDwqq6d7zSI5mi3NOtnyUi_qyehpotBwOvzKzFgLk1d5-QfvggEIvh2T9gcw3aE1m8jxxYyZV9LpXIEO" +
                "uW2Uz0fjqzG-jmtEg2WoHboSKpTj-AZt_PR5hGPefKPNbLZDZZvI7WLEvmQ";
        assertNull(authenticator.authenticate(getClass().getSimpleName(), token));

        logListener.assertWarn("expired token");
    }

    @Test
    public void testAuthenticateRejectsTokenWithIncorrectSubject() {
        String token = authenticator.sign("Fizz");
        assertNull(authenticator.authenticate("Buzz", token));

        logListener.assertWarn("Rejecting Fizz token from 1; expected Buzz");
    }

    @Test
    public void testAuthenticateRejectsTokenWithNoMatchingPublicKeyFromUnknownNode() {
        reset(keyRegistry);
        when(meshManager.findNodeById(NODE_ID)).thenReturn(null);

        String token = authenticator.sign("Fizz");
        assertNull(authenticator.authenticate("Fizz", token));

        verify(keyRegistry).findByFingerprint(eq(NODE_ID), eq(signingKey.getFingerprint()));
        verifyNoInteractions(sidebandManager);

        logListener.assertWarn("signature verification failed");
    }

    @Test
    public void testAuthenticateRequestsMissingKeyForNoMatchingPublicKeyFromKnownNode() {
        reset(keyRegistry);
        when(meshManager.findNodeById(NODE_ID)).thenReturn(new MeshNode.Builder(NODE_ID)
                .rpcUrl("127.0.0.1/rpc")
                .build());

        String token = authenticator.sign("Fizz");
        assertNull(authenticator.authenticate("Fizz", token));

        verify(keyRegistry).findByFingerprint(eq(NODE_ID), eq(signingKey.getFingerprint()));
        RpcNodeKeyRequest expectedNodeKeyRequest = RpcNodeKeyRequest.newBuilder()
                .setNodeId(NODE_ID)
                .setFingerprint(signingKey.getFingerprint())
                .build();
        RpcSidebandFollowerFragment expectedSidebandFollowerFragment = RpcSidebandFollowerFragment.newBuilder()
                .setNodeKeyRequest(expectedNodeKeyRequest).build();
        verify(sidebandManager).send(expectedSidebandFollowerFragment);

        logListener.assertWarn("signature verification failed");
    }

    @Test
    public void testAuthenticateDoesNotRequestMissingKeyForNoMatchingPublicKeyFromControlPlane() {
        reset(keyRegistry);
        when(nodeConfig.getId()).thenReturn(NID_CONTROL_PLANE);

        String token = authenticator.sign("Fizz");
        assertNull(authenticator.authenticate("Fizz", token));

        verify(keyRegistry).findByFingerprint(eq(NID_CONTROL_PLANE), eq(signingKey.getFingerprint()));
        verifyNoInteractions(sidebandManager);

        logListener.assertWarn("signature verification failed");
    }

    @Test
    public void testAuthenticateRateLimitsMissingPublicKeyRequests() throws InterruptedException {
        reset(keyRegistry);
        when(meshManager.findNodeById(NODE_ID)).thenReturn(new MeshNode.Builder(NODE_ID)
                .rpcUrl("127.0.0.1/rpc")
                .build());
        String nodeId2 = "2";
        when(meshManager.findNodeById(nodeId2)).thenReturn(new MeshNode.Builder(nodeId2)
                .rpcUrl("127.0.0.1/rpc")
                .build());
        String token1 = authenticator.sign("Fizz");
        when(nodeConfig.getId()).thenReturn(nodeId2);
        String token2 = authenticator.sign("Bar");

        assertNull(authenticator.authenticate("Fizz", token1));
        assertNull(authenticator.authenticate("Fizz", token1));
        assertNull(authenticator.authenticate("Bar", token2));
        assertNull(authenticator.authenticate("Bar", token2));
        verify(sidebandManager, times(1))
                .send(RpcSidebandFollowerFragment.newBuilder()
                        .setNodeKeyRequest(RpcNodeKeyRequest.newBuilder()
                                .setNodeId(NODE_ID)
                                .setFingerprint(signingKey.getFingerprint())
                                .build())
                        .build());
        verify(sidebandManager, times(1))
                .send(RpcSidebandFollowerFragment.newBuilder()
                        .setNodeKeyRequest(RpcNodeKeyRequest.newBuilder()
                                .setNodeId(nodeId2)
                                .setFingerprint(signingKey.getFingerprint())
                                .build())
                        .build());

        reset(sidebandManager);
        Thread.sleep(authenticationConfig.getKeyRequestInterval().plusMillis(100).toMillis());

        assertNull(authenticator.authenticate("Fizz", token1));
        verify(sidebandManager, times(1))
                .send(RpcSidebandFollowerFragment.newBuilder()
                        .setNodeKeyRequest(RpcNodeKeyRequest.newBuilder()
                                .setNodeId(NODE_ID)
                                .setFingerprint(signingKey.getFingerprint())
                                .build())
                        .build());
    }

    @Test
    public void testAuthenticateRejectsTokenWithNoSubject() throws Exception {
        registerPublicKey(NODE_ID);

        String token = "eyJmcCI6ImhKQWhCcmIrUGYrTVVwY1BtNWkzUi9pemIyZWI2bHlFREx5U2VDa3hDV28iLCJuaWQiOiIxIiwiYWxnI" +
                "joiUlMyNTYifQ.eyJpYXQiOjE2NDUzMjM3NTR9.NzQ-XF-lgHNfs6M7jZI-FLCcZPl9D9odQAd_XL_yyjRpFZMqQ2_GQOWo7" +
                "nWP-0uxhQSHDhzlsxnfqW07tvMnN8ohii0G-ABsQsW6P5Cd55_6bdgOIhQctl0Kvy4S9lkT7h9ZLluiuUEndl67zltw3HwKI" +
                "IS1tsiqaHF8EXLh-ffhulUZZiqgshynLwq--3gFQ26AZo_1N418OpB6YDH4GKk6aSLrT2bUO6TlDnDrrGSxieVYySmzQXoLm" +
                "bsTb1ku8yB-VlfHqrWxwtsUrwpby0et9t8ddTiL3KFxchtAhTm7aKRczqmIlo-KGZJ18M7ZjFUvFnKZUV2LJmFeqwp6dQ";
        assertNull(authenticator.authenticate(getClass().getSimpleName(), token));

        logListener.assertWarn("no subject was specified");
    }

    @Test
    public void testSignAndAuthenticate() {
        when(authenticationConfig.getExpiryInterval()).thenReturn(Duration.ofSeconds(5L));
        when(authenticationConfig.isExpiryEnabled()).thenReturn(true);

        String subject = getClass().getSimpleName();
        String token = authenticator.sign(subject);

        Principal principal = authenticator.authenticate(getClass().getSimpleName(), token);
        assertNotNull(principal);
        assertEquals(NODE_ID, principal.getId());
        assertSame(ApplicationRole.PEER, principal.getRole());
    }

    private void registerPublicKey(String nodeId) throws Exception {
        PublicKey publicKey = KeyUtils.parsePublicKeyFromDer(RSA_DER);

        reset(keyRegistry);
        when(keyRegistry.findByFingerprint(eq(nodeId), eq(RSA_FINGERPRINT)))
                .thenReturn(publicKey);
    }
}