package com.atlassian.bitbucket.internal.ssh.server;

import com.atlassian.bitbucket.auth.AuthenticationException;
import com.atlassian.bitbucket.auth.SshAuthenticationFailureContext;
import com.atlassian.bitbucket.auth.SshAuthenticationSuccessContext;
import com.atlassian.bitbucket.event.auth.AuthenticationFailureEvent;
import com.atlassian.bitbucket.i18n.SimpleI18nService;
import com.atlassian.bitbucket.internal.ssh.auth.PluginSshAuthenticationHandler;
import com.atlassian.bitbucket.scm.AuthenticationState;
import com.atlassian.bitbucket.server.ApplicationPropertiesService;
import com.atlassian.bitbucket.ssh.util.KeyUtils;
import com.atlassian.bitbucket.user.ApplicationUser;
import com.atlassian.bitbucket.user.DummySecurityService;
import com.atlassian.bitbucket.user.SecurityService;
import com.atlassian.event.api.EventPublisher;
import com.atlassian.johnson.Johnson;
import com.atlassian.plugin.servlet.ServletContextFactory;
import org.apache.mina.util.NamePreservingRunnable;
import org.apache.sshd.common.io.IoSession;
import org.apache.sshd.server.session.ServerSession;
import org.hamcrest.BaseMatcher;
import org.hamcrest.Description;
import org.junit.AfterClass;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.InjectMocks;
import org.mockito.Mock;
import org.mockito.Spy;
import org.mockito.invocation.InvocationOnMock;
import org.mockito.runners.MockitoJUnitRunner;
import org.mockito.stubbing.Answer;

import javax.servlet.ServletContext;
import java.net.SocketAddress;
import java.security.PublicKey;
import java.util.concurrent.*;

import static com.atlassian.bitbucket.internal.ssh.server.SessionAttributes.ATTRIBUTE_AUTH;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertTrue;
import static org.mockito.Mockito.*;

@RunWith(MockitoJUnitRunner.class)
public class DefaultPublicKeyAuthenticatorTest {

    private static final String PUBLIC_KEY_TEXT = "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQDBUqlkHktdkVsXbZ+OQnjgAWziTWB1S/enTZ9sisuzFfXVNlxujLjjKHjFSZRSzvFrgG3/qFhchGNqrtgskhQHnJmtQwMOJQIWKyb2N+6KWjZ/V5GEf+ovgHxwxfcq3AIIYrsr/IsIWG3cWirUaBUPkwRwoNwDmqgcV7JRxDlRgmsx1oyg8iMXth8qbYxMgBu/1NjVGM//do9e1sWAh7up7AJxfdzuCQ0xJeHSAU0jA1U9kdpAfi1Ek1SIWKxUkHEew90i6MdwW0GvWLxdJI8R6JobuylfErtrUMpPjXLGo8jCzhoLPVgQ/hat5N5JDLD1j205N3O/NqKwwH1VIBPD tpettersen@hoprocker.local";
    private static final String USERNAME = "username";

    @Mock
    private EventPublisher eventPublisher;
    @Mock
    private ServletContextFactory servletContextFactory;
    @Mock
    private PluginSshAuthenticationHandler authenticationHandler;
    @Mock
    private ExecutorService executorService;
    @Mock
    private ApplicationPropertiesService propertiesService;
    @Mock
    private ServerSession session;
    @Mock
    private ServletContext servletContext;
    @Spy
    private SimpleI18nService i18nService;
    @Mock
    private IoSession ioSession;
    @Mock
    private SocketAddress remoteAddress;
    @Mock
    private ApplicationUser user;
    @Spy // so mockito can inject it
    private SecurityService securityService = new DummySecurityService();

    @InjectMocks
    private DefaultPublicKeyAuthenticator authenticator;

    private PublicKey publicKey;

    @AfterClass
    public static void tearDownJohnson() {
        Johnson.terminate();
    }

    @Before
    public void setup() throws Exception {
        publicKey = KeyUtils.getPublicKey(PUBLIC_KEY_TEXT);

        when(propertiesService.getPluginProperty("plugin.ssh.auth.timeout", 30)).thenReturn(30);
        when(servletContextFactory.getServletContext()).thenReturn(servletContext);
        when(session.getIoSession()).thenReturn(ioSession);
        when(ioSession.getRemoteAddress()).thenReturn(remoteAddress);
        when(user.getName()).thenReturn("user");

        //onAuthenticationFailure handling
        when(authenticationHandler.hasFailureHandlers()).thenReturn(true);
        doAnswer(new Answer<Void>() {
            @Override
            public Void answer(InvocationOnMock invocationOnMock) throws Throwable {
                ((Runnable) invocationOnMock.getArguments()[0]).run();
                return null;
            }
        }).when(executorService).execute(any(Runnable.class));

        //Future authenticate handling
        final Future<?> handleFuture = mock(Future.class);
        when(executorService.submit(any(Runnable.class))).thenAnswer(new Answer<Future<?>>() {
            @Override
            public Future<?> answer(InvocationOnMock invocationOnMock) throws Throwable {
                ((Runnable) invocationOnMock.getArguments()[0]).run();
                return handleFuture;
            }
        });
        when(handleFuture.get(30, TimeUnit.SECONDS)).thenReturn(null);
    }

