package com.atlassian.bitbucket.internal.mirroring.mirror.auth.ssh;

import com.atlassian.bitbucket.AuthorisationException;
import com.atlassian.bitbucket.auth.AuthenticationContext;
import com.atlassian.bitbucket.auth.SshAuthenticationHandler;
import com.atlassian.bitbucket.i18n.I18nService;
import com.atlassian.bitbucket.internal.mirroring.mirror.InternalUpstreamService;
import com.atlassian.bitbucket.internal.mirroring.mirror.client.CloudAuthenticationResult;
import com.atlassian.bitbucket.internal.mirroring.mirror.client.InternalUpstreamClientFactory;
import com.atlassian.bitbucket.mirroring.mirror.IntegrationState;
import com.atlassian.bitbucket.mirroring.mirror.UpstreamServer;
import com.atlassian.bitbucket.permission.Permission;
import com.atlassian.bitbucket.repository.Repository;
import com.atlassian.bitbucket.scm.ssh.ExitCodeCallback;
import com.atlassian.bitbucket.scm.ssh.SshScmRequest;
import com.atlassian.bitbucket.scm.ssh.SshScmRequestHandler;
import com.atlassian.bitbucket.scm.ssh.SshScmRequestHandlerModuleDescriptor;
import com.atlassian.bitbucket.user.ApplicationUser;
import com.atlassian.bitbucket.user.SecurityService;
import com.atlassian.plugin.PluginAccessor;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.Optional;
import java.util.stream.Stream;

import static com.atlassian.bitbucket.internal.mirroring.mirror.MirrorDescriptionUtils.describe;
import static com.atlassian.bitbucket.mirroring.mirror.UpstreamServerType.BITBUCKET_CLOUD;
import static com.atlassian.bitbucket.permission.Permission.REPO_READ;
import static com.atlassian.bitbucket.util.ModuleDescriptorUtils.toSortedModules;
import static java.util.Objects.requireNonNull;
import static java.util.Optional.empty;

/**
 * This request handler is registered with a very low weight, to ensure it always gets called before all other SCM
 * request handlers.
 * It will intercept the SCM request, parse the repository URL from the command, and attempt to delegate the
 * authentication to Bitbucket Cloud.
 * On successful authentication it will pick the first viable {@link SshScmRequestHandler request handler} and delegate
 * to that for the actual command handling.
 *
 * The reason for doing this is that in order to delegate authentication to Bitbucket Cloud, we need to know the
 * repository the user is requesting to access. This cannot (at the time of writing) be done in
 * {@link SshAuthenticationHandler ssh auth handlers} as the requested command is not available during the SSH server's
 * execution of the ssh-userauth command.
 */
public class CloudAuthenticatingSshScmRequestHandler implements SshScmRequestHandler {

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

    private final AuthenticationContext authenticationContext;
    private final InternalUpstreamClientFactory clientFactory;
    private final I18nService i18nService;
    private final InternalUpstreamService upstreamService;
    private final PluginAccessor pluginAccessor;
    private final SecurityService securityService;

    CloudAuthenticatingSshScmRequestHandler(AuthenticationContext authenticationContext,
                                            InternalUpstreamClientFactory clientFactory, I18nService i18nService,
                                            InternalUpstreamService upstreamService, PluginAccessor pluginAccessor,
                                            SecurityService securityService) {
        this.authenticationContext = authenticationContext;
        this.clientFactory = clientFactory;
        this.i18nService = i18nService;
        this.upstreamService = upstreamService;
        this.pluginAccessor = pluginAccessor;
        this.securityService = securityService;
    }

    @Override
    public boolean supports(@Nonnull String command) {
        requireNonNull(command, "command");

        UpstreamServer upstreamServer = upstreamService.get();

        if (!isSupported(upstreamServer)) {
            return false;
        }

        // This class ends up delegating all scm requests to real handlers, so only pretend to support commands we can
        // delegate later.
        return getHandlers()
                .filter(CloudAuthenticatingSshScmRequestHandler::notThisClass)
                .anyMatch(handler -> handler.supports(command));
    }

