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

import javax.crypto.Cipher;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.security.Key;
import java.security.Provider;
import java.security.spec.AlgorithmParameterSpec;

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

/**
 * @author Alex Usun
 */
public class JceEncryption
        implements Encryption {

    private final Cipher cipher;

    public static Encryption encryptionFor(String name, String provider, byte[] keyData, byte[] ivData, int mode)
            throws OperationException {
        try {
            Cipher cipher = provider != null ? Cipher.getInstance(name, provider) : Cipher.getInstance(name);
            return encryptionFor(cipher, keyData, ivData, mode);
        } catch (Exception e) {
            throw new OperationException(e);
        }
    }

    public static Encryption encryptionFor(String name, Provider provider, byte[] keyData, byte[] ivData, int mode)
            throws OperationException {
        try {
            Cipher cipher = provider != null ? Cipher.getInstance(name, provider) : Cipher.getInstance(name);
            return encryptionFor(cipher, keyData, ivData, mode);
        } catch (Exception e) {
            throw new OperationException(e);
        }
    }

    private static Encryption encryptionFor(Cipher cipher, byte[] keyData, byte[] ivData, int mode)
            throws Exception {
        Key key = keyFor(cipher.getAlgorithm(), keyData);

        if (ivData.length > 0) {
            AlgorithmParameterSpec iv = new IvParameterSpec(ivData);
            cipher.init(mode, key, iv);
        } else {
            cipher.init(mode, key);
        }

        return new JceEncryption(cipher);
    }

    private static Key keyFor(String algorithm, byte[] keyData) {
        int separatorPosition = algorithm.indexOf('/');

        if (separatorPosition > 0) {
            algorithm = algorithm.substring(0, separatorPosition);
        }

        return new SecretKeySpec(keyData, algorithm);
    }

    public JceEncryption(Cipher cipher) {
        assertNotNull(cipher);
        this.cipher = cipher;
    }

    @Override
    public int blockLength() {
        return this.cipher.getBlockSize();
    }

    @Override
    public void apply(byte[] data, int offset, int length, byte[] buffer, int bufferOffset)
            throws OperationException {
        try {
            this.cipher.update(data, offset, length, buffer, bufferOffset);
        } catch (Exception e) {
            throw new OperationException(e);
        }
    }

    @Override
    public String toString() {
        return "JceEncryption {"
                + "cipher=" + this.cipher.getAlgorithm()
                + '}';
    }
}
