package com.atlassian.bitbucket.mesh.auth;

import com.atlassian.bitbucket.mesh.execution.GrpcExecutionManager;
import com.atlassian.bitbucket.mesh.git.exception.FailedPreconditionException;
import com.atlassian.bitbucket.mesh.io.MoreFiles;
import com.atlassian.bitbucket.mesh.junit.LogListener;
import com.atlassian.bitbucket.mesh.junit.TemporaryDirectory;
import com.atlassian.bitbucket.mesh.network.MeshNode;
import com.atlassian.bitbucket.mesh.network.MeshNodeRemovedEvent;
import com.atlassian.bitbucket.mesh.test.ExecutionContextRule;
import com.atlassian.bitbucket.mesh.util.KeyUtils;
import io.jsonwebtoken.SignatureAlgorithm;
import io.jsonwebtoken.security.Keys;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.ArgumentCaptor;
import org.mockito.InOrder;
import org.mockito.Mock;
import org.mockito.junit.MockitoJUnitRunner;
import org.springframework.beans.DirectFieldAccessor;
import org.springframework.scheduling.TaskScheduler;

import java.nio.file.Files;
import java.nio.file.Path;
import java.security.KeyPair;
import java.security.PublicKey;
import java.time.Duration;
import java.time.Instant;
import java.util.Optional;

import static com.atlassian.bitbucket.mesh.auth.DefaultAuthenticationManager.PEM_CONTROL_PLANE;
import static com.atlassian.bitbucket.mesh.auth.DefaultAuthenticationManager.PEM_SIGNING_KEY;
import static com.atlassian.bitbucket.mesh.rpc.RpcConstants.Token.NID_CONTROL_PLANE;
import static com.atlassian.bitbucket.mesh.rpc.RpcConstants.Token.NID_SIDECAR;
import static com.atlassian.bitbucket.mesh.test.MockitoUtils.anyOneOf;
import static com.atlassian.bitbucket.mesh.test.MockitoUtils.iterableOf;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.containsString;
import static org.junit.Assert.*;
import static org.mockito.Mockito.*;

@RunWith(MockitoJUnitRunner.class)
public class DefaultAuthenticationManagerTest {

    //This is base64-encoded binary containing an EC 256-bit public key
    private static final String EC_DER = "MFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEaucvQzbD4ZfCN+h8xuyCxFDa9tzRU2wikaQ" +
            "cuXS8gR/bIyzejKGQ3ylARv1oV4Pc5JWhOh4i+LB040JALAmZEw";
    //This is base64-encoded binary containing an RSA 2048-bit public key
    private static final String RSA_DER = "MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAspADKDS2sRRXyJaZ7l10aiftkP" +
            "OqYj2Qu/R7fU851GrSs0egLd0w50hNnWiV+5FOt8XScI7HbisyNE+AiPVb3oEozjGMxzU2EMH9vR6hFH3y20taEcGrMarc8K/0+s" +
            "12GXaDcn2DVSBZ9VMv3qdyorFFXfT4QwvtInhydpbHsO+VtkwC+7/0zZ46XljzyZhSz7oK34k0W39JO40GBm8fJe2TLiBg7te+8Q" +
            "sl0apsZK14TSZidM9PSI6q71HSWpU2A43Iaoue5X/eH6RFxtBByfpe/iAKZDkeMyHuAM6b475oi8allRTF19gXXaufbRtyRb5BaE" +
            "oYKNkKu4w65bp57QIDAQAB";

    @Rule
    public final ExecutionContextRule executionContextRule = new ExecutionContextRule(new GrpcExecutionManager());
    @Rule
    public final LogListener logListener = new LogListener(DefaultAuthenticationManager.class);
    @Rule
    public final TemporaryDirectory tempDir = new TemporaryDirectory();

    private Path configDir;
    @Mock
    private PublicKeyRegistry keyRegistry;
    private DefaultAuthenticationManager manager;
    @Mock
    private TaskScheduler taskScheduler;