    @Nonnull
    @Override
    public Optional<SshScmRequest> create(@Nonnull String command, @Nonnull InputStream in,
                                          @Nonnull OutputStream out, @Nonnull OutputStream error,
                                          @Nonnull ExitCodeCallback exitCodeCallback) {
        requireNonNull(command, "command");
        requireNonNull(in, "in");
        requireNonNull(out, "out");
        requireNonNull(error, "error");
        requireNonNull(exitCodeCallback, "exitCodeCallback");

        if (!upstreamService.isMirror()) {
            // this shouldn't happen - DelegatedSshAuthenticationHandler should not let us get here
            log.error("Delegated SSH authentication attempted for command {} but this is not a mirror", command);
            return empty();
        }

        UpstreamServer upstreamServer = upstreamService.get();
        if (upstreamServer == null) {
            // this shouldn't happen - DelegatedSshAuthenticationHandler should not let us get here
            log.error("Delegated SSH authentication attempted for command {} but there is no upstream server registered", command);
            return empty();
        }

        String upstreamServerDescription = describe(upstreamServer);
        if (!isSupported(upstreamServer)) {
            log.debug("{}: not proceeding with SSH auth delegation. Upstream server is not supported", upstreamServerDescription);
            return empty();
        }

        ApplicationUser user = authenticationContext.getCurrentUser();
        if (!(user instanceof PublicKeyUser)) {
            log.debug("{}: not proceeding with SSH auth delegation. Unexpected user type. (Expected: {}, but was: {})",
                    upstreamServerDescription, PublicKeyUser.class, user == null ? null : user.getClass());
            return empty();
        }

        log.trace("{}: creating delegate SCM request for command {}", upstreamServerDescription, command);
        Optional<SshScmRequest> delegateRequest =
                securityService.withPermission(REPO_READ, "Create SshScmRequest for delegated auth").call(
                        () -> getHandler(command).flatMap(
                                handler -> handler.create(command, in, out, error, exitCodeCallback)));

        if (delegateRequest.isPresent()) {
            Repository repository = securityService.withPermission(REPO_READ, "Repo info for delegated auth")
                    .call(() -> delegateRequest.get().getRepository());
            String publicKey = PublicKeyUser.class.cast(user).getPublicKey();
            log.trace("{}: performing actual delegated SSH authentication for {}", upstreamServerDescription, repository);
            CloudAuthenticationResult result = clientFactory.createForCloud(upstreamServer)
                    .authenticate(repository.getSlug(), publicKey)
                    .claim();

            if (result == CloudAuthenticationResult.OK) {
                log.trace("{}: SSH Delegated authentication successful. Granting access to {}.", repository);
                securityService.withPermission(REPO_READ, repository, "Successful auth delegation")
                        .withPermission(Permission.LICENSED_USER)
                        .applyToRequest();
            } else if (result == CloudAuthenticationResult.UNAUTHORIZED) {
                try {
                    error.write("Permission denied (publickey).".getBytes()); // Standard error, so don't i18n this
                    error.write('\n');
                    error.flush();
                } catch (IOException e) {
                   log.error("{}: Failed to write authorization failure to client.", upstreamServerDescription, e);
                }
                return Optional.of(new NoOpSshScmRequest(exitCodeCallback, repository));
            } else {
                // SSH key not recognized or could not perform authentication.
                // Don't grant access and throw an AuthorisationException.
                log.debug("{}: delegated authentication failed for public key {}", upstreamServerDescription, publicKey);
                throw new AuthorisationException(i18nService
                        .createKeyedMessage("bitbucket.mirroring.auth.delegation.not.permitted"));
            }
        } else {
            log.warn("{}: command not recognized. Denying execution and not perfoming delegated SSH authentication. (Command: {})", command);
        }

        return delegateRequest;
    }

    private Optional<SshScmRequestHandler> getHandler(String command) {
        return getHandlers()
                .filter(CloudAuthenticatingSshScmRequestHandler::notThisClass)
                .filter(handler -> handler.supports(command))
                .findFirst();
    }

    private Stream<SshScmRequestHandler> getHandlers() {
        return toSortedModules(pluginAccessor.getEnabledModuleDescriptorsByClass(SshScmRequestHandlerModuleDescriptor.class));
    }

    private static boolean isSupported(UpstreamServer upstreamServer) {
        return upstreamServer != null &&
                upstreamServer.getType() == BITBUCKET_CLOUD &&
                upstreamServer.getState() == IntegrationState.INSTALLED;
    }

    private static boolean notThisClass(Object obj) {
        return !CloudAuthenticatingSshScmRequestHandler.class.equals(obj.getClass());
    }
}