    @Test
    public void testAuthenticateFailInterrupted() throws Exception {
        Future<ApplicationUser> userFuture = mock(Future.class);
        doReturn(userFuture).when(executorService).submit(any(Callable.class));
        when(userFuture.get(30, TimeUnit.SECONDS)).thenThrow(new InterruptedException());

        assertFalse(authenticator.authenticate(USERNAME, publicKey, session));
        verify(eventPublisher, never()).publish(any(AuthenticationFailureEvent.class));
        verify(authenticationHandler, never()).onAuthenticationFailure(any(SshAuthenticationFailureContext.class));
        verify(authenticationHandler, never()).onAuthenticationSuccess(any(SshAuthenticationSuccessContext.class));
        verify(userFuture).cancel(true);
    }

    @Test
    public void testAuthenticateFailTimedOut() throws Exception {
        Future<ApplicationUser> userFuture = mock(Future.class);
        doReturn(userFuture).when(executorService).submit(any(Callable.class));
        when(userFuture.get(30, TimeUnit.SECONDS)).thenThrow(new TimeoutException());

        assertFalse(authenticator.authenticate(USERNAME, publicKey, session));
        verify(eventPublisher, never()).publish(any(AuthenticationFailureEvent.class));
        verify(userFuture).cancel(true);
    }

    @Test(expected = NullPointerException.class)
    public void testAuthenticateFailExecutionException() throws Exception {
        Future<ApplicationUser> userFuture = mock(Future.class);
        doReturn(userFuture).when(executorService).submit(any(Callable.class));
        when(userFuture.get(30, TimeUnit.SECONDS)).thenThrow(new ExecutionException(new NullPointerException()));

        try {
            authenticator.authenticate(USERNAME, publicKey, session);
        } finally {
            verify(eventPublisher, never()).publish(any(AuthenticationFailureEvent.class));
            verify(authenticationHandler, never()).onAuthenticationFailure(any(SshAuthenticationFailureContext.class));
        }
    }

    @Test
    public void testAuthenticateFailureException() throws Exception {
        Future<ApplicationUser> userFuture = mock(Future.class);
        doReturn(userFuture).when(executorService).submit(any(Callable.class));

        ExecutionException executionException = new ExecutionException(mock(AuthenticationException.class));
        when(userFuture.get(30, TimeUnit.SECONDS)).thenThrow(executionException);

        assertFalse(authenticator.authenticate(USERNAME, publicKey, session));
        verify(eventPublisher).publish(any(AuthenticationFailureEvent.class));
        verify(executorService).execute(any(NamePreservingRunnable.class));
        verify(authenticationHandler).onAuthenticationFailure(argThat(matchesFailedAuthenticationState(AuthenticationState.NOT_AUTHENTICATED)));
    }

    @Test
    public void testAuthenticate() throws Throwable {
        Future<SshAuthentication> authResultFuture = mock(Future.class);
        doReturn(authResultFuture).when(executorService).submit(any(Callable.class));

        SshAuthentication futureResult = new SshAuthentication(user, new SshCredentials(USERNAME, publicKey), new SshKeyDetails(null, null));
        when(authResultFuture.get(30, TimeUnit.SECONDS)).thenReturn(futureResult);
        when(user.getId()).thenReturn(1);
        when(user.getDisplayName()).thenReturn(USERNAME);

        assertTrue(authenticator.authenticate(USERNAME, publicKey, session));
        verify(session).setAttribute(ATTRIBUTE_AUTH, futureResult);
    }

    private BaseMatcher<SshAuthenticationFailureContext> matchesFailedAuthenticationState(final AuthenticationState state) {
        return new BaseMatcher<SshAuthenticationFailureContext>() {
            @Override
            public boolean matches(Object item) {
                return ((SshAuthenticationFailureContext) item).getAuthenticationState() == state;
            }

            @Override
            public void describeTo(Description description) {
                description.appendText("Failure context with authentication state " + state);
            }
        };
    }
}