    @Before
    public void setUp() throws Exception {
        configDir = tempDir.newFolder("config");

        manager = new DefaultAuthenticationManager(configDir, executionContextRule.getExecutionManager(),
                keyRegistry, taskScheduler);
    }

    @Test
    public void testGetPrincipalWithoutContextReturnsEmpty() {
        assertEquals(Optional.empty(), manager.getPrincipal());
    }

    @Test
    public void testGetPrincipalWithoutContextThrowsIfRequested() {
        assertThrows(FailedPreconditionException.class, manager::getPrincipalOrThrow);
    }

    @Test
    public void testIsProvisioned() {
        assertFalse(manager.isProvisioned());

        when(keyRegistry.isRegistered(eq(NID_CONTROL_PLANE))).thenReturn(true);
        assertTrue(manager.isProvisioned());
    }

    @Test
    public void testOnMeshNodeRemoved() {
        MeshNode node = mock(MeshNode.class);
        when(node.getId()).thenReturn("id");

        when(keyRegistry.removeAll(eq("id"))).thenReturn(true);

        manager.onMeshNodeRemoved(new MeshNodeRemovedEvent(node));

        verify(keyRegistry).removeAll(eq("id"));

        logListener.assertDebug("id: Removed all registered public keys");
    }

    @Test
    public void testProvision() throws Exception {
        PublicKey signingPublic = mock(PublicKey.class);

        SigningKey signingKey = mock(SigningKey.class);
        when(signingKey.getPublic()).thenReturn(signingPublic);

        //To keep this test focused on provisioning, we mock having initialized the signing key
        new DirectFieldAccessor(manager).setPropertyValue("signingKey", signingKey);

        PublicKey controlKey = KeyUtils.parsePublicKeyFromDer(RSA_DER);
        PublicKey nodeKey = manager.provision(controlKey);
        assertSame(signingPublic, nodeKey);

        Path pem = configDir.resolve(PEM_CONTROL_PLANE);
        assertEquals(controlKey, KeyUtils.readPublicKey(pem));

        logListener.assertInfo("Registered new control plane key");
    }

    @Test
    public void testProvisionDoesNotRegisterKeyIfWritingFails() throws Exception {
        MoreFiles.deleteRecursively(configDir); //Delete the directory, which will cause writing to fail

        FailedPreconditionException thrown = assertThrows(FailedPreconditionException.class,
                () -> manager.provision(KeyUtils.parsePublicKeyFromDer(RSA_DER)));
        assertThat(thrown.getMessage(), containsString("Failed to write"));

        verifyNoInteractions(keyRegistry); //Registration should only happen if written successfully
    }

    @Test
    public void testGetSigningKey() throws Exception {
        //Generate a key pair and write it to disk. We could "just" call prepareKeys and let it generate
        //a key pair itself, but this allows us to validate the round trip
        Path pem = configDir.resolve(PEM_SIGNING_KEY);
        KeyPair signingPair = Keys.keyPairFor(SignatureAlgorithm.RS256);
        KeyUtils.writeKeyPair(pem, signingPair);

        //Load the key pair
        manager.prepareKeys();

        //Verify the signing key accurately round-trips the generated key pair.
        SigningKey signingKey = manager.getSigningKey();
        assertSigningKey(signingPair, signingKey);
    }

    @Test
    public void testPrepareKeysGeneratesSigningKey() throws Exception {
        Path pem = configDir.resolve(PEM_SIGNING_KEY);
        assertFalse(Files.exists(pem));

        manager.prepareKeys();

        verifyNoInteractions(keyRegistry); //No control plane key, so nothing should be registered

        //Read the key pair from disk and compare it to the signing key
        KeyPair signingPair = KeyUtils.readKeyPair(pem);

        SigningKey signingKey = manager.getSigningKey();
        assertSigningKey(signingPair, signingKey);

        logListener.assertInfo("No control plane key");
        logListener.assertInfo("Generated signing key");
    }

