/*
 * Class: KeyFactory
 *
 * Author: Alex Usun
 *
 * Date: 31/8/2005
 */
package com.jscape.inet.ssh.util.keyreader.openssh;

import com.jscape.crypto.Hash;
import com.jscape.crypto.MD5;
import com.jscape.inet.ssh.util.keyreader.IKeyFactory;
import com.jscape.util.ArrayTools;
import com.jscape.util.Assert;

import javax.crypto.spec.SecretKeySpec;
import java.nio.charset.StandardCharsets;
import java.security.Key;

/**
 * OpenSSH format key factory.
 *
 * @author Alex Usun
 * @version 1.0
 * @see <a href="KeyFactory.java.html">Source code</a>
 */
public class KeyFactory
        implements IKeyFactory {

    private final byte[] passphrase;
    private final byte[] iv;

    /**
     * Creates a new factory instance.
     *
     * @param passphrase a passphrase
     * @param iv         a IV
     */
    public KeyFactory(String passphrase, byte[] iv) {
        Assert.notNull(passphrase);
        Assert.notNull(iv);

        this.passphrase = passphrase.getBytes(StandardCharsets.UTF_8);
        this.iv = iv.clone();
    }

    /**
     * Returns current passphrase.
     *
     * @return current passphrase
     */
    public String getPassphrase() {
        return new String(this.passphrase, StandardCharsets.UTF_8);
    }

    /**
     * Returns current IV.
     *
     * @return current IV
     */
    public byte[] getIv() {
        return (byte[]) this.iv.clone();
    }

    /**
     * @see IKeyFactory#createKey(String, int)
     */
    public Key createKey(String algorithm, int keySize) {
        Assert.isValidString(algorithm);
        Assert.greater(keySize, 0, "bad key size.");

        if ("AES".equalsIgnoreCase(algorithm)) {
            Hash hash = MD5.defaultHash();
            hash.update(this.passphrase);
            hash.update(this.iv, 0, 8);
            byte[] hashData = hash.compute();
            byte[] keyData = ArrayTools.copy(hashData, 0, keySize);

            return new SecretKeySpec(keyData, algorithm);
        } else {
            Hash hash = MD5.defaultHash();
            int digestLength = hash.getLength();
            byte[] keyData = new byte[keyDataSizeFor(keySize, digestLength)];

            for (int index = 0; index <= keyData.length - digestLength; index += digestLength) {
                hash.update(this.passphrase);
                hash.update(this.iv);
                byte[] tmp = hash.compute();
                System.arraycopy(tmp, 0, keyData, index, tmp.length);
                hash.update(tmp, 0, tmp.length);
            }

            return new SecretKeySpec(keyData, 0, keySize, algorithm);
        }
    }

    /**
     * @see IKeyFactory#createIv(int)
     */
    public byte[] createIv(int size) {
        Assert.greater(size, 0, "bad key size.");

        byte[] ivData = new byte[size];
        int dataLength = Math.min(ivData.length, this.iv.length);
        System.arraycopy(this.iv, 0, ivData, 0, dataLength);

        return ivData;
    }

    /**
     * Computes needed hash size.
     *
     * @return hash size
     */
    private int keyDataSizeFor(int keySize, int digestSize) {
        int hashSize = keySize > digestSize ? keySize : digestSize;
        int remainder = hashSize % digestSize;

        if (remainder != 0) {
            hashSize += digestSize - remainder;
        }

        return hashSize;
    }

}
