package com.atlassian.stash.internal.auth;

import com.atlassian.security.random.SecureTokenGenerator;
import com.atlassian.bitbucket.annotation.FixedCurrentTime;
import com.atlassian.stash.internal.ApplicationConstants;
import com.atlassian.stash.internal.user.InternalApplicationUser;
import com.atlassian.stash.internal.web.HttpRequestInfoHelper;
import com.atlassian.bitbucket.mockito.MockitoUtils;
import com.atlassian.bitbucket.rule.FixedCurrentTimeRule;
import org.apache.commons.lang3.StringUtils;
import org.hamcrest.BaseMatcher;
import org.hamcrest.Description;
import org.hamcrest.Matcher;
import org.joda.time.Instant;
import org.joda.time.Seconds;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.ArgumentCaptor;
import org.mockito.Captor;
import org.mockito.InOrder;
import org.mockito.Mock;
import org.mockito.runners.MockitoJUnitRunner;
import org.springframework.dao.DataAccessException;
import org.springframework.dao.DataIntegrityViolationException;
import org.springframework.transaction.PlatformTransactionManager;
import org.springframework.transaction.TransactionStatus;

import javax.servlet.http.Cookie;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;

import java.util.concurrent.TimeUnit;

import static org.junit.Assert.*;
import static org.mockito.Matchers.argThat;
import static org.mockito.Mockito.*;

@RunWith(MockitoJUnitRunner.class)
public class DefaultRememberMeServiceTest {

    private static final String COOKIE_NAME = "__cookie_name";
    private static final String REMOTE_ADDRESS = "192.168.171.13";
    private static final String SERIES = "remember-me-series";
    private static final String TOKEN = "remember-me-value";
    private static final int USER_ID = 37;
    private static final int EXPIRY_DAYS = 15;

    @Rule
    public FixedCurrentTimeRule fixedCurrentTimeRule = new FixedCurrentTimeRule();

    @Captor
    private ArgumentCaptor<Cookie> cookieCaptor;
    @Mock
    private RememberMeTokenDao dao;
    @Mock
    private HttpServletRequest request;
    @Mock
    private HttpRequestInfoHelper requestInfoHelper;
    @Mock
    private HttpServletResponse response;
    @Captor
    private ArgumentCaptor<InternalRememberMeToken> tokenCaptor;
    @Mock
    private SecureTokenGenerator tokenGenerator;
    @Mock
    private PlatformTransactionManager transactionManager;
    @Mock
    private InternalApplicationUser user;
    @Mock
    private InternalRememberMeToken token;

    private DefaultRememberMeService service;

    @Before
    public void setUp() {
        service = createService("optional");
        Cookie cookie = new Cookie(COOKIE_NAME, service.encodeCookie(SERIES, TOKEN));
        cookie.setMaxAge(Integer.MAX_VALUE);
        when(request.getCookies()).thenReturn(new Cookie[] {cookie});
        when(requestInfoHelper.getRemoteAddress(request)).thenReturn(REMOTE_ADDRESS);
        when(request.isSecure()).thenReturn(true);

        when(user.getId()).thenReturn(USER_ID);

        when(dao.create(any(InternalRememberMeToken.class))).thenAnswer(MockitoUtils.returnFirst());
    }

