package io.quarkus.smallrye.jwt.runtime.auth;

import static io.undertow.util.Headers.AUTHORIZATION;
import static io.undertow.util.Headers.COOKIE;
import static io.undertow.util.Headers.WWW_AUTHENTICATE;
import static io.undertow.util.StatusCodes.UNAUTHORIZED;

import java.util.List;
import java.util.Locale;

import org.eclipse.microprofile.jwt.JsonWebToken;

import io.smallrye.jwt.auth.principal.JWTAuthContextInfo;
import io.undertow.UndertowLogger;
import io.undertow.security.api.AuthenticationMechanism;
import io.undertow.security.api.SecurityContext;
import io.undertow.security.idm.Account;
import io.undertow.security.idm.IdentityManager;
import io.undertow.server.HttpServerExchange;

/**
 * An AuthenticationMechanism that validates a caller based on a MicroProfile JWT bearer token
 */
public class JWTAuthMechanism implements AuthenticationMechanism {

    private JWTAuthContextInfo authContextInfo;
    private IdentityManager identityManager;

    public JWTAuthMechanism(JWTAuthContextInfo authContextInfo, IdentityManager identityManager) {
        this.authContextInfo = authContextInfo;
        this.identityManager = identityManager;
    }

    /**
     * Extract the Authorization header and validate the bearer token if it exists. If it does, and is validated, this
     * builds the org.jboss.security.SecurityContext authenticated Subject that drives the container APIs as well as
     * the authorization layers.
     *
     * @param exchange - the http request exchange object
     * @param securityContext - the current security context that
     * @return one of AUTHENTICATED, NOT_AUTHENTICATED or NOT_ATTEMPTED depending on the header and authentication outcome.
     */
    @Override
    public AuthenticationMechanismOutcome authenticate(HttpServerExchange exchange, SecurityContext securityContext) {
        String jwtToken = getJwtToken(exchange);
        if (jwtToken != null) {
            try {
                JWTCredential credential = new JWTCredential(jwtToken, authContextInfo);
                if (UndertowLogger.SECURITY_LOGGER.isTraceEnabled()) {
                    UndertowLogger.SECURITY_LOGGER.tracef("Bearer token: %s", jwtToken);
                }
                // Install the JWT principal as the caller
                Account account = identityManager.verify(credential.getName(), credential);
                if (account != null) {
                    securityContext.authenticationComplete(account, "MP-JWT", false);
                    UndertowLogger.SECURITY_LOGGER.debugf("Authenticated caller(%s) for path(%s) with roles: %s",
                            credential.getName(), exchange.getRequestPath(), account.getRoles());
                    return AuthenticationMechanismOutcome.AUTHENTICATED;
                } else {
                    UndertowLogger.SECURITY_LOGGER.info("Failed to authenticate JWT bearer token");
                    return AuthenticationMechanismOutcome.NOT_AUTHENTICATED;
                }
            } catch (Exception e) {
                UndertowLogger.SECURITY_LOGGER.infof(e, "Failed to validate JWT bearer token");
                return AuthenticationMechanismOutcome.NOT_AUTHENTICATED;
            }
        }

        // No suitable header has been found in this request,
        return AuthenticationMechanismOutcome.NOT_ATTEMPTED;
    }

    private String getJwtToken(HttpServerExchange exchange) {
        String bearerToken = null;
        if (AUTHORIZATION.toString().equals(authContextInfo.getTokenHeader())) {
            String authScheme = exchange.getRequestHeaders().getFirst(authContextInfo.getTokenHeader());
            if (authScheme != null && authScheme.toLowerCase(Locale.ENGLISH).startsWith("bearer ")) {
                bearerToken = authScheme.substring(7);
            }
        } else if (COOKIE.toString().equals(authContextInfo.getTokenHeader())
                && authContextInfo.getTokenCookie() != null) {
            List<String> cookieValues = exchange.getRequestHeaders().get(authContextInfo.getTokenHeader());
            if (cookieValues != null) {
                for (String cookieValue : cookieValues) {
                    for (String cookie : cookieValue.split(";")) {
                        if (cookie.startsWith(authContextInfo.getTokenCookie())) {
                            String[] parts = cookie.split("=");
                            if (parts.length == 2) {
                                bearerToken = parts[1].trim();
                            }
                        }
                    }
                }
            }
        } else {
            bearerToken = exchange.getRequestHeaders().getFirst(authContextInfo.getTokenHeader());
        }
        return bearerToken;
    }

    @Override
    public ChallengeResult sendChallenge(HttpServerExchange exchange, SecurityContext securityContext) {
        exchange.getResponseHeaders().add(WWW_AUTHENTICATE, "Bearer {token}");
        UndertowLogger.SECURITY_LOGGER.debugf("Sending Bearer {token} challenge for %s", exchange);
        return new ChallengeResult(true, UNAUTHORIZED);
    }

}

@@ -278,87 +278,13 @@
 il.L
-ist;%0Aimport java.util.Locale;%0A%0Aimport org.eclipse.microprofile.jwt.JsonWebToken
+ocale
 ;%0A%0Ai
@@ -616,16 +616,59 @@
 xchange;
+%0Aimport io.undertow.server.handlers.Cookie;
 %0A%0A/**%0A *
@@ -3919,33 +3919,21 @@
     
-List%3CString%3E
+Cookie
  cookie
-Values
  = e
@@ -3946,30 +3946,30 @@
 e.getRequest
-Header
+Cookie
 s().get(auth
@@ -3980,38 +3980,38 @@
 extInfo.getToken
-Header
+Cookie
 ());%0A           
@@ -4021,22 +4021,16 @@
  (cookie
-Values
  != null
@@ -4053,475 +4053,40 @@
     
-for (String cookieValue : cookieValues) %7B%0A                    for (String cookie : cookieValue.split(%22;%22)) %7B%0A                        if (cookie.startsWith(authContextInfo.getTokenCookie())) %7B%0A                            String%5B%5D parts = cookie.split(%22=%22);%0A                            if (parts.length == 2) %7B%0A                                bearerToken = parts%5B1%5D.trim();%0A                            %7D%0A                        %7D%0A                    %7D%0A                %7D
+bearerToken = cookie.getValue();
 %0A   

Use Undertow to parse the cookies