package com.atlassian.stash.internal.auth;

import com.atlassian.security.random.SecureTokenGenerator;
import com.atlassian.bitbucket.auth.RememberMeMode;
import com.atlassian.stash.internal.ApplicationConstants;
import com.atlassian.stash.internal.InternalConverter;
import com.atlassian.stash.internal.spring.SpringTransactionUtils;
import com.atlassian.stash.internal.web.HttpRequestInfoHelper;
import com.atlassian.bitbucket.user.ApplicationUser;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.primitives.Ints;
import org.apache.commons.codec.binary.Base64;
import org.apache.commons.lang3.StringUtils;
import org.hibernate.HibernateException;
import org.joda.time.DateTimeUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.dao.DataAccessException;
import org.springframework.stereotype.Service;
import org.springframework.transaction.PlatformTransactionManager;
import org.springframework.transaction.annotation.Propagation;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.transaction.support.TransactionTemplate;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import javax.servlet.http.Cookie;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.Objects;
import java.util.concurrent.TimeUnit;

import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkNotNull;

@Service("rememberMeService")
@Transactional
public class DefaultRememberMeService implements RememberMeService {

    private static final Logger log = LoggerFactory.getLogger(DefaultRememberMeService.class);

    static final int MAX_ATTEMPTS = 5;
    static final String PARAMETER_REMEMBER_ME = "_atl_remember_me";

    private final String cookieName;
    private final RememberMeTokenDao dao;
    private final long expirySeconds;
    private final long gracePeriodSeconds;
    private final RememberMeMode rememberMeMode;
    private final HttpRequestInfoHelper requestInfoHelper;
    private final SecureTokenGenerator tokenGenerator;
    private final TransactionTemplate transactionTemplate;

    @Autowired
    public DefaultRememberMeService(@Value("${auth.remember-me.cookie.name}") String cookieName, RememberMeTokenDao dao,
                                    @Value("${auth.remember-me.token.expiry}") int expiryDays,
                                    @Value("${auth.remember-me.token.grace.period}") long gracePeriodSeconds,
                                    @Value("${auth.remember-me.enabled}") String rememberMeMode,
                                    HttpRequestInfoHelper requestInfoHelper, SecureTokenGenerator tokenGenerator,
                                    PlatformTransactionManager transactionManager) {
        this.cookieName = cookieName;
        this.dao = dao;
        this.expirySeconds = TimeUnit.DAYS.toSeconds(expiryDays);
        this.gracePeriodSeconds = gracePeriodSeconds;
        this.rememberMeMode = RememberMeMode.fromId(rememberMeMode);
        this.requestInfoHelper = requestInfoHelper;
        this.tokenGenerator = tokenGenerator;
        this.transactionTemplate = new TransactionTemplate(transactionManager, SpringTransactionUtils.REQUIRES_NEW);
    }

    @Nullable
    @Override
    public ApplicationUser authenticate(@Nonnull final HttpServletRequest request, @Nonnull HttpServletResponse response) {
        Cookie cookie = getCookie(request);
        if (cookie == null) {
            log.trace("Remember-me cookie not found, skipping remember-me authentication");
            return null;
        }

        InternalRememberMeToken token = toToken(cookie);
        if (token == null) {
            log.info("Invalid remember-me cookie detected (wrong format: '{}') - canceling the cookie", cookie.getValue());
            setCookie(request, response, null);
            return null;
        }

        if (rememberMeMode == RememberMeMode.NEVER) {
            log.trace("Remember-me has been disabled, skipping remember-me authentication");
            return null;
        }

        // Each RememberMeToken can be claimed exactly once. When that succeeds the token is marked as claimed
        InternalRememberMeToken validToken = claimToken(token, getRemoteAddress(request));

        if (validToken == null) {
            // provided token was invalid
            log.info("Invalid remember-me cookie detected (expired) - canceling the cookie");
            setCookie(request, response, null);
            return null;
        }

        if (!validToken.getToken().equals(token.getToken())) {
            // only update the the cookie if the token value has changed. The value won't change if the provided token
            // has already been claimed, but is still in its grace period.
            setCookie(request, response, validToken);
        }

        log.trace("Remember-me authentication succeeded for series '{}' and user '{}'",
                validToken.getSeries(), getUsername(validToken));

        return validToken.getUser();
    }

    @Override
    @Transactional(propagation = Propagation.SUPPORTS) // no transaction required
    public void cancelCookie(@Nonnull HttpServletRequest request, @Nonnull HttpServletResponse response) {
        Cookie cookie = getCookie(request);
        if (cookie != null) {
            if (log.isTraceEnabled()) {
                InternalRememberMeToken token = toToken(cookie);
                if (token == null) {
                    log.trace("Canceling invalid remember-me cookie");
                } else {
                    log.trace("Canceling remember-me cookie for series '{}'", token.getSeries());
                }
            }

            setCookie(request, response, null);
        }
    }

