package com.atlassian.bitbucket.mesh.auth;

import com.atlassian.bitbucket.mesh.AbstractGrpcTest;
import com.atlassian.bitbucket.mesh.MockResponseObserver;
import com.atlassian.bitbucket.mesh.config.ApplicationConfig;
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.network.Partition;
import com.atlassian.bitbucket.mesh.rpc.v1.*;
import com.atlassian.bitbucket.mesh.rpc.v1.ManagementServiceGrpc.ManagementServiceBlockingStub;
import com.google.common.collect.Collections2;
import com.google.protobuf.ByteString;
import io.grpc.ClientInterceptor;
import io.grpc.ManagedChannel;
import io.grpc.Status;
import io.grpc.StatusRuntimeException;
import io.grpc.stub.StreamObserver;
import io.jsonwebtoken.SignatureAlgorithm;
import io.jsonwebtoken.security.Keys;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.springframework.test.annotation.DirtiesContext;
import org.springframework.test.context.TestPropertySource;

import javax.annotation.Resource;
import java.security.KeyPair;
import java.security.PublicKey;
import java.util.concurrent.TimeoutException;

import static java.util.concurrent.TimeUnit.SECONDS;
import static org.mockito.Mockito.*;

//Enable authentication and configuration for this test only
@SuppressWarnings("ResultOfMethodCallIgnored")
@TestPropertySource(properties = {
        "authentication.enabled=true",
        "authentication.keyRequestInterval=3",
        "node.config-required=true",
        "grpc.server.port=7778" // use a different port to avoid conflicts with the other non-authenticated wired tests
})
public class GrpcAuthenticationAndAuthorisationTest extends AbstractGrpcTest {

    public static final String PEER_NODE_1 = "1001";
    public static final String PEER_NODE_2 = "1002";
    private static AuthenticationClientInterceptor asControlPlane;
    private static AuthenticationClientInterceptor asPeer1;
    private static AuthenticationClientInterceptor asPeer2;
    private static AuthenticationClientInterceptor asUnknown;
    private static KeyPair controlPlaneKeyPair;

    @Rule
    public final LogListener logListener = new LogListener("com.atlassian.bitbucket.mesh");

    @Resource
    private ApplicationConfig config;
    @Resource
    private PublicKeyRegistry keyRegistry;
    @Resource
    private ManagedChannel managedChannel;
    @Resource
    private MeshManager meshManager;
    @Resource
    private SidebandManager sidebandManager;

    @Before
    public void setUp() {
        maybeConfigure();
    }

    @Test
    public void testNonRestrictedRequestAnonymouslyThrows() {
        //ManagementService/HeartBeat is not restricted by @RequiresRole, but should still be rejected
        //if not signed
        StatusRuntimeException exception = assertThrows(StatusRuntimeException.class, () ->
                blockingStub().heartbeat(RpcHeartbeat.getDefaultInstance()));
        assertEquals(Status.UNAUTHENTICATED.getCode(), exception.getStatus().getCode());
    }

    @Test
    public void testProvisionAnonymouslyWhenAlreadyProvisionedThrows() {
        //ManagementService/Provision requires CONTROL_PLANE when already provisioned, so this should throw
        //when called anonymously
        StatusRuntimeException exception = assertThrows(StatusRuntimeException.class, () ->
                blockingStub().provisionNode(RpcProvisionNodeRequest.newBuilder()
                        .setNodeId(config.getNode().getId())
                        .setPublicKey(toPublicKey(controlPlaneKeyPair.getPublic()))
                        .build()));
        assertEquals(Status.UNAUTHENTICATED.getCode(), exception.getStatus().getCode());
    }

    @Test
    public void testProvisionAsControlPlane() {
        //ManagementService/Provision requires CONTROL_PLANE when already provisioned, so this should
        //not throw when called from the control plane
        blockingStub(asControlPlane).provisionNode(RpcProvisionNodeRequest.newBuilder()
                .setNodeId(config.getNode().getId())
                .setPublicKey(toPublicKey(controlPlaneKeyPair.getPublic()))
                .build());
        logListener.assertInfo("Registered new control plane key");
    }

    @Test
    public void testProvisionAsPeerThrows() {
        //ManagementService/Provision requires CONTROL_PLANE when already provisioned, so this should throw
        //when called from a PEER node
        StatusRuntimeException exception = assertThrows(StatusRuntimeException.class, () ->
                blockingStub(asPeer1).provisionNode(RpcProvisionNodeRequest.newBuilder()
                        .setNodeId(config.getNode().getId())
                        .setPublicKey(toPublicKey(controlPlaneKeyPair.getPublic()))
                        .build()));
        assertEquals(Status.PERMISSION_DENIED.getCode(), exception.getStatus().getCode());
        logListener.assertWarn("Blocked a (re-)provisioning attempt by " + PEER_NODE_1);
    }

