/*
 * Copyright (C) 2018, 2022 Thomas Wolf <thomas.wolf@paranor.ch> and others
 *
 * This program and the accompanying materials are made available under the
 * terms of the Eclipse Distribution License v. 1.0 which is available at
 * https://www.eclipse.org/org/documents/edl-v10.php.
 *
 * SPDX-License-Identifier: BSD-3-Clause
 */
package org.eclipse.jgit.internal.transport.sshd;

import static java.text.MessageFormat.format;
import static org.apache.sshd.core.CoreModuleProperties.PASSWORD_PROMPTS;
import static org.apache.sshd.core.CoreModuleProperties.PREFERRED_AUTHS;
import static org.eclipse.jgit.internal.transport.ssh.OpenSshConfigFile.positive;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.Proxy;
import java.net.SocketAddress;
import java.nio.file.Files;
import java.nio.file.InvalidPathException;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.security.GeneralSecurityException;
import java.security.KeyPair;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.Objects;
import java.util.function.Supplier;
import java.util.stream.Collectors;

import org.apache.sshd.agent.SshAgentFactory;
import org.apache.sshd.client.SshClient;
import org.apache.sshd.client.config.hosts.HostConfigEntry;
import org.apache.sshd.client.future.ConnectFuture;
import org.apache.sshd.client.future.DefaultConnectFuture;
import org.apache.sshd.client.session.ClientSessionImpl;
import org.apache.sshd.client.session.SessionFactory;
import org.apache.sshd.common.AttributeRepository;
import org.apache.sshd.common.config.keys.FilePasswordProvider;
import org.apache.sshd.common.future.SshFutureListener;
import org.apache.sshd.common.io.IoConnectFuture;
import org.apache.sshd.common.io.IoSession;
import org.apache.sshd.common.keyprovider.AbstractResourceKeyPairProvider;
import org.apache.sshd.common.keyprovider.KeyIdentityProvider;
import org.apache.sshd.common.session.SessionContext;
import org.apache.sshd.common.session.helpers.AbstractSession;
import org.apache.sshd.common.util.ValidateUtils;
import org.apache.sshd.common.util.net.SshdSocketAddress;
import org.eclipse.jgit.internal.transport.sshd.JGitClientSession.ChainingAttributes;
import org.eclipse.jgit.internal.transport.sshd.JGitClientSession.SessionAttributes;
import org.eclipse.jgit.internal.transport.sshd.proxy.HttpClientConnector;
import org.eclipse.jgit.internal.transport.sshd.proxy.Socks5ClientConnector;
import org.eclipse.jgit.transport.CredentialsProvider;
import org.eclipse.jgit.transport.SshConstants;
import org.eclipse.jgit.transport.sshd.KeyCache;
import org.eclipse.jgit.transport.sshd.ProxyData;
import org.eclipse.jgit.transport.sshd.ProxyDataFactory;
import org.eclipse.jgit.util.StringUtils;

/**
 * Customized {@link SshClient} for JGit. It creates specialized
 * {@link JGitClientSession}s that know about the {@link HostConfigEntry} they
 * were created for, and it loads all KeyPair identities lazily.
 */
public class JGitSshClient extends SshClient {

	/**
	 * We need access to this during the constructor of the ClientSession,
	 * before setConnectAddress() can have been called. So we have to remember
	 * it in an attribute on the SshClient, from where we can then retrieve it.
	 */
	static final AttributeKey<HostConfigEntry> HOST_CONFIG_ENTRY = new AttributeKey<>();

	static final AttributeKey<InetSocketAddress> ORIGINAL_REMOTE_ADDRESS = new AttributeKey<>();

	/**
	 * An attribute key for the comma-separated list of default preferred
	 * authentication mechanisms.
	 */
	public static final AttributeKey<String> PREFERRED_AUTHENTICATIONS = new AttributeKey<>();

	/**
	 * An attribute key for the home directory.
	 */
	public static final AttributeKey<Path> HOME_DIRECTORY = new AttributeKey<>();

	/**
	 * An attribute key for storing an alternate local address to connect to if
	 * a local forward from a ProxyJump ssh config is present. If set,
	 * {@link #connect(HostConfigEntry, AttributeRepository, SocketAddress)}
	 * will not connect to the address obtained from the {@link HostConfigEntry}
	 * but to the address stored in this key (which is assumed to forward the
	 * {@code HostConfigEntry} address).
	 */
	public static final AttributeKey<SshdSocketAddress> LOCAL_FORWARD_ADDRESS = new AttributeKey<>();

