package com.atlassian.bitbucket.mesh.auth;

import com.atlassian.bitbucket.mesh.annotation.Profiled;
import com.atlassian.bitbucket.mesh.config.ApplicationConfig;
import com.atlassian.bitbucket.mesh.config.AuthenticationConfig;
import com.atlassian.bitbucket.mesh.config.NodeConfig;
import com.atlassian.bitbucket.mesh.management.SidebandManager;
import com.atlassian.bitbucket.mesh.network.MeshManager;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcNodeKeyRequest;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcSidebandFollowerFragment;
import com.atlassian.bitbucket.mesh.util.Gate;
import io.jsonwebtoken.*;
import io.jsonwebtoken.security.SignatureException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import java.security.Key;
import java.security.PublicKey;
import java.util.Date;

import static com.atlassian.bitbucket.mesh.rpc.RpcConstants.Token.*;
import static com.google.common.base.MoreObjects.firstNonNull;
import static java.util.Objects.requireNonNull;

@Profiled
public class JwtAuthenticator implements Authenticator {

    //When a node ID isn't provided or isn't available (some types of exceptions don't expose the JWS headers)
    //these keys are used as placeholders for the failure gate. They don't have any _meaning_; the intention
    //is simply that they be a) different from each other, so messages will be gated separately and b) obscure
    //enough that they're unlikely to be used as actual node IDs
    private static final String KEY_MALFORMED = "*";
    private static final String KEY_SIGNATURE = "$";

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

    private final AuthenticationConfig authenticationConfig;
    private final AuthenticationManager authenticationManager;
    private final Gate<String> failureGate;
    private final SigningKeyResolver keyResolver;
    private final MeshManager meshManager;
    private final NodeConfig nodeConfig;
    private final SidebandManager sidebandManager;
    private final Gate<String> keyRequestGate;

    public JwtAuthenticator(AuthenticationManager authenticationManager, ApplicationConfig config,
                            PublicKeyRegistry keyRegistry, MeshManager meshManager, SidebandManager sidebandManager) {
        this.authenticationManager = authenticationManager;

        authenticationConfig = config.getAuthentication();
        failureGate = new Gate<>(authenticationConfig.getFailureLogInterval());
        keyRequestGate = new Gate<>(authenticationConfig.getKeyRequestInterval());
        keyResolver = new RegistrySigningKeyResolver(keyRegistry);
        this.meshManager = meshManager;
        nodeConfig = config.getNode();
        this.sidebandManager = sidebandManager;
    }

    @Override
    public Principal authenticate(@Nonnull String subject, @Nonnull String token) {
        requireNonNull(subject, "subject");
        requireNonNull(token, "token");

        try {
            Jws<Claims> jws = Jwts.parserBuilder()
                    .requireSubject(subject)
                    .setAllowedClockSkewSeconds(authenticationConfig.getAllowedClockSkew().getSeconds())
                    .setSigningKeyResolver(keyResolver)
                    .build()
                    .parseClaimsJws(token);

            return new JwsPrincipal(jws);
        } catch (ExpiredJwtException e) {
            //The JWT iat+exp combination expired, factoring in allowed clock skew, before it was parsed
            String nodeId = getNodeId(e);
            logFailure(nodeId, "Rejecting expired token from {}", nodeId);
        } catch (IncorrectClaimException e) {
            //The JWT is for another endpoint
            String nodeId = getNodeId(e);
            logFailure(nodeId, "Rejecting {} token from {}; expected {}", e.getClaims().getSubject(), nodeId, subject);
        } catch (MalformedJwtException | UnsupportedJwtException e) {
            //The JWT isn't parseable or isn't understood
            logFailure(KEY_MALFORMED, "Rejecting invalid token", e);
        } catch (MissingClaimException e) {
            //The JWT didn't specify a subject
            String nodeId = getNodeId(e);
            logFailure(nodeId, "Rejecting token from {}; no subject was specified", nodeId);
        } catch (MissingNodeKeyException e) {
            logFailure(KEY_SIGNATURE, "Rejecting token; signature verification failed due to missing public " +
                    "key", e);
            String nodeId = e.getNodeId();
            if (!NID_CONTROL_PLANE.equals(nodeId) && meshManager.findNodeById(nodeId) != null) {
                if (keyRequestGate.maybeRun(nodeId, () -> sidebandManager.send(RpcSidebandFollowerFragment.newBuilder()
                        .setNodeKeyRequest(RpcNodeKeyRequest.newBuilder()
                                .setNodeId(nodeId)
                                .setFingerprint(e.getFingerprint())
                                .build())
                        .build()))) {
                    logFailure(KEY_SIGNATURE, "Sent request to the control plane to configure public key for " +
                            "the node with ID {}", nodeId);
                }
            }
        } catch (SignatureException e) {
            //The JWT signature couldn't be verified
            logFailure(KEY_SIGNATURE, "Rejecting token; signature verification failed", e);
        }

        return null;
    }

