package be.iminds.ilabt.jfed.util;

import be.iminds.ilabt.jfed.lowlevel.connection.JFedConnection;
import com.hierynomus.sshj.signature.SignatureEdDSA;
import net.schmizz.sshj.DefaultConfig;
import net.schmizz.sshj.SSHClient;
import net.schmizz.sshj.common.SSHPacket;
import net.schmizz.sshj.connection.Connection;
import net.schmizz.sshj.connection.channel.direct.AbstractDirectChannel;
import net.schmizz.sshj.connection.channel.direct.LocalPortForwarder;
import net.schmizz.sshj.signature.SignatureDSA;
import net.schmizz.sshj.signature.SignatureRSA;
import net.schmizz.sshj.transport.verification.PromiscuousVerifier;
import net.schmizz.sshj.userauth.UserAuthException;
import net.schmizz.sshj.userauth.keyprovider.KeyPairWrapper;
import org.slf4j.LoggerFactory;
import org.xbill.DNS.Address;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import javax.net.SocketFactory;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.*;

public class SshProxySocketImplFactory implements SocketImplFactory {
    private static final org.slf4j.Logger LOG = LoggerFactory.getLogger(SshProxySocketImplFactory.class);

    protected final SSHClient sshClient;

    private final static boolean singleConnectionPerSShConnection = true;//TCP connection setup
    protected static final int CONNECT_TIMEOUT_MS = 1000;//full SSH connection setup
    protected static final int KEX_TIMEOUT_MS = 8000;

    protected final String forcedCreateSocketHostname;
    protected final Integer forcedCreateSocketPort;

    public synchronized static SshProxySocketImplFactory getFactory(JFedConnection.SshProxyInfo sshProxyInfo,
                                                                    String forcedCreateSocketHostname,
                                                                    Integer forcedCreateSocketPort) throws IOException {
        //see also SSHSocketImpl.close. since singleConnectionPerSShConnection == true, we can close SSH connection there!
        if (singleConnectionPerSShConnection)
            return new SshProxySocketImplFactory(sshProxyInfo, null, forcedCreateSocketHostname, forcedCreateSocketPort);
        else
            throw new RuntimeException("SSH connection reuse yet implemented");
    }

    public synchronized static SshProxySocketImplFactory getFactory(JFedConnection.SshProxyInfo sshProxyInfo) throws IOException {
        //see also SSHSocketImpl.close. since singleConnectionPerSShConnection == true, we can close SSH connection there!
        if (singleConnectionPerSShConnection)
            return new SshProxySocketImplFactory(sshProxyInfo, null, null/*forcedCreateSocketHostname*/, null/*forcedCreateSocketPort*/);
        else
            throw new RuntimeException("SSH connection reuse yet implemented");
    }

    public synchronized static SshProxySocketImplFactory getFactory(JFedConnection.SshProxyInfo sshProxyInfo, @Nullable
    SocketFactory customSocketFactory) throws IOException {
        //see also SSHSocketImpl.close. since singleConnectionPerSShConnection == true, we can close SSH connection there!
        if (singleConnectionPerSShConnection)
            return new SshProxySocketImplFactory(sshProxyInfo, customSocketFactory, null/*forcedCreateSocketHostname*/, null/*forcedCreateSocketPort*/);
        else
            throw new RuntimeException("SSH connection reuse yet implemented");
    }

    public synchronized static SshProxySocketImplFactory getFactory(JFedConnection.SshProxyInfo sshProxyInfo,
                                                                    @Nullable SocketFactory customSocketFactory,
                                                                    String forcedCreateSocketHostname,
                                                                    int forcedCreateSocketPort) throws IOException {
        //see also SSHSocketImpl.close. since singleConnectionPerSShConnection == true, we can close SSH connection there!
        if (singleConnectionPerSShConnection)
            return new SshProxySocketImplFactory(sshProxyInfo, customSocketFactory, forcedCreateSocketHostname, forcedCreateSocketPort);
        else
            throw new RuntimeException("SSH connection reuse yet implemented");
    }