    @Test
    @DirtiesContext(methodMode = DirtiesContext.MethodMode.AFTER_METHOD)
    public void testProvisionAsUnknownThrows() {
        try {
            //ManagementService/Provision requires CONTROL_PLANE when already provisioned, so this should throw
            //when called from an unknown node
            StatusRuntimeException exception = assertThrows(StatusRuntimeException.class, () ->
                    blockingStub(asUnknown).provisionNode(RpcProvisionNodeRequest.newBuilder()
                            .setNodeId(config.getNode().getId())
                            .setPublicKey(toPublicKey(controlPlaneKeyPair.getPublic()))
                            .build()));
            assertEquals(Status.UNAUTHENTICATED.getCode(), exception.getStatus().getCode());
        } finally {
            controlPlaneKeyPair = null;
        }
    }

    @Test
    public void testPublicRpcSignedByControlPlane() {
        //ManagementService/GetNodeInfo is not anonymously accessible and should completely disregard the signature
        blockingStub(asControlPlane).getNodeInfo(RpcGetNodeInfoRequest.getDefaultInstance());
    }

    @Test
    public void testPublicRpcSignedByPeer() {
        //ManagementService/GetNodeInfo is not anonymously accessible and should completely disregard the signature
        blockingStub(asControlPlane).getNodeInfo(RpcGetNodeInfoRequest.getDefaultInstance());
    }

    @Test
    public void testPublicRpcSignedByUnknownKey() {
        //ManagementService/GetNodeInfo is not anonymously accessible and should completely disregard the signature
        blockingStub(asControlPlane).getNodeInfo(RpcGetNodeInfoRequest.getDefaultInstance());
    }

    @Test
    public void testUnrestrictedRpcSignedByControlPlane() {
        //ManagementService/HeartBeat is not restricted by @RequiresRole
        blockingStub(asControlPlane).heartbeat(RpcHeartbeat.getDefaultInstance());
    }

    @Test
    public void testUnrestrictedRpcSignedByPeer() {
        //ManagementService/HeartBeat is not restricted by @RequiresRole
        blockingStub(asPeer1).heartbeat(RpcHeartbeat.getDefaultInstance());
    }

    @Test
    @DirtiesContext(methodMode = DirtiesContext.MethodMode.AFTER_METHOD)
    public void testUnrestrictedRpcSignedByUnknownKeyThrows() {
        try {
            //ManagementService/HeartBeat is not restricted by @RequiresRole, but should still be rejected
            //if signed with an unknown key
            StatusRuntimeException exception = assertThrows(StatusRuntimeException.class, () ->
                    blockingStub(asUnknown).heartbeat(RpcHeartbeat.getDefaultInstance()));
            assertEquals(Status.UNAUTHENTICATED.getCode(), exception.getStatus().getCode());

            logListener.assertWarn("Rejecting token; signature verification failed");
        } finally {
            controlPlaneKeyPair = null;
        }
    }

    @Test
    @DirtiesContext(methodMode = DirtiesContext.MethodMode.AFTER_METHOD)
    public void testUnrestrictedRpcSignedByKnownPeerSendsSidebandSignalWhenPublicKeyIsMissing()
            throws InterruptedException, TimeoutException {
        ManagementServiceGrpc.ManagementServiceStub managedStub = ManagementServiceGrpc.newStub(managedChannel)
                .withInterceptors(asControlPlane);
        MockResponseObserver<RpcSidebandFollowerFragment> responseObserver = new MockResponseObserver<>();
        StreamObserver<RpcSidebandLeaderFragment> sideband = managedStub.sideband(responseObserver);
        try {
            blockingStub(asControlPlane).setConfiguration(RpcSetConfigurationRequest.newBuilder()
                    .setTopology(RpcTopology.newBuilder()
                            .addNodes(RpcNode.newBuilder()
                                    .setId(PEER_NODE_2)
                                    .setRpcUrl("http://localhost:7886")
                                    .setName("peer-node-2")
                                    .build())
                            .build()).build());
            StatusRuntimeException exception = assertThrows(StatusRuntimeException.class, () ->
                    blockingStub(asPeer2).heartbeat(RpcHeartbeat.getDefaultInstance()));
            assertEquals(Status.UNAUTHENTICATED.getCode(), exception.getStatus().getCode());
            // Calling this method multiple times to test rate limiting which should result in only one fragment
            // received on sideband
            assertThrows(StatusRuntimeException.class, () ->
                    blockingStub(asPeer2).heartbeat(RpcHeartbeat.getDefaultInstance()));
            // authentication.keyRequestInterval is configured as 3 seconds, hence sleeping for a little more than that.
            Thread.sleep(3500);
            RpcSidebandFollowerFragment followerFragment = responseObserver.awaitFragment(0, 5, SECONDS);
            assertEquals(PEER_NODE_2, followerFragment.getNodeKeyRequest().getNodeId());
            assertNotNull(followerFragment.getNodeKeyRequest().getFingerprint());
            assertEquals(1, responseObserver.getFragments().size());
            assertThrows(StatusRuntimeException.class, () ->
                    blockingStub(asPeer2).heartbeat(RpcHeartbeat.getDefaultInstance()));
            responseObserver.awaitFragment(1, 5, SECONDS);
            assertEquals(2, responseObserver.getFragments().size());
        } finally {
            sideband.onCompleted();
            controlPlaneKeyPair = null;
        }
    }

