package com.atlassian.bitbucket.mesh.auth;

import com.atlassian.bitbucket.mesh.rpc.CapabilityServiceGrpc;
import com.atlassian.bitbucket.mesh.rpc.v1.ManagementServiceGrpc;
import io.grpc.Metadata;
import io.grpc.ServerCall;
import io.grpc.ServerCallHandler;
import io.grpc.Status;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.ArgumentCaptor;
import org.mockito.Captor;
import org.mockito.InjectMocks;
import org.mockito.Mock;
import org.mockito.junit.MockitoJUnitRunner;

import static com.atlassian.bitbucket.mesh.auth.AuthenticationServerInterceptor.METHOD_PROVISION_NODE;
import static com.atlassian.bitbucket.mesh.rpc.MetadataKeys.KEY_TOKEN;
import static org.junit.Assert.*;
import static org.mockito.ArgumentMatchers.*;
import static org.mockito.Mockito.*;

@RunWith(MockitoJUnitRunner.class)
public class AuthenticationServerInterceptorTest {

    static final String METHOD_HEARTBEAT = ManagementServiceGrpc.getHeartbeatMethod().getFullMethodName();

    @Mock
    private AuthenticationManager authenticationManager;
    @Mock
    private Authenticator authenticator;
    @Mock
    private ServerCall<Object, Object> call;
    @Mock
    private ServerCallHandler<Object, Object> handler;
    @InjectMocks
    private AuthenticationServerInterceptor interceptor;
    @Mock
    private ServerCall.Listener<Object> listener;
    @Captor
    private ArgumentCaptor<Status> statusCaptor;

    @Before
    public void setup() {
        doReturn(ManagementServiceGrpc.getHeartbeatMethod()).when(call).getMethodDescriptor();
        when(handler.startCall(same(call), any())).thenReturn(listener);
    }

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

        when(authenticationManager.withPrincipal(same(principal), notNull())).thenAnswer(invocation -> {
            Principal actual = invocation.getArgument(0);

            return invocation.<PrincipalCallback<?>>getArgument(1).callWithPrincipal(actual);
        });
        when(authenticator.authenticate(notNull(), notNull())).thenReturn(principal);

        Metadata metadata = new Metadata();
        metadata.put(KEY_TOKEN, "token");

        assertSame(listener, interceptor.interceptCall(call, metadata, handler));

        verify(authenticationManager).withPrincipal(same(principal), notNull());
        verify(authenticator).authenticate(eq(METHOD_HEARTBEAT), eq("token"));
    }

    @Test
    public void testInterceptCallAllowsAnonymousForGetCapabilities() {
        doReturn(CapabilityServiceGrpc.getGetCapabilitiesMethod()).when(call).getMethodDescriptor();

        assertSame(listener, interceptor.interceptCall(call, new Metadata(), handler));

        verifyNoInteractions(authenticationManager, authenticator);
    }

    @Test
    public void testInterceptCallAllowsAnonymousForProvisionNodeIfNotProvisioned() {
        doReturn(ManagementServiceGrpc.getProvisionNodeMethod()).when(call).getMethodDescriptor();

        assertSame(listener, interceptor.interceptCall(call, new Metadata(), handler));

        verify(authenticationManager).isProvisioned();
        verifyNoInteractions(authenticator);
    }

    @Test
    public void testInterceptCallRejectsInvalidTokenForHeartbeat() {
        Metadata metadata = new Metadata();
        metadata.put(KEY_TOKEN, "token");

        interceptor.interceptCall(call, metadata, handler);

        verify(authenticator).authenticate(eq(METHOD_HEARTBEAT), eq("token"));
        verify(call).close(statusCaptor.capture(), any());

        Status status = statusCaptor.getValue();
        assertSame(Status.Code.UNAUTHENTICATED, status.getCode());
        assertEquals("The [" + KEY_TOKEN.name() + "] is invalid", status.getDescription());
    }

    @Test
    public void testInterceptCallRejectsInvalidTokenForProvisionNodeIfProvisioned() {
        Metadata metadata = new Metadata();
        metadata.put(KEY_TOKEN, "token");

        doReturn(ManagementServiceGrpc.getProvisionNodeMethod()).when(call).getMethodDescriptor();
        when(authenticationManager.isProvisioned()).thenReturn(true);

        interceptor.interceptCall(call, metadata, handler);

        verify(authenticationManager).isProvisioned();
        verify(authenticator).authenticate(eq(METHOD_PROVISION_NODE), eq("token"));
        verify(call).close(statusCaptor.capture(), any());

        Status status = statusCaptor.getValue();
        assertSame(Status.Code.UNAUTHENTICATED, status.getCode());
        assertEquals("The [" + KEY_TOKEN.name() + "] is invalid", status.getDescription());
    }

    @Test
    public void testInterceptCallRejectsNullTokenWithoutAuthenticatingForHeartbeat() {
        interceptor.interceptCall(call, new Metadata(), handler);

        verify(call).close(statusCaptor.capture(), any());
        verifyNoInteractions(authenticator);

        Status status = statusCaptor.getValue();
        assertSame(Status.Code.UNAUTHENTICATED, status.getCode());
        assertEquals("No [" + KEY_TOKEN.name() + "] was sent", status.getDescription());
    }

    @Test
    public void testInterceptCallRejectsNullTokenWithoutAuthenticatingForProvisionNodeIfProvisioned() {
        doReturn(ManagementServiceGrpc.getProvisionNodeMethod()).when(call).getMethodDescriptor();
        when(authenticationManager.isProvisioned()).thenReturn(true);

        interceptor.interceptCall(call, new Metadata(), handler);

        verify(authenticationManager).isProvisioned();
        verify(call).close(statusCaptor.capture(), any());
        verifyNoInteractions(authenticator);

        Status status = statusCaptor.getValue();
        assertSame(Status.Code.UNAUTHENTICATED, status.getCode());
        assertEquals("No [" + KEY_TOKEN.name() + "] was sent", status.getDescription());
    }
}