    public static class SocketWithCustomSocketFactory extends Socket {
        public SocketWithCustomSocketFactory(SocketImplFactory socketImplFactory, String host, int port) throws IOException {
            super(socketImplFactory.createSocketImpl());
            connect(host != null ?
                            new InetSocketAddress(host, port) :
                            new InetSocketAddress(InetAddress.getByName(null), port)
            );
        }

        public SocketWithCustomSocketFactory(SocketImplFactory socketImplFactory) throws SocketException {
            super(socketImplFactory.createSocketImpl());
        }

        public SocketWithCustomSocketFactory(SocketImplFactory socketImplFactory, InetAddress host, int port) throws IOException {
            super(socketImplFactory.createSocketImpl());
            connect(new InetSocketAddress(host, port));
        }
    }

    public static class SocketFactoryWithCustomSocketImplFactory extends SocketFactory {

        private final SocketImplFactory socketImplFactory;

        public SocketFactoryWithCustomSocketImplFactory(SocketImplFactory socketImplFactory) {
            this.socketImplFactory = socketImplFactory;
        }


        @Override
        public Socket createSocket() throws IOException {
            return new SocketWithCustomSocketFactory(socketImplFactory);
        }

        @Override
        public Socket createSocket(String host, int port) throws IOException {
            return new SocketWithCustomSocketFactory(socketImplFactory, host, port);
        }

        @Override
        public Socket createSocket(String host, int port, InetAddress localAddress, int localPort) throws IOException {
            throw new UnsupportedOperationException("No support for binding sockets to a local address and port");
        }

        @Override
        public Socket createSocket(InetAddress host, int port) throws IOException {
            return new SocketWithCustomSocketFactory(socketImplFactory, host, port);
        }

        @Override
        public Socket createSocket(InetAddress host, int port, InetAddress localAddress, int localPort) throws IOException {
            throw new UnsupportedOperationException("No support for binding sockets to a local address and port");
        }
    }


    private SshProxySocketImplFactory(@Nonnull final JFedConnection.SshProxyInfo sshProxyInfo,
                                      @Nullable final SocketFactory customSocketFactory,
                                      @Nullable final String forcedCreateSocketHostname,
                                      @Nullable final Integer forcedCreateSocketPort) throws IOException {
        this.forcedCreateSocketHostname = forcedCreateSocketHostname;
        this.forcedCreateSocketPort = forcedCreateSocketPort;

        //remove support for ECDSA signatures
        DefaultConfig config = new DefaultConfig() {
            @Override
            protected void initSignatureFactories() {
                setSignatureFactories(new SignatureRSA.Factory(), new SignatureDSA.Factory(), new SignatureEdDSA.Factory());

            }
        };
        sshClient = new SSHClient(config);

        if (customSocketFactory != null)
            sshClient.setSocketFactory(customSocketFactory);

        if (sshProxyInfo.getHostKey() != null) {
            sshClient.addHostKeyVerifier(new OpenSSHKnownHostsInMemory(sshProxyInfo.getHostKey()));
        } else {
            LOG.warn("Cannot verify proxy host as no known_hosts-line is specified");
            sshClient.addHostKeyVerifier(new PromiscuousVerifier());
        }


        //timeouts in milliseconds. Zero means no timeout.

        sshClient.setTimeout(KEX_TIMEOUT_MS);
        sshClient.setConnectTimeout(CONNECT_TIMEOUT_MS);

        LOG.debug("Connecting to " + sshProxyInfo.getHostname() + ":" + sshProxyInfo.getPort() + " as user "+sshProxyInfo.getUsername()+" with timeout " + CONNECT_TIMEOUT_MS+" ms");
        sshClient.connect(sshProxyInfo.getHostname(), sshProxyInfo.getPort());

        LOG.debug("Connect done. Will now authenticate.");

        try {
            assert sshProxyInfo.getUsername() != null;
            assert sshProxyInfo.getSshKeyInfo() != null;
            assert sshProxyInfo.getSshKeyInfo().getPrivateKey() != null;
            assert KeyUtil.matchingKeys(sshProxyInfo.getSshKeyInfo().getPublicKey(), sshProxyInfo.getSshKeyInfo().getPrivateKey());

            sshClient.authPublickey(sshProxyInfo.getUsername(),
                    new KeyPairWrapper(sshProxyInfo.getSshKeyInfo().getPublicKey(), sshProxyInfo.getSshKeyInfo().getPrivateKey()));
            LOG.debug("Connected and authenticated to {}:{}", sshProxyInfo.getHostname(), sshProxyInfo.getPort());

        } catch (UserAuthException ex) {
            LOG.warn("Connected to {}:{}, but authentication failed", sshProxyInfo.getHostname(), sshProxyInfo.getPort());
            throw new IOException("Failed to authenticate user \"" + sshProxyInfo.getUsername() + "\"" +
                    " on server " + sshProxyInfo.getHostname() + ":" + sshProxyInfo.getPort() + "", ex);
        }
    }

