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

import com.jscape.inet.ssh.util.keyreader.FormatException;
import com.jscape.inet.ssh.util.keyreader.KeyPairFormat;
import com.jscape.util.Assert;
import com.jscape.util.ExceptionTools;
import com.jscape.util.crypto.CryptoServicesProvider;

import java.io.*;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.security.KeyPair;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;

/**
 * Key pair reader. Provides reading services for several key formats.
 * <p/>
 * All methods throw <code>NullPointerException</code> if null-references were passed in arguments.
 *
 * @author Alex Usun
 */
@SuppressWarnings("unchecked")
public class KeyPairAssembler {

    private static final int BUFFER_SIZE = 128;

    private final List formats = new ArrayList();

    /**
     * Creates a new assembler instance.
     */
    public KeyPairAssembler() {
        addFormat(new com.jscape.inet.ssh.util.keyreader.openssh.RSAFormat());
        addFormat(new com.jscape.inet.ssh.util.keyreader.openssh.DSAFormat());
        addFormat(new com.jscape.inet.ssh.util.keyreader.openssh.OpenSshNewFormat());
        addFormat(new com.jscape.inet.ssh.util.keyreader.putty.RSAFormat());
        addFormat(new com.jscape.inet.ssh.util.keyreader.putty.DSAFormat());
        addFormat(new com.jscape.inet.ssh.util.keyreader.sshcom.RSAFormat());
        addFormat(new com.jscape.inet.ssh.util.keyreader.sshcom.DSAFormat());
        addFormat(new com.jscape.inet.ssh.util.keyreader.vandyke.RSAFormat());
        addFormat(new com.jscape.inet.ssh.util.keyreader.vandyke.DSAFormat());
        addFormat(new com.jscape.inet.ssh.util.keyreader.Pkcs8KeyPairFormat());
        addFormat(new com.jscape.inet.ssh.util.keyreader.KeyStoreKeyPairFormat("JKS", null));
        addFormat(new com.jscape.inet.ssh.util.keyreader.KeyStoreKeyPairFormat("PKCS12", CryptoServicesProvider.instance()));
    }

    /**
     * Check if format is already added.
     *
     * @param format target format
     *
     * @return <code>true</code> if format is already added; <code>false</code> otherwise
     */
    public boolean containsFormat(KeyPairFormat format) {
        Assert.notNull(format);
        return this.formats.contains(format);
    }

    /**
     * Adds a new format to the reader.
     *
     * @param format target format
     *
     * @throws IllegalArgumentException if format is already added
     */
    public void addFormat(KeyPairFormat format) {
        Assert.isTrue(!containsFormat(format), "format already added");
        this.formats.add(format);
    }

    /**
     * Removes format from the reader.
     *
     * @param format target format
     */
    public void removeFormat(KeyPairFormat format) {
        this.formats.remove(format);
    }

    /**
     * Restores key pair from the specified byte array.
     *
     * @param data       key pair data
     * @param passphrase passphrase
     *
     * @return restored key pair
     *
     * @throws com.jscape.inet.ssh.util.keyreader.FormatException
     *          if key pair cannot be restored
     */
    public KeyPair restoreKeyPair(byte[] data, String passphrase)
            throws FormatException {
        Assert.notNull(data);
        Assert.notNull(passphrase);
        FormatException error = null;

        for (Iterator i = this.formats.iterator(); i.hasNext(); ) {
            KeyPairFormat format = (KeyPairFormat) i.next();

            try {
                return format.restoreKeyPair(data, passphrase);
            } catch (FormatException e) {
                error = e;
            } catch (Exception e) {
                error = new FormatException(ExceptionTools.messageOf(e), e);
            }
        }

        throw error != null ? error : new FormatException("cannot restore key pair");
    }

    /**
     * Reads key pair from the specified input stream.
     * Assumes, that input stream contains only one key pair.
     *
     * @param in         input stream
     * @param passphrase passphrase
     *
     * @return readed key pair
     *
     * @throws FormatException if key pair cannot be restored
     * @throws IOException     if I/O error occurs
     */
    public KeyPair restoreKeyPair(InputStream in, String passphrase)
            throws FormatException, IOException {
        Assert.notNull(in);
        Assert.notNull(passphrase);

        ByteArrayOutputStream out = new ByteArrayOutputStream();
        byte[] buf = new byte[BUFFER_SIZE];

        for (int read = in.read(buf); read != -1; read = in.read(buf)) {
            out.write(buf, 0, read);
        }

        return restoreKeyPair(out.toByteArray(), passphrase);
    }

    /**
     * Reads key pair from the specified file.
     *
     * @param file       key pair file
     * @param passphrase passphrase
     *
     * @return readed key pair
     *
     * @throws FormatException if key pair cannot be restored
     * @throws IOException     if I/O error occurs
     */
    public KeyPair restoreKeyPair(File file, String passphrase)
            throws IOException {
        try (InputStream in = Files.newInputStream(file.toPath())) {
            return restoreKeyPair(in, passphrase);
        }
    }

    /**
     * Reads key pair from the specified key string
     *
     * @param key        the private key
     * @param passphrase the private key password
     *
     * @return key pair
     *
     * @throws FormatException if key pair cannot be restored
     * @throws IOException     if I/O error occurs
     */
    public KeyPair restoreKeyPair(String key, String passphrase)
            throws FormatException {
        try {
            return this.restoreKeyPair(key.getBytes(StandardCharsets.UTF_8), passphrase);
        } catch (FormatException e) {
            throw e;
        }
    }
}
