package com.jscape.inet.ssh.protocol.v2.marshaling.algorithms;

import javax.crypto.Cipher;
import java.security.Provider;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

import static com.jscape.util.Parameters.*;

/**
 * JCE-based encryption algorithm factory.
 *
 * @author Alex Usun
 */
public class JceEncryptionFactory
        implements EncryptionFactory {

    private final Map<String, Entry> entries;

    /**
     * Creates a new factory instance.
     *
     * @param entries entries
     */
    public JceEncryptionFactory(Entry... entries) {
        this.entries = new LinkedHashMap<>();
        set(entries);
    }

    /**
     * Sets algorithm entries.
     *
     * @param entries entries
     *
     * @return current factory instance
     */
    public JceEncryptionFactory set(Entry... entries) {
        for (Entry entry : entries) {
            this.entries.put(entry.name, entry);
        }

        return this;
    }

    @Override
    public List<String> encryptions() {
        return new ArrayList<>(this.entries.keySet());
    }

    @Override
    public Encryption clientToServerEncryptionFor(String name, KeyFactory keyFactory, Mode mode)
            throws OperationException {
        Entry entry = entryFor(name);
        return entry.clientToServer(keyFactory, mode);
    }

    @Override
    public Encryption serverToClientEncryptionFor(String name, KeyFactory keyFactory, Mode mode)
            throws OperationException {
        Entry entry = entryFor(name);
        return entry.serverToClient(keyFactory, mode);
    }

    @Override
    public String toString() {
        return "JceEncryptionFactory {"
                + "entries=" + this.entries
                + '}';
    }

    private Entry entryFor(String name) {
        Entry entry = this.entries.get(name);

        if (entry == null) {
            throw new IllegalArgumentException("Unsupported algorithm: " + name);
        }

        return entry;
    }

    /**
     * An algorithm entry.
     */
    public static class Entry {

        public static final JceEncryptionFactory.Entry NONE = new JceEncryptionFactory.Entry("none", "", null, 8, 1) {

            @Override
            public Encryption clientToServer(KeyFactory keyFactory, EncryptionFactory.Mode mode) {
                return Encryption.NULL;
            }

            @Override
            public Encryption serverToClient(KeyFactory keyFactory, EncryptionFactory.Mode mode) {
                return Encryption.NULL;
            }
        };

        public final String name;
        public final String algorithm;
        public final Object provider;
        public final int blockLength;
        public final int keyLength;

        public static Entry entryFor(String name, String algorithm, String provider, int blockLength, int keyLength) {
            return new Entry(name, algorithm, provider, blockLength, keyLength);
        }

        public static Entry entryFor(String name, String algorithm, Provider provider, int blockLength, int keyLength) {
            return new Entry(name, algorithm, provider, blockLength, keyLength);
        }

        /**
         * Creates a new algorithm entry.
         *
         * @param name a SSH algorithm name
         * @param algorithm a JCE algorithm name
         * @param provider a JCE provider name
         * @param blockLength the cipher block length in bytes
         * @param keyLength the key length in bytes
         */
        protected Entry(String name, String algorithm, Object provider, int blockLength, int keyLength) {
            assertNotNull(name);
            this.name = name;

            assertNotNull(algorithm);
            this.algorithm = algorithm;

            this.provider = provider;

            assertGreaterOrEqual(blockLength, 0, "Bad block length.");
            this.blockLength = blockLength;

            assertGreater(keyLength, 0, "Bad key length.");
            this.keyLength = keyLength;
        }

        public Encryption clientToServer(KeyFactory keyFactory, Mode mode)
                throws OperationException {
            try {
                byte[] ivData = keyFactory.ivClientToServer(this.blockLength);
                byte[] keyData = keyFactory.encryptionKeyClientToServer(this.keyLength);

                int cipherMode = mode == Mode.ENCRYPTION ? Cipher.ENCRYPT_MODE : Cipher.DECRYPT_MODE;

                Encryption encryption = encryptionFor(keyData, ivData, cipherMode);
                adjust(encryption);

                return encryption;
            } catch (Exception e) {
                throw new OperationException(e);
            }
        }

        public Encryption serverToClient(KeyFactory keyFactory, Mode mode)
                throws OperationException {
            try {
                byte[] ivData = keyFactory.ivServerToClient(this.blockLength);
                byte[] keyData = keyFactory.encryptionKeyServerToClient(this.keyLength);

                int cipherMode = mode == Mode.ENCRYPTION ? Cipher.ENCRYPT_MODE : Cipher.DECRYPT_MODE;
                Encryption encryption = encryptionFor(keyData, ivData, cipherMode);
                adjust(encryption);

                return encryption;
            } catch (Exception e) {
                throw new OperationException(e);
            }
        }

        @Override
        public String toString() {
            return "Entry {"
                    + "name='" + this.name + '\''
                    + ", algorithm='" + this.algorithm + '\''
                    + ", provider='" + this.provider + '\''
                    + ", blockLength=" + this.blockLength
                    + ", keyLength=" + this.keyLength
                    + '}';
        }

        private Encryption encryptionFor(byte[] keyData, byte[] ivData, int cipherMode)
                throws Exception {
            return this.provider instanceof String
                   ? JceEncryption.encryptionFor(this.algorithm, (String) this.provider, keyData, ivData, cipherMode)
                   : this.provider instanceof Provider ? JceEncryption.encryptionFor(this.algorithm, (Provider) this.provider, keyData, ivData, cipherMode)
                                                       : JceEncryption.encryptionFor(this.algorithm, (String) null, keyData, ivData, cipherMode);
        }

        private void adjust(Encryption encryption)
                throws Exception {
            if ("arcfour128".equals(this.name) || "arcfour256".equals(this.name)) {
                adjustRC4(encryption);
            }
        }

        private void adjustRC4(Encryption encryption)
                throws Exception {
            byte[] buffer = new byte[1536];
            encryption.apply(buffer, 0, buffer.length, buffer, 0);
        }
    }
}