    @Test
    public void testAuthenticate() {
        InternalRememberMeToken validToken = new InternalRememberMeToken.Builder()
                .expiresAfter(10, TimeUnit.HOURS)
                .series(SERIES)
                .token(TOKEN)
                .user(user)
                .build();

        when(dao.findBySeriesAndToken(SERIES, TOKEN)).thenReturn(validToken);
        when(dao.claim(argThat(isSeriesAndToken(SERIES, TOKEN)))).thenReturn(true);
        String nextToken = "time-for-a-new-token";
        when(tokenGenerator.generateToken()).thenReturn(nextToken);

        assertEquals(user, service.authenticate(request, response));

        verify(dao).findBySeriesAndToken(SERIES, TOKEN);
        verify(dao).claim(argThat(isSeriesAndToken(SERIES, TOKEN)));
        verify(dao).create(tokenCaptor.capture());
        verifyNoMoreInteractions(dao);

        // validate the token
        InternalRememberMeToken token = tokenCaptor.getValue();
        assertEquals(SERIES, token.getSeries());
        assertEquals(nextToken, token.getToken());
        assertFalse(token.isClaimed());
        assertFalse(token.isExpired());
        assertNotNull(token.getExpiryDate());
        long expiry = token.getExpiryDate().getTime() - System.currentTimeMillis();
        long expectedExpiry = TimeUnit.DAYS.toMillis(EXPIRY_DAYS);
        assertTrue(expiry + 5000L > expectedExpiry && expiry <= expectedExpiry); // 5s margin for test duration
        assertNull(token.getRemoteAddress());
        assertEquals(user, token.getUser());

        // validate the cookie
        assertCookie(SERIES, nextToken, (int) expectedExpiry / 1000);
    }

    @Test
    public void testAuthenticateClaimedToken() {
        // token is found, but has already be claimed. It's still in its grace period though
        InternalRememberMeToken claimedToken = new InternalRememberMeToken.Builder()
                .claimed()
                .expiresAfter(1, TimeUnit.MINUTES)
                .remoteAddress(REMOTE_ADDRESS)
                .series(SERIES)
                .token(TOKEN)
                .user(user)
                .build();
        when(dao.findBySeriesAndToken(SERIES, TOKEN)).thenReturn(claimedToken);

        // authentication should succeed
        assertEquals(user, service.authenticate(request, response));

        // verify that no tokens are updated
        verify(dao).findBySeriesAndToken(SERIES, TOKEN);
        verifyNoMoreInteractions(dao);

        // verify that the cookie is not updated on the response
        verifyZeroInteractions(response);
    }

    @Test
    public void testAuthenticateClaimedTokenFromOtherRemoteAddress() {
        // token is found, but has already be claimed. It's still in its grace period though
        InternalRememberMeToken claimedToken = new InternalRememberMeToken.Builder()
                .claimed()
                .expiresAfter(1, TimeUnit.MINUTES)
                .remoteAddress("127.0.0.1") // different remote address!
                .series(SERIES)
                .token(TOKEN)
                .user(user)
                .build();

        when(dao.findBySeriesAndToken(SERIES, TOKEN)).thenReturn(claimedToken);
        when(dao.deleteForSeries(SERIES)).thenReturn(2);

        // authentication should not succeed
        assertNull(service.authenticate(request, response));

        // verify that the series has been deleted
        verify(dao).deleteForSeries(SERIES);

        // verify that the cookie has been reset
        verifyCookieCleared();
    }

    @Test
    public void testAuthenticateConcurrentClaim() {
        InternalRememberMeToken token = spy(new InternalRememberMeToken.Builder()
                .expiresAfter(10, TimeUnit.HOURS)
                .series(SERIES)
                .token(TOKEN)
                .user(user)
                .build());

        when(dao.findBySeriesAndToken(SERIES, TOKEN)).thenReturn(token);
        when(dao.claim(argThat(isSeriesAndToken(SERIES, TOKEN)))).thenReturn(false);

        // isClaimed is called once before and once after the refresh
        when(token.isClaimed()).thenReturn(false, true);
        when(token.getRemoteAddress()).thenReturn(REMOTE_ADDRESS);

        assertEquals(user, service.authenticate(request, response));

        InOrder inOrder = inOrder(dao, token);

        inOrder.verify(dao).findBySeriesAndToken(SERIES, TOKEN);
        inOrder.verify(token).isClaimed();
        inOrder.verify(dao).claim(argThat(isSeriesAndToken(SERIES, TOKEN)));
        inOrder.verify(dao).refresh(token);
        inOrder.verify(token).isClaimed();
        verifyNoMoreInteractions(dao);
    }