	private KeyCache keyCache;

	private CredentialsProvider credentialsProvider;

	private ProxyDataFactory proxyDatabase;

	private Supplier<SshAgentFactory> agentFactorySupplier = () -> null;

	@Override
	protected SessionFactory createSessionFactory() {
		// Override the parent's default
		return new JGitSessionFactory(this);
	}

	@Override
	public ConnectFuture connect(HostConfigEntry hostConfig,
			AttributeRepository context, SocketAddress localAddress)
			throws IOException {
		if (connector == null) {
			throw new IllegalStateException("SshClient not started."); //$NON-NLS-1$
		}
		Objects.requireNonNull(hostConfig, "No host configuration"); //$NON-NLS-1$
		String originalHost = ValidateUtils.checkNotNullAndNotEmpty(
				hostConfig.getHostName(), "No target host"); //$NON-NLS-1$
		int originalPort = hostConfig.getPort();
		ValidateUtils.checkTrue(originalPort > 0, "Invalid port: %d", //$NON-NLS-1$
				originalPort);
		InetSocketAddress originalAddress = new InetSocketAddress(originalHost,
				originalPort);
		InetSocketAddress targetAddress = originalAddress;
		String userName = hostConfig.getUsername();
		String id = userName + '@' + originalAddress;
		AttributeRepository attributes = chain(context, this);
		SshdSocketAddress localForward = attributes
				.resolveAttribute(LOCAL_FORWARD_ADDRESS);
		if (localForward != null) {
			targetAddress = new InetSocketAddress(localForward.getHostName(),
					localForward.getPort());
			id += '/' + targetAddress.toString();
		}
		ConnectFuture connectFuture = new DefaultConnectFuture(id, null);
		SshFutureListener<IoConnectFuture> listener = createConnectCompletionListener(
				connectFuture, userName, originalAddress, hostConfig);
		attributes = sessionAttributes(attributes, hostConfig, originalAddress);
		// Proxy support
		if (localForward == null) {
			ProxyData proxy = getProxyData(targetAddress);
			if (proxy != null) {
				targetAddress = configureProxy(proxy, targetAddress);
				proxy.clearPassword();
			}
		}
		connector.connect(targetAddress, attributes, localAddress)
				.addListener(listener);
		return connectFuture;
	}

	private AttributeRepository chain(AttributeRepository self,
			AttributeRepository parent) {
		if (self == null) {
			return Objects.requireNonNull(parent);
		}
		if (parent == null || parent == self) {
			return self;
		}
		return new ChainingAttributes(self, parent);
	}

	private AttributeRepository sessionAttributes(AttributeRepository parent,
			HostConfigEntry hostConfig, InetSocketAddress originalAddress) {
		// sshd needs some entries from the host config already in the
		// constructor of the session. Put those into a dedicated
		// AttributeRepository for the new session where it will find them.
		// We can set the host config only once the session object has been
		// created.
		Map<AttributeKey<?>, Object> data = new HashMap<>();
		data.put(HOST_CONFIG_ENTRY, hostConfig);
		data.put(ORIGINAL_REMOTE_ADDRESS, originalAddress);
		data.put(TARGET_SERVER, new SshdSocketAddress(originalAddress));
		String preferredAuths = hostConfig.getProperty(
				SshConstants.PREFERRED_AUTHENTICATIONS,
				resolveAttribute(PREFERRED_AUTHENTICATIONS));
		if (!StringUtils.isEmptyOrNull(preferredAuths)) {
			data.put(SessionAttributes.PROPERTIES,
					Collections.singletonMap(
							PREFERRED_AUTHS.getName(),
							preferredAuths));
		}
		return new SessionAttributes(
				AttributeRepository.ofAttributesMap(data),
				parent, this);
	}

	private ProxyData getProxyData(InetSocketAddress remoteAddress) {
		ProxyDataFactory factory = getProxyDatabase();
		return factory == null ? null : factory.get(remoteAddress);
	}