    @Test
    public void testPrepareKeysLoadsControlPlaneAndSigningKeysAndRegistersControlPlaneKey() throws Exception {
        Path controlPem = configDir.resolve(PEM_CONTROL_PLANE);
        PublicKey controlKey = KeyUtils.parsePublicKeyFromDer(RSA_DER);
        KeyUtils.writePublicKey(controlPem, controlKey);

        Path signingPem = configDir.resolve(PEM_SIGNING_KEY);
        KeyPair signingPair = Keys.keyPairFor(SignatureAlgorithm.RS256);
        KeyUtils.writeKeyPair(signingPem, signingPair);

        manager.prepareKeys();

        verify(keyRegistry).add(eq(NID_CONTROL_PLANE), eq(controlKey));

        logListener.assertDebug("Registered provisioned control plane key");
        logListener.assertDebug("Loaded signing key");
    }

    @Test
    public void testUpdatePublicKeysAllowsEmpty() {
        manager.updatePublicKeys(new UpdatePublicKeysParameters.Builder().build());

        logListener.assertDebug("No updated control plane keys");
    }

    @Test
    public void testUpdatePublicKeysUsesAddWhenIncomplete() throws Exception {
        Path pem = configDir.resolve(PEM_CONTROL_PLANE);
        assertFalse("A control plane key already exists on disk", Files.exists(pem));

        PublicKey controlEc = KeyUtils.parsePublicKeyFromDer(EC_DER);

        KeyPair node1New = Keys.keyPairFor(SignatureAlgorithm.RS256);
        PublicKey node1Old = KeyUtils.parsePublicKeyFromDer(RSA_DER);

        KeyPair node2 = Keys.keyPairFor(SignatureAlgorithm.RS256);

        Instant now = Instant.now(); //Capture this before calling updatePublicKeys so we can use it later
        manager.updatePublicKeys(new UpdatePublicKeysParameters.Builder()
                .controlPlaneKey(ExpirablePublicKey.of(controlEc)) //Should be written to disk
                .nodeKey("1", new ExpirablePublicKey.Builder(node1Old)
                        .expiryInterval(Duration.ofMinutes(1L))
                        .build())
                .nodeKey("1", ExpirablePublicKey.of(node1New.getPublic()))
                .nodeKey("2", ExpirablePublicKey.of(node2.getPublic()))
                .sidecarKey(ExpirablePublicKey.of(node1Old))
                .build());

        ArgumentCaptor<Runnable> taskCaptor = ArgumentCaptor.forClass(Runnable.class);

        InOrder ordered = inOrder(keyRegistry, taskScheduler);
        ordered.verify(keyRegistry).add(eq(NID_CONTROL_PLANE), eq(controlEc));
        ordered.verify(keyRegistry).add(eq(NID_SIDECAR), eq(node1Old));
        //times(2) plus anyOneOf allows old and new to appear in any order as long as they're sequential
        ordered.verify(keyRegistry, times(2)).add(eq("1"), argThat(anyOneOf(node1Old, node1New.getPublic())));
        ordered.verify(taskScheduler).schedule(taskCaptor.capture(),
                //The task should be scheduled at least one minute in the future. >= is used to make the
                //test a little more forgiving
                argThat((Instant then) -> Duration.between(now, then).toMinutes() >= 1));
        ordered.verify(keyRegistry).add(eq("2"), eq(node2.getPublic()));
        ordered.verifyNoMoreInteractions();

        //Verify the task we captured is to remove the old key for node 1
        taskCaptor.getValue().run();
        verify(keyRegistry).remove(eq("1"), eq(node1Old));

        //This test verifies persistence roundtripping for an EC control plane key
        PublicKey controlKey = KeyUtils.readPublicKey(pem);
        assertNotNull("The control plane key should have been written to disk", controlKey);
        assertEquals("The wrong control plane key was persisted", controlEc, controlKey);
    }

