/*
 * Copyright 2016 Red Hat, Inc. and/or its affiliates
 * and other contributors as indicated by the @author tags.
 *
 * 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 org.keycloak;

import org.keycloak.common.VerificationException;
import org.keycloak.exceptions.TokenNotActiveException;
import org.keycloak.exceptions.TokenSignatureInvalidException;
import org.keycloak.jose.jws.AlgorithmType;
import org.keycloak.jose.jws.JWSHeader;
import org.keycloak.jose.jws.JWSInput;
import org.keycloak.jose.jws.JWSInputException;
import org.keycloak.jose.jws.crypto.HMACProvider;
import org.keycloak.jose.jws.crypto.RSAProvider;
import org.keycloak.representations.AccessToken;
import org.keycloak.representations.JsonWebToken;
import org.keycloak.util.TokenUtil;

import javax.crypto.SecretKey;
import java.security.PublicKey;
import java.util.*;

/**
 * @author <a href="mailto:bill@burkecentral.com">Bill Burke</a>
 * @version $Revision: 1 $
 */
public class TokenVerifier<T extends JsonWebToken> {

    // This interface is here as JDK 7 is a requirement for this project.
    // Once JDK 8 would become mandatory, java.util.function.Predicate would be used instead.

    // @FunctionalInterface
    public static interface Predicate<T extends JsonWebToken> {
        /**
         * Performs a single check on the given token verifier.
         * @param t Token, guaranteed to be non-null.
         * @return
         * @throws VerificationException
         */
        boolean test(T t) throws VerificationException;
    }

    public static final Predicate<JsonWebToken> SUBJECT_EXISTS_CHECK = new Predicate<JsonWebToken>() {
        @Override
        public boolean test(JsonWebToken t) throws VerificationException {
            String subject = t.getSubject();
            if (subject == null) {
                throw new VerificationException("Subject missing in token");
            }

            return true;
        }
    };

    public static final Predicate<JsonWebToken> IS_ACTIVE = new Predicate<JsonWebToken>() {
        @Override
        public boolean test(JsonWebToken t) throws VerificationException {
            if (! t.isActive()) {
                throw new TokenNotActiveException("Token is not active");
            }

            return true;
        }
    };

    public static class RealmUrlCheck implements Predicate<JsonWebToken> {

        private static final RealmUrlCheck NULL_INSTANCE = new RealmUrlCheck(null);

        private final String realmUrl;

        public RealmUrlCheck(String realmUrl) {
            this.realmUrl = realmUrl;
        }

        @Override
        public boolean test(JsonWebToken t) throws VerificationException {
            if (this.realmUrl == null) {
                throw new VerificationException("Realm URL not set");
            }

            if (! this.realmUrl.equals(t.getIssuer())) {
                throw new VerificationException("Invalid token issuer. Expected '" + this.realmUrl + "', but was '" + t.getIssuer() + "'");
            }

            return true;
        }
    };

    public static class TokenTypeCheck implements Predicate<JsonWebToken> {

        private static final TokenTypeCheck INSTANCE_BEARER = new TokenTypeCheck(TokenUtil.TOKEN_TYPE_BEARER);

        private final String tokenType;

        public TokenTypeCheck(String tokenType) {
            this.tokenType = tokenType;
        }

        @Override
        public boolean test(JsonWebToken t) throws VerificationException {
            if (! tokenType.equalsIgnoreCase(t.getType())) {
                throw new VerificationException("Token type is incorrect. Expected '" + tokenType + "' but was '" + t.getType() + "'");
            }
            return true;
        }
    };

    private String tokenString;
    private Class<? extends T> clazz;
    private PublicKey publicKey;
    private SecretKey secretKey;
    private String realmUrl;
    private String expectedTokenType = TokenUtil.TOKEN_TYPE_BEARER;
    private boolean checkTokenType = true;
    private boolean checkRealmUrl = true;
    private final LinkedList<Predicate<? super T>> checks = new LinkedList<>();

    private JWSInput jws;
    private T token;

    protected TokenVerifier(String tokenString, Class<T> clazz) {
        this.tokenString = tokenString;
        this.clazz = clazz;
    }

    protected TokenVerifier(T token) {
        this.token = token;
    }

    /**
     * Creates a {@code TokenVerifier<AccessToken> instance. The method is here for backwards compatibility.
     * @param tokenString
     * @return
     * @deprecated use {@link #create(java.lang.String, java.lang.Class) } instead
     */
    public static TokenVerifier<AccessToken> create(String tokenString) {
        return create(tokenString, AccessToken.class);
    }

    public static <T extends JsonWebToken> TokenVerifier<T> create(String tokenString, Class<T> clazz) {
        return new TokenVerifier(tokenString, clazz)
          .check(RealmUrlCheck.NULL_INSTANCE)
          .check(SUBJECT_EXISTS_CHECK)
          .check(TokenTypeCheck.INSTANCE_BEARER)
          .check(IS_ACTIVE);
    }