    @Test
    public void testAuthenticateConcurrentClaimInDatabase() {
        InternalRememberMeToken token = new InternalRememberMeToken.Builder()
                .expiresAfter(10, TimeUnit.HOURS)
                .series(SERIES)
                .token(TOKEN)
                .user(user)
                .build();
        String nextValue = "next-token";
        when(tokenGenerator.generateToken()).thenReturn(nextValue);
        when(dao.findBySeriesAndToken(SERIES, TOKEN)).thenReturn(token);
        when(dao.claim(argThat(isSeriesAndToken(SERIES, TOKEN)))).thenReturn(true);

        // mock exceptions being thrown 4 times - claim should be retried
        DataAccessException ex = new DataIntegrityViolationException("testing");
        doThrow(ex).doThrow(ex).doThrow(ex).doThrow(ex).doNothing().when(transactionManager)
                .commit(any(TransactionStatus.class));

        assertEquals(user, service.authenticate(request, response));
        verify(response).addCookie(cookieCaptor.capture());
        String tokenValue = service.decodeCookieValue(cookieCaptor.getValue().getValue())[1];
        assertEquals(nextValue, tokenValue);
        verify(dao, never()).deleteForSeries(SERIES);
        verify(dao, never()).deleteForUser(USER_ID);
    }

    @Test(timeout = 5000L)
    public void testAuthenticateConcurrentClaimInDatabaseFails() {
        InternalRememberMeToken token = new InternalRememberMeToken.Builder()
                .expiresAfter(10, TimeUnit.HOURS)
                .series(SERIES)
                .token(TOKEN)
                .user(user)
                .build();
        String nextValue = "next-token";
        when(tokenGenerator.generateToken()).thenReturn(nextValue);
        when(dao.findBySeriesAndToken(SERIES, TOKEN)).thenReturn(token);
        when(dao.claim(argThat(isSeriesAndToken(SERIES, TOKEN)))).thenReturn(true);

        // mock exceptions being thrown all the time - claim should be retried up to 5 times and authentication should fail
        DataAccessException ex = new DataIntegrityViolationException("testing");
        doThrow(ex).when(transactionManager).commit(any(TransactionStatus.class));

        assertNull(service.authenticate(request, response));
        verifyCookieCleared();
        verify(dao, times(DefaultRememberMeService.MAX_ATTEMPTS)).claim(argThat(isSeriesAndToken(SERIES, TOKEN)));
        verify(dao, never()).deleteForSeries(SERIES);
        verify(dao, never()).deleteForUser(USER_ID);
    }

    @Test
    public void testAuthenticateExpiredClaimedToken() {
        // token could not be claimed
        when(dao.claim(argThat(isSeriesAndToken(SERIES, TOKEN)))).thenReturn(false);
        // token is found, but has already be claimed and expired
        InternalRememberMeToken claimedToken = new InternalRememberMeToken.Builder()
                .claimed()
                .expiresAfter(-10, TimeUnit.SECONDS)
                .remoteAddress(REMOTE_ADDRESS)
                .series(SERIES)
                .token(TOKEN)
                .user(user)
                .build();
        when(dao.findBySeriesAndToken(SERIES, TOKEN)).thenReturn(claimedToken);
        when(dao.deleteForSeries(SERIES)).thenReturn(5);

        assertNull(service.authenticate(request, response));
        // the whole series should be invalidated when an authentication attempts is made with an invalid token
        verify(dao).deleteForSeries(SERIES);
        verifyCookieCleared();
    }



    @Test
    public void testAuthenticateExpiredUnclaimedToken() {
        // token could not be claimed
        when(dao.claim(argThat(isSeriesAndToken(SERIES, TOKEN)))).thenReturn(false);
        // token is found, but has already expired
        InternalRememberMeToken token = new InternalRememberMeToken.Builder()
                .expiresAfter(-10, TimeUnit.SECONDS)
                .series(SERIES)
                .token(TOKEN)
                .user(user)
                .build();
        when(dao.findBySeriesAndToken(SERIES, TOKEN)).thenReturn(token);

        assertNull(service.authenticate(request, response));
        verifyCookieCleared();
    }