    @Nonnull
    @Override
    public String sign(@Nonnull String subject) {
        requireNonNull(subject, "subject");

        SigningKey signingKey = authenticationManager.getSigningKey();
        Date iat = new Date();

        JwtBuilder builder = Jwts.builder()
                .setHeaderParam(HDR_FP, signingKey.getFingerprint())
                .setHeaderParam(HDR_NID, nodeConfig.getId())
                .setIssuedAt(iat)
                .setSubject(subject)
                .signWith(signingKey.getPrivate()); //Automatically selects algorithm based on key type/strength
        if (authenticationConfig.isExpiryEnabled()) {
            builder.setExpiration(new Date(iat.getTime() + authenticationConfig.getExpiryInterval().toMillis()));
        }
        return builder.compact();
    }

    /**
     * Gets the {@code "nid"} header from the provided exception. If the header is not included the control plane
     * ({@code "control-plane"}) is assumed.
     *
     * @param e the exception to extract the header from
     * @return the extracted header, or {@code "control-plane"} if the header was not a string
     */
    private static String getNodeId(ClaimJwtException e) {
        Object nodeId = e.getHeader().get(HDR_NID);

        return nodeId instanceof String ? (String) nodeId : NID_CONTROL_PLANE;
    }

    /**
     * Logs a failure, automatically selecting between {@code WARN} and {@code DEBUG} levels based on whether another
     * failure has been logged recently.
     *
     * @param key        the gate key, used to differentiate between clients
     * @param message    the message, which may include format placeholders
     * @param parameters parameters, if any, for the message
     */
    private void logFailure(String key, String message, Object... parameters) {
        if (!failureGate.maybeRun(key, () -> log.warn(message, parameters))) {
            log.debug(message, parameters);
        }
    }

    private static class JwsPrincipal implements Principal {

        private final String nodeId;
        private final ApplicationRole role;

        JwsPrincipal(Jws<Claims> jws) {
            nodeId = (String) jws.getHeader().get(HDR_NID);
            role = nodeId == null || NID_CONTROL_PLANE.equals(nodeId) ?
                    ApplicationRole.CONTROL_PLANE :
                    ApplicationRole.PEER;
        }

        @Override
        public String getId() {
            return nodeId;
        }

        @Nonnull
        @Override
        public ApplicationRole getRole() {
            return role;
        }
    }

    private static class RegistrySigningKeyResolver implements SigningKeyResolver {

        private final PublicKeyRegistry registry;

        RegistrySigningKeyResolver(PublicKeyRegistry registry) {
            this.registry = registry;
        }

        @Override
        public Key resolveSigningKey(JwsHeader header, Claims claims) {
            Object fp = header.get(HDR_FP);
            Object nid = header.get(HDR_NID); //May be omitted if the request is from the control plane
            if (fp instanceof String && (nid == null || nid instanceof String)) {
                String nodeId = firstNonNull((String) nid, NID_CONTROL_PLANE);
                String fingerprint = (String) fp;
                PublicKey key = registry.findByFingerprint(nodeId, fingerprint);
                if (key != null) {
                    return key;
                }
                throw new MissingNodeKeyException(nodeId, fingerprint);
            }

            //The SigningKeyResolver SPI does not indicate nullability, but the calling JJWT code expects
            //it to always return a key or throw an exception; it can't return null
            throw new SignatureException("No signing key could be resolved for " + nid + " with fingerprint " + fp);
        }

        @Override
        public Key resolveSigningKey(JwsHeader header, String plaintext) {
            throw new UnsupportedJwtException("Plaintext tokens are not supported");
        }
    }
}
