/*
 * Copyright (C) 2022 Authlete, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
 * either express or implied. See the License for the specific
 * language governing permissions and limitations under the
 * License.
 */
package com.authlete.jaxrs.server.api;


import javax.ws.rs.WebApplicationException;
import javax.ws.rs.core.CacheControl;
import javax.ws.rs.core.MediaType;
import javax.ws.rs.core.Response;
import javax.ws.rs.core.Response.Status;
import com.authlete.common.api.AuthleteApi;
import com.authlete.common.dto.TokenCreateRequest;
import com.authlete.common.dto.TokenCreateResponse;
import com.authlete.common.dto.TokenResponse;
import com.authlete.common.types.GrantType;
import com.nimbusds.jwt.JWT;
import com.nimbusds.jwt.JWTParser;
import com.nimbusds.jwt.SignedJWT;


/**
 * A sample implementation of processing a token request which uses the grant
 * type {@code "urn:ietf:params:oauth:grant-type:jwt-bearer"} (<a href=
 * "https://www.rfc-editor.org/rfc/rfc7523.html">RFC 7523</a>).
 *
 * <p>
 * The token request contains an {@code assertion} request parameter. Its value
 * is a JWT. However, RFC 7523 does not define details about how the JWT is
 * generated by whom. As a result, it is not defined in the specification how
 * to obtain the key whereby to verify the signature of the JWT. Therefore,
 * each deployment has to define their own rules which are necessary to
 * determine the key for signature verification.
 * </p>
 *
 * <p>
 * Note that your system must verify the signature of the assertion JWT by
 * itself. The JavaDoc of <a href=
 * "https://authlete.github.io/authlete-java-common/com/authlete/common/dto/TokenResponse.html"
 * >TokenResponse</a> explains (1) what validation steps Authlete performs on
 * behalf of your system and (2) why Authlete does not (can not) verify the
 * signature of the assertion JWT.
 * </p>
 *
 * @see <a href="https://www.rfc-editor.org/rfc/rfc7521.html">RFC 7521
 *      Assertion Framework for OAuth 2.0 Client Authentication and
 *      Authorization Grants</a>
 *
 * @see <a href="https://www.rfc-editor.org/rfc/rfc7523.html">RFC 7523
 *      JSON Web Token (JWT) Profile for OAuth 2.0 Client Authentication
 *      and Authorization Grants</a>
 */
class JwtAuthzGrantProcessor
{
    private final AuthleteApi mAuthleteApi;
    private final TokenResponse mTokenResponse;


    public JwtAuthzGrantProcessor(AuthleteApi authleteApi, TokenResponse tokenResponse)
    {
        mAuthleteApi   = authleteApi;
        mTokenResponse = tokenResponse;
    }


    public Response process()
    {
        try
        {
            return createResponse();
        }
        catch (WebApplicationException cause)
        {
            return cause.getResponse();
        }
    }


    private Response createResponse() throws WebApplicationException
    {
        // Validate the assertion.
        SignedJWT jwt = validateAssertion();

        // Client ID to assign.
        long clientId = determineClientId();

        // Scopes to assign.
        String[] scopes = determineScopes();

        // Subject to assign.
        String subject = determineSubject(jwt);

        // Create an access token.
        TokenCreateResponse tcResponse =
                createAccessToken(clientId, scopes, subject);

        // Create a successful token response.
        return createSuccessfulResponse(tcResponse);
    }


    private SignedJWT validateAssertion()
    {
        // The value of the 'assertion' request parameter.
        String assertion = mTokenResponse.getAssertion();

        // This implementation requires that the assertion is a signed JWT.
        SignedJWT jwt = parseAsSignedJwt(assertion);

        // When the assertion is a signed JWT, all validation steps common to
        // signed JWTs have been done on Authlete side except verification of
        // the signature. See the JavaDoc of TokenResponse class for details
        // about the validation steps performed on Authlete side.
        //
        //   https://authlete.github.io/authlete-java-common/com/authlete/common/dto/TokenResponse.html
        //

        // Verify the signature of the JWT.
        verifySignature(jwt);

        return jwt;
    }


    private SignedJWT parseAsSignedJwt(String assertion)
    {
        JWT jwt;

        try
        {
            // Parse the assertion as a JWT.
            jwt = JWTParser.parse(assertion);
        }
        catch (Exception cause)
        {
            throw invalidGrant("The assertion failed to be parsed as a JWT.");
        }

        // If the JWT is not a signed JWT.
        if (!(jwt instanceof SignedJWT))
        {
            throw invalidGrant(
                    "This authorization server requires that the assertion be a signed JWT.");
        }

        return (SignedJWT)jwt;
    }


    private void verifySignature(SignedJWT jwt)
    {
        // Because RFC 7523 does not define details about how the assertion
        // JWT is generated by whom. As a result, it is not defined in the
        // specification how to obtain the key whereby to verify the signature
        // of the JWT. Therefore, each deployment has to define their own rules
        // which are necessary to determine the key for signature verification.

        // Your system must define additional requirements about the assertion
        // so that your system can determine how to obtain the key for signature
        // verification.
        //
        // For example, your system may define a rule like below.
        //
        //   The value of the 'assertion' request parameter must be an ID Token
        //   issued by "https://example.com".
        //
        // If the assertion is an ID Token, it is possible to find the key for
        // signature verification by (1) getting the server configuration from
        // the discovery endpoint, (2) getting the JWK Set document from the
        // location indicated by the "jwks_uri" property in the server
        // configuration, and (3) selecting a key from among the JWK Set document.

        // TODO
        // In any case, your implementation must verify the signature of the JWT.
    }