	private InetSocketAddress configureProxy(ProxyData proxyData,
			InetSocketAddress remoteAddress) {
		Proxy proxy = proxyData.getProxy();
		if (proxy.type() == Proxy.Type.DIRECT
				|| !(proxy.address() instanceof InetSocketAddress)) {
			return remoteAddress;
		}
		InetSocketAddress address = (InetSocketAddress) proxy.address();
		if (address.isUnresolved()) {
			address = new InetSocketAddress(address.getHostName(),
					address.getPort());
		}
		switch (proxy.type()) {
		case HTTP:
			setClientProxyConnector(
					new HttpClientConnector(address, remoteAddress,
							proxyData.getUser(), proxyData.getPassword()));
			return address;
		case SOCKS:
			setClientProxyConnector(
					new Socks5ClientConnector(address, remoteAddress,
							proxyData.getUser(), proxyData.getPassword()));
			return address;
		default:
			log.warn(format(SshdText.get().unknownProxyProtocol,
					proxy.type().name()));
			return remoteAddress;
		}
	}

	private SshFutureListener<IoConnectFuture> createConnectCompletionListener(
			ConnectFuture connectFuture, String username,
			InetSocketAddress address, HostConfigEntry hostConfig) {
		return new SshFutureListener<>() {

			@Override
			public void operationComplete(IoConnectFuture future) {
				if (future.isCanceled()) {
					connectFuture.cancel();
					return;
				}
				Throwable t = future.getException();
				if (t != null) {
					connectFuture.setException(t);
					return;
				}
				IoSession ioSession = future.getSession();
				try {
					JGitClientSession session = createSession(ioSession,
							username, address, hostConfig);
					connectFuture.setSession(session);
				} catch (RuntimeException e) {
					connectFuture.setException(e);
					ioSession.close(true);
				}
			}

			@Override
			public String toString() {
				return "JGitSshClient$ConnectCompletionListener[" + username //$NON-NLS-1$
						+ '@' + address + ']';
			}
		};
	}

	private JGitClientSession createSession(IoSession ioSession,
			String username, InetSocketAddress address,
			HostConfigEntry hostConfig) {
		AbstractSession rawSession = AbstractSession.getSession(ioSession);
		if (!(rawSession instanceof JGitClientSession)) {
			throw new IllegalStateException("Wrong session type: " //$NON-NLS-1$
					+ rawSession.getClass().getCanonicalName());
		}
		JGitClientSession session = (JGitClientSession) rawSession;
		session.setUsername(username);
		session.setConnectAddress(address);
		session.setHostConfigEntry(hostConfig);
		if (session.getCredentialsProvider() == null) {
			session.setCredentialsProvider(getCredentialsProvider());
		}
		int numberOfPasswordPrompts = getNumberOfPasswordPrompts(hostConfig);
		PASSWORD_PROMPTS.set(session, Integer.valueOf(numberOfPasswordPrompts));
		List<Path> identities = hostConfig.getIdentities().stream()
				.map(s -> {
					try {
						return Paths.get(s);
					} catch (InvalidPathException e) {
						log.warn(format(SshdText.get().configInvalidPath,
								SshConstants.IDENTITY_FILE, s), e);
						return null;
					}
				}).filter(p -> p != null && Files.exists(p))
				.collect(Collectors.toList());
		CachingKeyPairProvider ourConfiguredKeysProvider = new CachingKeyPairProvider(
				identities, keyCache);
		FilePasswordProvider passwordProvider = getFilePasswordProvider();
		ourConfiguredKeysProvider.setPasswordFinder(passwordProvider);
		if (hostConfig.isIdentitiesOnly()) {
			session.setKeyIdentityProvider(ourConfiguredKeysProvider);
		} else {
			KeyIdentityProvider defaultKeysProvider = getKeyIdentityProvider();
			if (defaultKeysProvider instanceof AbstractResourceKeyPairProvider<?>) {
				((AbstractResourceKeyPairProvider<?>) defaultKeysProvider)
						.setPasswordFinder(passwordProvider);
			}
			KeyIdentityProvider combinedProvider = new CombinedKeyIdentityProvider(
					ourConfiguredKeysProvider, defaultKeysProvider);
			session.setKeyIdentityProvider(combinedProvider);
		}
		return session;
	}