    @Test
    public void testAuthenticateInvalidCookie() {
        Cookie cookie = new Cookie(COOKIE_NAME, "not-encoded-properly");
        cookie.setMaxAge(Integer.MAX_VALUE);
        when(request.getCookies()).thenReturn(new Cookie[] {cookie});

        assertNull(service.authenticate(request, response));
        verifyCookieCleared();
    }

    @Test
    public void testAuthenticateNoCookie() {
        when(request.getCookies()).thenReturn(new Cookie[] {});

        assertNull(service.authenticate(request, response));
        verifyZeroInteractions(response);
    }

    @Test
    public void testAuthenticateOverlyLongRemoteAddress() {
        // InternalRememberMeToken.remoteAddress can be at most 255 chars, validate that the limit is enforced
        String remoteAddress = "199.168.1.1,192.168.1.2,192.168.2.114,192.168.8.123,192.168.8.123,192.168.8.123," +
                "192.168.8.123,192.168.8.123,192.168.8.123,192.168.8.123,192.168.8.123,192.168.8.123,192.168.8.123," +
                "192.168.8.123,192.168.8.123,192.168.8.123,192.168.8.123,192.168.8.123,192.255.8.123,192.168.8.123," +
                "192.168.8.123,192.168.8.200";
        when(requestInfoHelper.getRemoteAddress(request)).thenReturn(remoteAddress);

        InternalRememberMeToken validToken = new InternalRememberMeToken.Builder()
                .expiresAfter(10, TimeUnit.HOURS)
                .series(SERIES)
                .token(TOKEN)
                .user(user)
                .build();

        when(dao.findBySeriesAndToken(SERIES, TOKEN)).thenReturn(validToken);
        when(dao.claim(argThat(isSeriesAndToken(SERIES, TOKEN)))).thenReturn(true);
        when(tokenGenerator.generateToken()).thenReturn("next-token");

        // auth success
        assertEquals(user, service.authenticate(request, response));

        // validate that the remote address is properly trimmed before stored in the database
        verify(dao).claim(tokenCaptor.capture());
        String actualAddress = tokenCaptor.getValue().getRemoteAddress();
        assertEquals(255, actualAddress.length());
        assertTrue(remoteAddress.startsWith(actualAddress));
    }

    @Test
    public void testAuthenticateRememberMeNever() {
        service = createService("never");

        assertNull(service.authenticate(request, response));
        verifyZeroInteractions(response, dao);
    }

    @Test
    public void testAuthenticateUnknownToken() {
        // token could not be claimed
        when(dao.claim(argThat(isSeriesAndToken(SERIES, TOKEN)))).thenReturn(false);
        // token is not found
        when(dao.findBySeriesAndToken(SERIES, TOKEN)).thenReturn(null);

        assertNull(service.authenticate(request, response));
        verifyCookieCleared();
    }

    @Test
    public void testCancelCookie() {
        service.cancelCookie(request, response);

        verifyCookieCleared();
    }

    @Test
    public void testCreateCookie() {
        when(request.getCookies()).thenReturn(new Cookie[] {});
        when(tokenGenerator.generateToken()).thenReturn("new-series", "new-value", "moar-tokens");

        service.createCookie(user, request, response);

        verify(dao).create(argThat(isSeriesAndToken("new-series", "new-value")));
        assertCookie("new-series", "new-value", (int) TimeUnit.DAYS.toSeconds(EXPIRY_DAYS));
    }

    @Test(expected = IllegalStateException.class)
    public void testCreateCookieWhenCookieProvided() {
        service.createCookie(user, request, response);
    }

    @Test
    public void testDeleteExpiredTokens() {
        service.deleteExpiredTokens(3, TimeUnit.HOURS);

        ArgumentCaptor<Long> expiryCaptor = ArgumentCaptor.forClass(Long.class);
        ArgumentCaptor<TimeUnit> unitCaptor = ArgumentCaptor.forClass(TimeUnit.class);

        verify(dao).deleteExpiredTokens(expiryCaptor.capture(), unitCaptor.capture());
        assertEquals(TimeUnit.HOURS.toSeconds(3), unitCaptor.getValue().toSeconds(expiryCaptor.getValue()));
    }