    //note: multiple SocketImpl can be needed per factory. example: When using ganymed itself using this factory, it needs multiple
    public SocketImpl createSocketImpl() {
        if (sshClient == null)
            throw new RuntimeException("The SSH connection is not ready. Cannot create socket over it.");
        SSHSocketImpl res = new SSHSocketImpl(sshClient, forcedCreateSocketHostname, forcedCreateSocketPort);
        if (forcedCreateSocketHostname != null && forcedCreateSocketPort != null)
            LOG.debug("returning new SSHSocketImpl with forced destination "+forcedCreateSocketHostname+":"+forcedCreateSocketPort);
        else
            LOG.debug("returning new SSHSocketImpl");
        return res;

    }

    public static class DirectTCPIPChannel extends AbstractDirectChannel {
        protected final LocalPortForwarder.Parameters parameters;

        public DirectTCPIPChannel(Connection conn, LocalPortForwarder.Parameters parameters) {
            super(conn, "direct-tcpip");
            this.parameters = parameters;
        }

        @Override
        protected SSHPacket buildOpenReq() {
            return super.buildOpenReq()
                    .putString(parameters.getRemoteHost())
                    .putUInt32(parameters.getRemotePort())
                    .putString(parameters.getLocalHost())
                    .putUInt32(parameters.getLocalPort());
        }

    }


    /**
     * We aren't going to actually create any new connection, we will forward
     * things to SSH.
     */
    private static class SSHSocketImpl extends SocketImpl {
        private SSHClient sshConnection;
        private DirectTCPIPChannel channel;

        protected final String forcedCreateSocketHostname;
        protected final Integer forcedCreateSocketPort;

        InputStream is;

        OutputStream os;

        SSHSocketImpl(SSHClient sshConnection, String forcedCreateSocketHostname, Integer forcedCreateSocketPort) {
            this.sshConnection = sshConnection;
            this.forcedCreateSocketHostname = forcedCreateSocketHostname;
            this.forcedCreateSocketPort = forcedCreateSocketPort;
        }

        @Override
        protected void accept(SocketImpl s) throws IOException {
            throw new IOException("SSHSocketImpl does not implement accept");
        }

        @Override
        protected int available() throws IOException {
            if (is == null) {
                throw new ConnectException("Not connected");
            }
            return is.available();
        }

        @Override
        protected void bind(InetAddress host, int port) throws IOException {
            if ((host != null && !host.isAnyLocalAddress()) || port != 0) {
                throw new IOException("SSHSocketImpl does not implement bind");
            }
        }

        @Override
        protected void close() throws IOException {
            LOG.debug("Closing SSH Proxy socket");
            if (channel != null) {
                channel.close();
                channel = null;
                is = null;
                os = null;
            }

            //also close SSH connection:
            if (singleConnectionPerSShConnection) {
                sshConnection.close();
                sshConnection = null;
            } else
                throw new RuntimeException("connection close not yet implemented for SSH connection reuse");
        }

        @Override
        protected void connect(String host, int port) throws IOException {
            if (forcedCreateSocketHostname != null && forcedCreateSocketPort != null) {
                connectToForced(300000);
                return;
            }

            InetAddress addr = Address.getByName(host);
            connect(addr, port);
        }

        @Override
        protected void connect(InetAddress address, int port) throws IOException {
            if (forcedCreateSocketHostname != null && forcedCreateSocketPort != null) {
                connectToForced(300000);
                return;
            }
            connect(new InetSocketAddress(address, port), 300000);
        }