    @Test
    public void testUpdatePublicKeysUsesSetWhenComplete() throws Exception {
        Path pem = configDir.resolve(PEM_CONTROL_PLANE);
        assertFalse("A control plane key already exists on disk", Files.exists(pem));

        PublicKey controlEc = KeyUtils.parsePublicKeyFromDer(EC_DER);
        PublicKey controlRsa = KeyUtils.parsePublicKeyFromDer(RSA_DER);

        KeyPair node1 = Keys.keyPairFor(SignatureAlgorithm.RS256);
        KeyPair node2 = Keys.keyPairFor(SignatureAlgorithm.RS256);

        Instant now = Instant.now(); //Capture this before calling updatePublicKeys so we can use it later
        manager.updatePublicKeys(new UpdatePublicKeysParameters.Builder()
                .complete(true) //Keys should be registered using set instead of add
                .controlPlaneKey(new ExpirablePublicKey.Builder(controlEc)
                        .expiryInterval(Duration.ofMinutes(1L))
                        .build())
                .controlPlaneKey(ExpirablePublicKey.of(controlRsa)) //Should be written to disk
                .nodeKey("1", ExpirablePublicKey.of(node1.getPublic()))
                .nodeKey("2", ExpirablePublicKey.of(node2.getPublic()))
                .sidecarKey(ExpirablePublicKey.of(controlEc))
                .build());

        ArgumentCaptor<Runnable> taskCaptor = ArgumentCaptor.forClass(Runnable.class);

        InOrder ordered = inOrder(keyRegistry, taskScheduler);
        ordered.verify(keyRegistry).set(eq(NID_CONTROL_PLANE), argThat(iterableOf(controlEc, controlRsa)));
        ordered.verify(taskScheduler).schedule(taskCaptor.capture(),
                //The task should be scheduled at least one minute in the future. >= is used to make the
                //test a little more forgiving
                argThat((Instant then) -> Duration.between(now, then).toMinutes() >= 1));
        ordered.verify(keyRegistry).set(eq(NID_SIDECAR), argThat(iterableOf(controlEc)));
        ordered.verify(keyRegistry).set(eq("1"), argThat(iterableOf(node1.getPublic())));
        ordered.verify(keyRegistry).set(eq("2"), argThat(iterableOf(node2.getPublic())));
        ordered.verifyNoMoreInteractions();

        //Verify the task we captured is to remove the control plane's EC key
        taskCaptor.getValue().run();
        verify(keyRegistry).remove(eq(NID_CONTROL_PLANE), eq(controlEc));

        //This test verifies persistence roundtripping for an RSA control plane key
        PublicKey controlKey = KeyUtils.readPublicKey(pem);
        assertNotNull("The control plane key should have been written to disk", controlKey);
        assertEquals("The wrong control plane key was persisted", controlRsa, controlKey);
    }

    @Test
    public void testWithPrincipalCreatesExecutionContextIfNecessary() {
        Principal principal = mock(Principal.class);

        executionContextRule.runWithoutContext(() -> {
            String result = manager.withPrincipal(principal, actual -> {
                assertSame(principal, actual);
                assertSame(principal, manager.getPrincipalOrThrow());

                return "+";
            });
            assertEquals("+", result);
        });

        //The context created inside withPrincipal should have been closed when withPrincipal returned,
        //so the principal should no longer be set
        assertEquals(Optional.empty(), manager.getPrincipal());
    }

    @Test
    public void testWithPrincipalInExistingExecutionContext() {
        Principal principal = mock(Principal.class);

        executionContextRule.runInNewContext(() -> {
            String result = manager.withPrincipal(principal, actual -> {
                assertSame(principal, actual);
                assertSame(principal, manager.getPrincipalOrThrow());

                return "+";
            });
            assertEquals("+", result);

            //The principal should still be set, because the ExecutionContext is still open
            assertSame(principal, manager.getPrincipalOrThrow());
        });

        //Now that the ExecutionContext is closed, the principal should be cleared
        assertEquals(Optional.empty(), manager.getPrincipal());
    }

    private static void assertSigningKey(KeyPair expected, SigningKey actual) {
        assertNotNull(actual);
        assertEquals(expected.getPublic(), actual.getPublic());
        assertEquals(expected.getPrivate(), actual.getPrivate());
        assertEquals(KeyUtils.getFingerprint(expected.getPublic()), actual.getFingerprint());
    }
}