package com.atlassian.bitbucket.mesh.auth;

import com.atlassian.bitbucket.mesh.health.rpc.v1.HealthGrpc;
import com.atlassian.bitbucket.mesh.rpc.CapabilityServiceGrpc;
import com.atlassian.bitbucket.mesh.rpc.v1.ManagementServiceGrpc;
import io.grpc.*;

import static com.atlassian.bitbucket.mesh.rpc.MetadataKeys.KEY_TOKEN;

/**
 * A {@link ServerInterceptor} which extracts a token header from incoming requests and uses an {@link Authenticator}
 * to authenticate it. If the token fails validation, the request is rejected.
 */
public class AuthenticationServerInterceptor implements ServerInterceptor {

    static final String METHOD_GET_CAPABILITIES =
            CapabilityServiceGrpc.getGetCapabilitiesMethod().getFullMethodName();
    static final String METHOD_HEALTH_CHECK = HealthGrpc.getCheckMethod().getFullMethodName();
    static final String METHOD_HEALTH_WATCH = HealthGrpc.getWatchMethod().getFullMethodName();
    static final String METHOD_PROVISION_NODE =
            ManagementServiceGrpc.getProvisionNodeMethod().getFullMethodName();

    private final AuthenticationManager authenticationManager;
    private final Authenticator authenticator;

    public AuthenticationServerInterceptor(AuthenticationManager authenticationManager, Authenticator authenticator) {
        this.authenticationManager = authenticationManager;
        this.authenticator = authenticator;
    }

    @Override
    public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(ServerCall<ReqT, RespT> call, Metadata headers,
                                                                 ServerCallHandler<ReqT, RespT> next) {
        String methodName = call.getMethodDescriptor().getFullMethodName();
        if (allowAnonymous(methodName)) {
            //If the method being called allows anonymous access, don't bother validating the token
            return next.startCall(call, headers);
        }

        String token = headers.get(KEY_TOKEN);
        if (token == null) {
            return rejected(call, "No [" + KEY_TOKEN.name() + "] was sent");
        }

        Principal principal = authenticator.authenticate(methodName, token);
        if (principal == null) {
            return rejected(call, "The [" + KEY_TOKEN.name() + "] is invalid");
        }

        return authenticationManager.withPrincipal(principal, ignored -> next.startCall(call, headers));
    }

    private static <ReqT, RespT> ServerCall.Listener<ReqT> rejected(ServerCall<ReqT, RespT> call, String reason) {
        call.close(Status.UNAUTHENTICATED.withDescription(reason), new Metadata());

        //Nothing to do; we've already rejected the request
        return new ServerCall.Listener<ReqT>() { };
    }

    private boolean allowAnonymous(String methodName) {
        //CapabilityService/GetCapabilities is always allowed because it provides capabilities information a client
        //may need in order to decide what calls can be made
        return METHOD_GET_CAPABILITIES.equals(methodName) ||
                METHOD_HEALTH_CHECK.equals(methodName) ||
                METHOD_HEALTH_WATCH.equals(methodName) ||
                //TODO: Do we want to secure the initial exchange of trust? If so, how?
                //For now, the initial exchange of trust is made unsecured. Once keys have been exchanged,
                //subsequent provision requests (if any) must be signed
                METHOD_PROVISION_NODE.equals(methodName) && !authenticationManager.isProvisioned();
    }
}