    @Override
    public void createCookie(@Nonnull ApplicationUser user, @Nonnull HttpServletRequest request, @Nonnull HttpServletResponse response) {
        // To prevent security incidents, never give out a remember-me cookie for a new series if a cookie was provided
        // on the request. By maintaining this invariant, it's impossible to use a stolen cookie to obtain a fresh token
        // that's not part of the intercepted series. As long as the series remains the same, the cookie theft will be
        // detected the next time the real user tries to authenticate with his remember-me cookie. When that happens,
        // all cookies for that series are cleared
        Cookie cookie = getCookie(request);
        if (cookie != null) {
            cancelCookie(request, response);
            InternalRememberMeToken token = toToken(cookie);
            throw new IllegalStateException("A remember-me cookie for series '+" +
                    (token != null ? token.getSeries() : "invalid") +
                    "' is already present. Cannot provide a remember-me cookie for a new series. " +
                    "Canceling the existing cookie");
        }

        InternalRememberMeToken token = dao.create(new InternalRememberMeToken.Builder()
                .series(tokenGenerator.generateToken())
                .token(tokenGenerator.generateToken())
                .user(InternalConverter.convertToInternalUser(user))
                .expiresAfter(expirySeconds, TimeUnit.SECONDS)
                .build());

        setCookie(request, response, token);
        log.debug("Created new remember-me series '{}' for user '{}'", token.getSeries(), user.getName());
    }

    @Override
    public void deleteExpiredTokens(long maxExpiredTime, @Nonnull TimeUnit timeUnit) {
        checkArgument(maxExpiredTime >= 0, "Negative expiry times are not allowed");
        int count = dao.deleteExpiredTokens(maxExpiredTime, timeUnit);
        log.debug("Cleaned up {} expired remember-me tokens", count);
    }

    @Override
    public void deleteForUser(@Nonnull ApplicationUser user) {
        dao.deleteForUser(checkNotNull(user, "user").getId());
    }

    @Transactional(propagation = Propagation.SUPPORTS) // no transaction required
    public boolean isCookieRequested(@Nonnull HttpServletRequest request) {
        String uri = StringUtils.substringBefore(request.getRequestURI(), ";");
        if (!uri.endsWith(ApplicationConstants.URL_SECURITY_CHECK)) {
            return false;
        }

        Cookie cookie = getCookie(request);
        if (cookie != null) {
            // should never hand out a new cookie (for a new series) if a remember-me cookie was already provided on
            // the request
            log.debug("Not providing a remember-me cookie because a cookie was already present on the request");
            return false;
        }

        switch (rememberMeMode) {
            case ALWAYS:
                return true;
            case NEVER:
                return false;
            default:
                String rememberMe = request.getParameter(PARAMETER_REMEMBER_ME);
                return "on".equals(rememberMe) || "checked".equals(rememberMe) || "1".equals(rememberMe) ||
                        "true".equalsIgnoreCase(rememberMe);
        }
    }

    @Override
    public void logout(@Nonnull HttpServletRequest request, @Nonnull HttpServletResponse response) {
        Cookie cookie = getCookie(request);
        if (cookie != null) {
            InternalRememberMeToken token = toToken(cookie);
            if (token != null) {
                log.trace("Canceling remember-me series '{}' on logout", token.getSeries());

                dao.deleteForSeries(token.getSeries());
                cancelCookie(request, response);
            }
        }
    }

    @VisibleForTesting
    protected String[] decodeCookieValue(String cookieValue) {
        String base64 = cookieValue + StringUtils.repeat('=', cookieValue.length() % 4);
        if (!Base64.isBase64(base64)) {
            return null;
        }
        String decoded = new String(Base64.decodeBase64(base64.getBytes()));
        return StringUtils.split(decoded, ':');
    }

    @VisibleForTesting
    protected String encodeCookie(String... cookieTokens) {
        String joined = StringUtils.join(cookieTokens, ":");
        String encoded = new String(Base64.encodeBase64(joined.getBytes()));
        return StringUtils.stripEnd(encoded, "=");
    }

    private InternalRememberMeToken claimToken(final InternalRememberMeToken token, final String remoteAddress) {
        int i = 0;
        while (i < MAX_ATTEMPTS) {
            try {
                return transactionTemplate.execute(status -> doClaimToken(token, remoteAddress));
            } catch (HibernateException | DataAccessException e) {
                if (++i < MAX_ATTEMPTS) {
                    log.debug("Could not claim token for series {}: '{}'. Retrying", token.getSeries(), e.getMessage());
                } else {
                    log.error("Could not claim token for series {} after {} attempts", token.getSeries(), i, e);
                }
            }
        }
        return null;
    }