	private int getNumberOfPasswordPrompts(HostConfigEntry hostConfig) {
		String prompts = hostConfig
				.getProperty(SshConstants.NUMBER_OF_PASSWORD_PROMPTS);
		if (prompts != null) {
			prompts = prompts.trim();
			int value = positive(prompts);
			if (value > 0) {
				return value;
			}
			log.warn(format(SshdText.get().configInvalidPositive,
					SshConstants.NUMBER_OF_PASSWORD_PROMPTS, prompts));
		}
		return PASSWORD_PROMPTS.getRequiredDefault().intValue();
	}

	/**
	 * Set a cache for loaded keys. Newly discovered keys will be added when
	 * IdentityFile host entries from the ssh config file are used during
	 * session authentication.
	 *
	 * @param cache
	 *            to use
	 */
	public void setKeyCache(KeyCache cache) {
		keyCache = cache;
	}

	/**
	 * Sets a {@link ProxyDataFactory} for connecting through proxies.
	 *
	 * @param factory
	 *            to use, or {@code null} if proxying is not desired or
	 *            supported
	 */
	public void setProxyDatabase(ProxyDataFactory factory) {
		proxyDatabase = factory;
	}

	/**
	 * Retrieves the {@link ProxyDataFactory}.
	 *
	 * @return the factory, or {@code null} if none is set
	 */
	protected ProxyDataFactory getProxyDatabase() {
		return proxyDatabase;
	}

	/**
	 * Sets the {@link CredentialsProvider} for this client.
	 *
	 * @param provider
	 *            to set
	 */
	public void setCredentialsProvider(CredentialsProvider provider) {
		credentialsProvider = provider;
	}

	/**
	 * Retrieves the {@link CredentialsProvider} set for this client.
	 *
	 * @return the provider, or {@code null} if none is set.
	 */
	public CredentialsProvider getCredentialsProvider() {
		return credentialsProvider;
	}

	@Override
	public SshAgentFactory getAgentFactory() {
		return agentFactorySupplier.get();
	}

	@Override
	protected void checkConfig() {
		// The super class requires channel factories for agent forwarding if a
		// factory for an SSH agent is set. We haven't implemented this yet, and
		// we don't do SSH agent forwarding for now. Unfortunately, there is no
		// way to bypass this check in the super class except making
		// getAgentFactory() return null until after the check.
		super.checkConfig();
		agentFactorySupplier = super::getAgentFactory;
	}

	/**
	 * A {@link SessionFactory} to create our own specialized
	 * {@link JGitClientSession}s.
	 */
	private static class JGitSessionFactory extends SessionFactory {

		public JGitSessionFactory(JGitSshClient client) {
			super(client);
		}

		@Override
		protected ClientSessionImpl doCreateSession(IoSession ioSession)
				throws Exception {
			return new JGitClientSession(getClient(), ioSession);
		}
	}

	/**
	 * A {@link KeyIdentityProvider} that iterates over the {@link Iterable}s
	 * returned by other {@link KeyIdentityProvider}s.
	 */
	private static class CombinedKeyIdentityProvider
			implements KeyIdentityProvider {

		private final List<KeyIdentityProvider> providers;

		public CombinedKeyIdentityProvider(KeyIdentityProvider... providers) {
			this(Arrays.stream(providers).filter(Objects::nonNull)
					.collect(Collectors.toList()));
		}

		public CombinedKeyIdentityProvider(
				List<KeyIdentityProvider> providers) {
			this.providers = providers;
		}

		@Override
		public Iterable<KeyPair> loadKeys(SessionContext context) {
			return () -> new Iterator<>() {

				private Iterator<KeyIdentityProvider> factories = providers
						.iterator();
				private Iterator<KeyPair> current;

				private Boolean hasElement;

				@Override
				public boolean hasNext() {
					if (hasElement != null) {
						return hasElement.booleanValue();
					}
					while (current == null || !current.hasNext()) {
						if (factories.hasNext()) {
							try {
								current = factories.next().loadKeys(context)
										.iterator();
							} catch (IOException | GeneralSecurityException e) {
								throw new RuntimeException(e);
							}
						} else {
							current = null;
							hasElement = Boolean.FALSE;
							return false;
						}
					}
					hasElement = Boolean.TRUE;
					return true;
				}

				@Override
				public KeyPair next() {
					if ((hasElement == null && !hasNext())
							|| !hasElement.booleanValue()) {
						throw new NoSuchElementException();
					}
					hasElement = null;
					KeyPair result;
					try {
						result = current.next();
					} catch (NoSuchElementException e) {
						result = null;
					}
					return result;
				}

			};
		}
	}
}