    @Test(expected = IllegalArgumentException.class)
    public void testDeleteExpiredTokensIllegalExpiry() {
        service.deleteExpiredTokens(-3, TimeUnit.DAYS);
    }

    @Test
    public void testDeleteForUser() {
        service.deleteForUser(user);

        verify(dao).deleteForUser(USER_ID);
    }

    @Test
    public void testEncodeDecodeCookie() {
        for (String[] tokens : new String[][] {
                new String[] {"series", "value"},
                new String[] {"series1", "value"},
                new String[] {"series12", "value"},
                new String[] {"series123", "value"},
                new String[] {"series1234", "value"},
                new String[] {"series12345", "value"},
                new String[] {StringUtils.repeat('s', 512), StringUtils.repeat('v', 512)}
        }) {
            String encoded = service.encodeCookie(tokens);
            assertNotNull("Cookie value for " + StringUtils.join(tokens[0], ':') + " was null", encoded);
            String[] decoded = service.decodeCookieValue(encoded);
            assertArrayEquals(tokens, decoded);
        }
    }

    @Test
    public void testIsCookieRequested() {
        when(request.getRequestURI()).thenReturn("/stash/" + ApplicationConstants.URL_SECURITY_CHECK);
        when(request.getParameter(DefaultRememberMeService.PARAMETER_REMEMBER_ME)).thenReturn("on");
        when(request.getCookies()).thenReturn(new Cookie[]{});

        assertTrue(service.isCookieRequested(request));
    }

    @Test
    public void testIsCookieRequestedCookieProvided() {
        when(request.getRequestURI()).thenReturn("/stash/" + ApplicationConstants.URL_SECURITY_CHECK);
        when(request.getParameter(DefaultRememberMeService.PARAMETER_REMEMBER_ME)).thenReturn("on");

        assertFalse(service.isCookieRequested(request));
    }

    @Test
    public void testIsCookieRequestedNoParameterProvided() {
        when(request.getRequestURI()).thenReturn("/stash/" + ApplicationConstants.URL_SECURITY_CHECK);
        when(request.getCookies()).thenReturn(new Cookie[]{});

        assertFalse(service.isCookieRequested(request));
    }

    @Test
    public void testIsCookieRequestedRememberMeAlways() {
        when(request.getRequestURI()).thenReturn("/stash/" + ApplicationConstants.URL_SECURITY_CHECK);
        when(request.getCookies()).thenReturn(new Cookie[]{});
        service = createService("always");
        // no parameter set

        assertTrue(service.isCookieRequested(request));
    }

    @Test
    public void testIsCookieRequestedRememberMeNever() {
        when(request.getRequestURI()).thenReturn("/stash/" + ApplicationConstants.URL_SECURITY_CHECK);
        when(request.getParameter(DefaultRememberMeService.PARAMETER_REMEMBER_ME)).thenReturn("on");
        when(request.getCookies()).thenReturn(new Cookie[]{});
        service = createService("never");

        assertFalse(service.isCookieRequested(request));
    }

    @Test
    public void testIsCookieRequestedUnsupportedParameterProvided() {
        when(request.getRequestURI()).thenReturn("/stash/" + ApplicationConstants.URL_SECURITY_CHECK);
        when(request.getParameter(DefaultRememberMeService.PARAMETER_REMEMBER_ME)).thenReturn("yes please");
        when(request.getCookies()).thenReturn(new Cookie[]{});

        assertFalse(service.isCookieRequested(request));
    }

    @Test
    public void testIsCookieRequestedWrongUri() {
        when(request.getRequestURI()).thenReturn("/stash/login");
        when(request.getParameter(DefaultRememberMeService.PARAMETER_REMEMBER_ME)).thenReturn("on");
        when(request.getCookies()).thenReturn(new Cookie[]{});

        assertFalse(service.isCookieRequested(request));
    }

    @Test
    public void testLogout() {
        service.logout(request, response);

        verifyCookieCleared();
        verify(dao).deleteForSeries(SERIES);
    }

