package x.y.z;

import lombok.extern.log4j.Log4j2;
import org.springframework.security.core.Authentication;
import org.springframework.security.web.authentication.SimpleUrlAuthenticationSuccessHandler;

import javax.servlet.ServletException;
import javax.servlet.http.Cookie;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;

@Log4j2
public class AuthenticationSuccessHandlerImpl extends SimpleUrlAuthenticationSuccessHandler {
    @Override
    public void onAuthenticationSuccess(HttpServletRequest request, HttpServletResponse response, Authentication authentication)
            throws IOException, ServletException {
        addSameSiteAttributeToCookie(request, response);
        handle(request, response, authentication);
        clearAuthenticationAttributes(request);
    }

    private void addSameSiteAttributeToCookie(HttpServletRequest request, HttpServletResponse response) {
        for (final Cookie cookie : request.getCookies()) {
            if (cookie.getName().equals("JSESSIONID")) {
                final String value = cookie.getValue();

                log.info("Cookie value: {}", value);
                response.addHeader("Set-Cookie", getCookieWithSameSite(value, "Strict"));
            }
        }
    }

    private String getCookieWithSameSite(String value, String mode) {
        StringBuilder builder = new StringBuilder();
        builder.append("JSESSIONID").append('=').append(value)
                .append(";Path=/")
                .append(";HttpOnly")
                .append(";Secure")
                .append(";SameSite=").append(mode);
        return builder.toString();
    }
}