    public static <T extends JsonWebToken> TokenVerifier<T> from(T token) {
        return new TokenVerifier(token)
          .check(RealmUrlCheck.NULL_INSTANCE)
          .check(SUBJECT_EXISTS_CHECK)
          .check(TokenTypeCheck.INSTANCE_BEARER)
          .check(IS_ACTIVE);
    }

    private void removeCheck(Class<? extends Predicate<?>> checkClass) {
        for (Iterator<Predicate<? super T>> it = checks.iterator(); it.hasNext();) {
            if (it.next().getClass() == checkClass) {
                it.remove();
            }
        }
    }

    private void removeCheck(Predicate<? super T> check) {
        checks.remove(check);
    }

    private <P extends Predicate<? super T>> TokenVerifier<T> replaceCheck(Class<? extends Predicate<?>> checkClass, boolean active, P predicate) {
        removeCheck(checkClass);
        if (active) {
            checks.add(predicate);
        }
        return this;
    }

    private <P extends Predicate<? super T>> TokenVerifier<T> replaceCheck(Predicate<? super T> check, boolean active, P predicate) {
        removeCheck(check);
        if (active) {
            checks.add(predicate);
        }
        return this;
    }

    /**
     * Resets all preset checks and will test the given checks in {@link #verify()} method.
     * @param checks
     * @return
     */
    public TokenVerifier<T> checkOnly(Predicate<? super T>... checks) {
        this.checks.clear();
        if (checks != null) {
            this.checks.addAll(Arrays.asList(checks));
        }
        return this;
    }

    /**
     * Will test the given checks in {@link #verify()} method in addition to already set checks.
     * @param checks
     * @return
     */
    public TokenVerifier<T> check(Predicate<? super T>... checks) {
        if (checks != null) {
            this.checks.addAll(Arrays.asList(checks));
        }
        return this;
    }

    public TokenVerifier<T> publicKey(PublicKey publicKey) {
        this.publicKey = publicKey;
        return this;
    }

    public TokenVerifier<T> secretKey(SecretKey secretKey) {
        this.secretKey = secretKey;
        return this;
    }

    public TokenVerifier<T> realmUrl(String realmUrl) {
        this.realmUrl = realmUrl;
        return replaceCheck(RealmUrlCheck.class, checkRealmUrl, new RealmUrlCheck(realmUrl));
    }

    public TokenVerifier<T> checkTokenType(boolean checkTokenType) {
        this.checkTokenType = checkTokenType;
        return replaceCheck(TokenTypeCheck.class, this.checkTokenType, new TokenTypeCheck(expectedTokenType));
    }

    public TokenVerifier<T> tokenType(String tokenType) {
        this.expectedTokenType = tokenType;
        return replaceCheck(TokenTypeCheck.class, this.checkTokenType, new TokenTypeCheck(expectedTokenType));
    }

    public TokenVerifier<T> checkActive(boolean checkActive) {
        return replaceCheck(IS_ACTIVE, checkActive, IS_ACTIVE);
    }

    public TokenVerifier<T> checkRealmUrl(boolean checkRealmUrl) {
        this.checkRealmUrl = checkRealmUrl;
        return replaceCheck(RealmUrlCheck.class, this.checkRealmUrl, new RealmUrlCheck(realmUrl));
    }

    public TokenVerifier<T> parse() throws VerificationException {
        if (jws == null) {
            if (tokenString == null) {
                throw new VerificationException("Token not set");
            }

            try {
                jws = new JWSInput(tokenString);
            } catch (JWSInputException e) {
                throw new VerificationException("Failed to parse JWT", e);
            }


            try {
                token = jws.readJsonContent(clazz);
            } catch (JWSInputException e) {
                throw new VerificationException("Failed to read access token from JWT", e);
            }
        }
        return this;
    }

    public T getToken() throws VerificationException {
        if (token == null) {
            parse();
        }
        return token;
    }

    public JWSHeader getHeader() throws VerificationException {
        parse();
        return jws.getHeader();
    }

    public void verifySignature() throws VerificationException {
        AlgorithmType algorithmType = getHeader().getAlgorithm().getType();

        if (null == algorithmType) {
            throw new VerificationException("Unknown or unsupported token algorithm");
        } else switch (algorithmType) {
            case RSA:
                if (publicKey == null) {
                    throw new VerificationException("Public key not set");
                }
                if (!RSAProvider.verify(jws, publicKey)) {
                    throw new TokenSignatureInvalidException("Invalid token signature");
                }   break;
            case HMAC:
                if (secretKey == null) {
                    throw new VerificationException("Secret key not set");
                }
                if (!HMACProvider.verify(jws, secretKey)) {
                    throw new TokenSignatureInvalidException("Invalid token signature");
                }   break;
            default:
                throw new VerificationException("Unknown or unsupported token algorithm");
        }
    }

    public TokenVerifier<T> verify() throws VerificationException {
        if (getToken() == null) {
            parse();
        }
        if (jws != null) {
            verifySignature();
        }

        for (Predicate<? super T> check : checks) {
            if (! check.test(getToken())) {
                throw new VerificationException("JWT check failed for check " + check);
            }
        }

        return this;
    }

}