    private long determineClientId()
    {
        // The client ID of the client that made the token request.
        long clientId = mTokenResponse.getClientId();

        // If 'Service.jwtGrantByIdentifiableClientsOnly' is false, token
        // requests that contain no client identifier are not rejected.
        // In that case, 'clientId' here becomes 0.
        //
        // However, this authorization server implementation does not allow
        // unidentifiable clients to make token requests with the grant type
        // "urn:ietf:params:oauth:grant-type:jwt-bearer" regardless of whether
        // 'Service.jwtGrantByIdentifiableClientsOnly' is true or false.
        if (clientId == 0)
        {
            throw invalidRequest(
                    "This authorization server does not allow unidentifiable " +
                    "clients to make token requests with the grant type " +
                    "'urn:ietf:params:oauth:grant-type:jwt-bearer'.");
        }

        // This simple implementation uses the client ID of the client
        // that made the token request.
        return clientId;
    }


    private String[] determineScopes()
    {
        // This simple implementation uses the scopes specified by the token request.
        return mTokenResponse.getScopes();
    }


    private String determineSubject(SignedJWT jwt)
    {
        try
        {
            // Get the value of the "sub" claim from the payload of the JWT.
            //
            // RFC 7523 requires that an assertion used with the grant type
            // "urn:ietf:params:oauth:grant-type:jwt-bearer" have the "sub"
            // claim.
            return jwt.getJWTClaimsSet().getSubject();
        }
        catch (Exception cause)
        {
            throw invalidGrant(
                    "The value of the 'sub' claim failed to be extracted " +
                    "from the payload of the assertion.");
        }
    }


    private TokenCreateResponse createAccessToken(
            long clientId, String[] scopes, String subject)
    {
        // A request to Authlete's /auth/token/create API.
        TokenCreateRequest request = new TokenCreateRequest()
                .setGrantType(GrantType.JWT_BEARER)
                .setClientId(clientId)
                .setScopes(scopes)
                .setSubject(subject)
                ;

        try
        {
            // Call Authlete's /auth/token/create API to create an access token.
            return mAuthleteApi.tokenCreate(request);
        }
        catch (Exception cause)
        {
            // API call to /auth/token/create failed.
            cause.printStackTrace();
            throw serverError("API call to /auth/token/create failed.");
        }
    }


    private Response createSuccessfulResponse(TokenCreateResponse tcResponse)
    {
        // The content of a successful token response that conforms to RFC 6749.
        String content = String.format(
                "{\n" +
                "  \"access_token\":\"%s\",\n" +
                "  \"token_type\":\"Bearer\",\n" +
                "  \"expires_in\":%d,\n" +
                "  \"scope\":\"%s\"\n" +
                "}\n",
                extractAccessToken(tcResponse),
                tcResponse.getExpiresIn(),
                buildScope(tcResponse)
                );

        return toJsonResponse(Status.OK, content);
    }


    private String extractAccessToken(TokenCreateResponse tcResponse)
    {
        // If a JWT access token has been issued, it takes precedence over
        // a random-string access token.

        // An access token in the JWT format. This response parameter holds
        // a non-null value when Service.accessTokenSignAlg is not null.
        String at = tcResponse.getJwtAccessToken();

        // If an access token in the JWT format has not been issued.
        if (at == null)
        {
            // An access token whose format is just a random string.
            at = tcResponse.getAccessToken();
        }

        // The newly issued access token.
        return at;
    }


    private String buildScope(TokenCreateResponse tcResponse)
    {
        String[] scopes = tcResponse.getScopes();

        if (scopes == null)
        {
            return "";
        }

        return String.join(" ", scopes);
    }


    private Response toJsonResponse(Status status, String content)
    {
        CacheControl cacheControl = new CacheControl();
        cacheControl.setNoCache(true);
        cacheControl.setNoStore(true);

        return Response
                .status(status)
                .type(MediaType.APPLICATION_JSON_TYPE)
                .cacheControl(cacheControl)
                .entity(content)
                .build();
    }


    private WebApplicationException toException(Status status, String error, String description)
    {
        String content = String.format(
                "{\n" +
                "  \"error\":\"%s\",\n" +
                "  \"error_description\":\"%s\"\n" +
                "}\n",
                error, description);

        Response response = toJsonResponse(status, content);

        return new WebApplicationException(response);
    }


    private WebApplicationException invalidGrant(String message)
    {
        return toException(Status.BAD_REQUEST, "invalid_grant", message);
    }


    private WebApplicationException invalidRequest(String message)
    {
        return toException(Status.BAD_REQUEST, "invalid_request", message);
    }


    private WebApplicationException serverError(String message)
    {
        return toException(Status.INTERNAL_SERVER_ERROR, "server_error", message);
    }
}