    private static RpcPublicKey toPublicKey(PublicKey publicKey) {
        return RpcPublicKey.newBuilder()
                .setDer(ByteString.copyFrom(publicKey.getEncoded()))
                .build();
    }

    private AuthenticationClientInterceptor authInterceptor(String nodeId, KeyPair signingKey) {
        AuthenticationManager authenticationManager = mock(AuthenticationManager.class);
        when(authenticationManager.getSigningKey()).thenReturn(new SimpleSigningKey(signingKey));
        when(authenticationManager.isProvisioned()).thenAnswer(invocation -> controlPlaneKeyPair != null);
        ApplicationConfig configSpy = spy(config);
        NodeConfig nodeConfigSpy = spy(config.getNode());
        when(configSpy.getNode()).thenReturn(nodeConfigSpy);
        when(nodeConfigSpy.getId()).thenReturn(nodeId);
        return new AuthenticationClientInterceptor(new JwtAuthenticator(authenticationManager, configSpy, keyRegistry,
                meshManager, sidebandManager));
    }

    private ManagementServiceBlockingStub blockingStub(ClientInterceptor... interceptors) {
        return ManagementServiceGrpc.newBlockingStub(managedChannel).withInterceptors(interceptors);
    }

    private RpcTopology createTopology() {
        RpcTopology.Builder builder = RpcTopology.newBuilder();
        builder.addAllNodes(Collections2.transform(meshManager.getNodes(), MeshNode::toNode));
        builder.addAllPartitions(Collections2.transform(meshManager.getPartitions(), Partition::toPartition));
        return builder.build();
    }

    private void maybeConfigure() {
        if (controlPlaneKeyPair == null) {
            controlPlaneKeyPair = Keys.keyPairFor(SignatureAlgorithm.RS256);
            KeyPair peerKeyPair1 = Keys.keyPairFor(SignatureAlgorithm.RS256);
            KeyPair peerKeyPair2 = Keys.keyPairFor(SignatureAlgorithm.RS256);
            KeyPair unknownKeyPair = Keys.keyPairFor(SignatureAlgorithm.RS256);

            //Provision the node with the control plane key
            blockingStub().provisionNode(RpcProvisionNodeRequest.newBuilder()
                    .setPublicKey(toPublicKey(controlPlaneKeyPair.getPublic()))
                    .setNodeId(config.getNode().getId())
                    .build());

            asControlPlane = authInterceptor("control-plane", controlPlaneKeyPair);
            asPeer1 = authInterceptor(PEER_NODE_1, peerKeyPair1);
            asPeer2 = authInterceptor(PEER_NODE_2, peerKeyPair2);
            asUnknown = authInterceptor("unknown-node", unknownKeyPair);

            //Configure the node with the control plane key and the peer public key to ensure requests signed with
            //the controlPlaneKeyPair and the peerKeyPair1 are recognized
            blockingStub(asControlPlane).setConfiguration(RpcSetConfigurationRequest.newBuilder()
                    .setNodeConfiguration(RpcNodeConfiguration.newBuilder()
                            .putProperties("node.id", config.getNode().getId())
                            .putProperties("node.name", "auth-test-node"))
                    .setPublicKeyConfiguration(RpcPublicKeyConfiguration.newBuilder()
                            .setControlPlaneKeys(RpcPublicKeyList.newBuilder()
                                    .addPublicKeys(toPublicKey(controlPlaneKeyPair.getPublic())))
                            .putNodePublicKeys(PEER_NODE_1, RpcPublicKeyList.newBuilder()
                                    .addPublicKeys(toPublicKey(peerKeyPair1.getPublic()))
                                    .build()))
                    .setTopology(createTopology())
                    .build());
        }
    }
}