    @Test
    public void testLogoutInvalidCookie() {
        Cookie cookie = new Cookie(COOKIE_NAME, "not-properly-encoded");
        when(request.getCookies()).thenReturn(new Cookie[]{cookie});

        service.logout(request, response);

        verifyZeroInteractions(dao, response);
    }

    @Test
    public void testLogoutNoRememberMeCookie() {
        when(request.getCookies()).thenReturn(new Cookie[]{});

        service.logout(request, response);

        verifyZeroInteractions(dao, response);
    }

    @Test
    @FixedCurrentTime
    public void testCreateCookieWithMaxAge1Version0Set() {
        when(token.getExpiryDate()).thenReturn(Instant.now().plus(Seconds.ONE.toStandardDuration()).toDate());

        service.setCookie(request, response, token);

        assertCookieMaxAgeAndVersion(1, 0);
    }

    @Test
    @FixedCurrentTime
    public void testCreateCookieWithMaxAge0Version1Set() {
        when(token.getExpiryDate()).thenReturn(Instant.now().toDate());

        service.setCookie(request, response, token);

        assertCookieMaxAgeAndVersion(0, 1);
    }

    @Test
    @FixedCurrentTime
    public void testCreateCookieWithNegativeMaxAgeVersion1Set() {
        when(token.getExpiryDate()).thenReturn(Instant.now().minus(Seconds.THREE.toStandardDuration()).toDate());

        service.setCookie(request, response, token);

        assertCookieMaxAgeAndVersion(-3, 1);
    }

    @Test
    public void testCreateCookieWithNullToken() {
        service.setCookie(request, response, null);

        assertCookieMaxAgeAndVersion(0, 1);
    }

    private void assertCookie(String series, String value, int expirySecs) {
        verify(response).addCookie(cookieCaptor.capture());
        Cookie cookie = cookieCaptor.getValue();

        assertEquals(COOKIE_NAME, cookie.getName());
        assertTrue(cookie.getMaxAge() <= expirySecs && cookie.getMaxAge() + 5 > expirySecs);
        assertTrue(cookie.isHttpOnly());
        assertTrue(cookie.getSecure());
        assertNotNull(cookie.getValue());
        String[] decoded = service.decodeCookieValue(cookie.getValue());
        assertEquals(series, decoded[0]);
        assertEquals(value, decoded[1]);
    }

    private void assertCookieMaxAgeAndVersion(int expectedMaxAge, int expectedVersion) {
        verify(response).addCookie(cookieCaptor.capture());
        Cookie cookie = cookieCaptor.getValue();

        assertEquals("Cookie has wrong max age", expectedMaxAge, cookie.getMaxAge());
        assertEquals("Cookie has wrong version", expectedVersion, cookie.getVersion());
    }

    private DefaultRememberMeService createService(String rememberMeMode) {
        return new DefaultRememberMeService(COOKIE_NAME, dao, EXPIRY_DAYS, 60, rememberMeMode, requestInfoHelper,
                tokenGenerator, transactionManager);
    }

    private void verifyCookieCleared() {
        verify(response).addCookie(cookieCaptor.capture());

        Cookie cookie = cookieCaptor.getValue();
        assertEquals(COOKIE_NAME, cookie.getName());
        assertEquals(0, cookie.getMaxAge());
        assertNull(cookie.getValue());
    }

    private static Matcher<InternalRememberMeToken> isSeriesAndToken(final String series, final String value) {
        return new BaseTokenMatcher() {
            @Override
            protected boolean tokenMatches(InternalRememberMeToken token) {
                return token.getSeries().matches(series) && token.getToken().matches(value);
            }
        };
    }

    private abstract static class BaseTokenMatcher extends BaseMatcher<InternalRememberMeToken> {
        @Override
        public boolean matches(Object item) {
            if (item instanceof InternalRememberMeToken) {
                InternalRememberMeToken token = (InternalRememberMeToken) item;
                return tokenMatches(token);
            }
            return false;
        }

        @Override
        public void describeTo(Description description) {
        }

        protected abstract  boolean tokenMatches(InternalRememberMeToken token);
    }
}