        @Override
        protected void connect(SocketAddress address, int timeout)
                throws IOException {
            if (forcedCreateSocketHostname != null && forcedCreateSocketPort != null) {
                connectToForced(timeout);
                return;
            }

            LOG.debug("Creating channel to " + ((InetSocketAddress) address).getHostName() + ":" + ((InetSocketAddress) address).getPort());

            channel = new DirectTCPIPChannel(sshConnection.getConnection(),
                    new LocalPortForwarder.Parameters("", 0, ((InetSocketAddress) address).getHostName(), ((InetSocketAddress) address).getPort()));

            try {
                channel.open();
            } catch (IOException e) {
                net.schmizz.sshj.common.IOUtils.closeQuietly(channel);
                throw e;
            }

            is = channel.getInputStream();
            os = channel.getOutputStream();
        }

        protected void connectToForced(int timeout)
                throws IOException {
            assert forcedCreateSocketHostname != null;
            assert forcedCreateSocketPort != null;

            LOG.debug("Creating channel to forced " + forcedCreateSocketHostname + ":" + forcedCreateSocketPort);

            channel = new DirectTCPIPChannel(sshConnection.getConnection(),
                    new LocalPortForwarder.Parameters("", 0, forcedCreateSocketHostname, forcedCreateSocketPort));

            try {
                channel.open();
            } catch (IOException e) {
                net.schmizz.sshj.common.IOUtils.closeQuietly(channel);
                throw e;
            }

            is = channel.getInputStream();
            os = channel.getOutputStream();
        }

        @Override
        protected void create(boolean stream) throws IOException {
            if (!stream) {
                throw new IOException("Cannot handle UDP streams");
            }
        }

        @Override
        protected InputStream getInputStream() throws IOException {
            return is;
        }

        @Override
        protected OutputStream getOutputStream() throws IOException {
            return os;
        }

        @Override
        protected void listen(int backlog) throws IOException {
            throw new IOException("SSHSocketImpl does not implement listen");
        }

        @Override
        protected void sendUrgentData(int data) throws IOException {
            throw new IOException("SSHSocketImpl does not implement sendUrgentData");
        }

        public Object getOption(int optID) throws SocketException {
            //fake answer to some common option requests if possible.
            if (optID == SocketOptions.TCP_NODELAY) return false;
            if (optID == SocketOptions.SO_BINDADDR)
                throw new SocketException("SSHSocketImpl does not implement getOption for SO_BINDADDR");
            if (optID == SocketOptions.SO_REUSEADDR) return false;
            if (optID == SocketOptions.SO_BROADCAST) return false;
            if (optID == SocketOptions.IP_MULTICAST_IF)
                throw new SocketException("SSHSocketImpl does not implement getOption for IP_MULTICAST_IF");
            if (optID == SocketOptions.IP_MULTICAST_IF2)
                throw new SocketException("SSHSocketImpl does not implement getOption for IP_MULTICAST_IF2");
            if (optID == SocketOptions.IP_MULTICAST_LOOP)
                throw new SocketException("SSHSocketImpl does not implement getOption for IP_MULTICAST_LOOP");
            if (optID == SocketOptions.IP_TOS)
                throw new SocketException("SSHSocketImpl does not implement getOption for IP_TOS");
            if (optID == SocketOptions.SO_LINGER) return 0; //no linger
            if (optID == SocketOptions.SO_TIMEOUT) return 0; //?
            if (optID == SocketOptions.SO_SNDBUF)
                throw new SocketException("SSHSocketImpl does not implement getOption for SO_SNDBUF"); //actual send buffer size
            if (optID == SocketOptions.SO_RCVBUF)
                throw new SocketException("SSHSocketImpl does not implement getOption for SO_RCVBUF");
            //actual recv buffer size
            if (optID == SocketOptions.SO_KEEPALIVE) return false;
            if (optID == SocketOptions.SO_OOBINLINE) return false;

            throw new SocketException("SSHSocketImpl does not implement getOption for " + optID);
        }

        /**
         * We silently ignore setOptions because they do happen, but there is
         * nothing that we can do about it.
         */
        public void setOption(int optID, Object value) throws SocketException {
        }

    }
}