    private InternalRememberMeToken doClaimToken(InternalRememberMeToken token, String remoteAddress) {
        InternalRememberMeToken current = dao.findBySeriesAndToken(token.getSeries(), token.getToken());
        if (current == null) {
            return null;
        }

        if (current.isClaimed()) {
            return handleClaimedToken(current, remoteAddress);
        }

        String username = getUsername(current);
        if (current.isExpired()) {
            log.debug("Provided token for remember-me series '{}' for user '{}' has expired. Ignoring the token",
                    token.getSeries(), username);
            return null;
        }

        if (dao.claim(new InternalRememberMeToken.Builder(current)
                .claimed()
                .expiresAfter(gracePeriodSeconds, TimeUnit.SECONDS)
                .remoteAddress(remoteAddress)
                .build())) {

            log.debug("Refreshed token for remember-me series '{}' for user '{}'", token.getSeries(), username);
            // Token has been claimed and it's grace period has started. Create a new token for the series and
            // return it
            return dao.create(new InternalRememberMeToken.Builder(current)
                    .id(ApplicationConstants.UNSAVED_ENTITY)
                    .token(tokenGenerator.generateToken())
                    .expiresAfter(expirySeconds, TimeUnit.SECONDS)
                    .build());
        }

        // claim failed - another thread must have claimed the token concurrently
        dao.refresh(current);
        if (current.isClaimed()) {
            return handleClaimedToken(current, remoteAddress);
        }

        return null;
    }

    private Cookie getCookie(HttpServletRequest request) {
        Cookie[] cookies = request.getCookies();
        if (cookies == null) {
            return null;
        }
        for (Cookie cookie : cookies) {
            if (cookieName.equals(cookie.getName())) {
                return cookie;
            }
        }
        return null;
    }

    private String getRemoteAddress(HttpServletRequest request) {
        String remoteAddress = requestInfoHelper.getRemoteAddress(request);
        // if the remote address > 255, only return the first (most significant) 255 chars
        return remoteAddress.length() < 256 ? remoteAddress : remoteAddress.substring(0, 255);
    }

    private String getUsername(InternalRememberMeToken token) {
        return StringUtils.defaultString(token.getUser() == null ? null : token.getUser().getName(), "<unknown>");
    }

    private InternalRememberMeToken handleClaimedToken(InternalRememberMeToken token, String remoteAddress) {
        boolean remoteAddressMatch = Objects.equals(remoteAddress, token.getRemoteAddress());
        if (remoteAddressMatch && !token.isExpired()) {
            log.trace("Accepting remember-me token for series '{}' for user '{}' that was already used for " +
                            "authentication from '{}'",
                    token.getSeries(), getUsername(token), token.getRemoteAddress());

            // Token has already been used for authentication and is still valid. Because we're paranoid, we don't want
            // to return the new token again. Presumably a parallel request from the same browser has already processed
            // it and updated the remember-me cookie. We'll return the old token here, which should result in no
            // Set-Cookie header being set on the response.
            return token;
        }

        // Token has already been used for authentication. This could mean that the cookie has been stolen and
        // used to authenticate. As a safety precaution, we'll kill the whole series and force the user to
        // authenticate again.
        int deletedTokens = dao.deleteForSeries(token.getSeries());
        if (deletedTokens > 0) {
            // at least one token dropped
            if (remoteAddressMatch) {
                log.info("Expired remember-me token detected for series '{}' for user '{}' (used from '{}'). As a safety " +
                                "precaution, all ({}) tokens from that series have been canceled.",
                        token.getSeries(), getUsername(token), token.getRemoteAddress(), deletedTokens);
            } else {
                log.info("Remember-me token detected for series '{}' for user '{}' (used from '{}'). Current request is " +
                        "from a different address ('{}')! As a safety precaution, all ({}) tokens from that series " +
                        "have been canceled.", token.getSeries(), getUsername(token), token.getRemoteAddress(),
                        remoteAddress, deletedTokens);
            }
        }
        return null;
    }

    @VisibleForTesting
    void setCookie(HttpServletRequest request, HttpServletResponse response, InternalRememberMeToken token) {
        String cookieValue = null;
        if (token != null) {
            cookieValue = encodeCookie(token.getSeries(), token.getToken());
        }

        Cookie cookie = new Cookie(cookieName, cookieValue);
        cookie.setHttpOnly(true);
        cookie.setSecure(request.isSecure());

        String path = request.getContextPath();
        cookie.setPath(StringUtils.isNotBlank(path) ? path : "/");
        // DateTimeUtils is used to get current time as it could be configured in unit tests to cover edge values of maxAge
        long maxAge = token == null ? 0 :
                TimeUnit.MILLISECONDS.toSeconds(token.getExpiryDate().getTime() - DateTimeUtils.currentTimeMillis());

        cookie.setMaxAge(Ints.saturatedCast(maxAge));

        if (maxAge < 1) {
            cookie.setVersion(1);
        }

        response.addCookie(cookie);
    }

    private InternalRememberMeToken toToken(Cookie cookie) {
        String[] tokens = decodeCookieValue(cookie.getValue());
        if (tokens == null || tokens.length != 2) {
            return null;
        }

        return new InternalRememberMeToken.Builder()
                .series(tokens[0])
                .token(tokens[1])
                .build();
    }
}
