/*
 * Class: OpenSSHFormat
 *
 * Author: Alex Usun
 *
 * Date: 1/9/2005
 */
package com.jscape.inet.ssh.util.keyreader.openssh;

import com.jscape.inet.ssh.util.keyreader.CipherFactory;
import com.jscape.inet.ssh.util.keyreader.FormatException;
import com.jscape.inet.ssh.util.keyreader.KeyPairFormat;
import com.jscape.inet.ssh.util.keyreader.KeyPairSpec;
import com.jscape.inet.util.Base64;
import com.jscape.util.Assert;
import com.jscape.util.StringTools;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.StringReader;
import java.nio.charset.StandardCharsets;
import java.security.GeneralSecurityException;
import java.security.KeyPair;
import java.util.ArrayList;
import java.util.List;

/**
 * OpenSSH key pair format.
 *
 * @author Alex Usun
 */
@SuppressWarnings("unchecked")
public abstract class OpenSSHFormat
        implements KeyPairFormat {

    private static final String START_LINE_PREFIX = "-----BEGIN ";
    private static final String STOP_LINE_PREFIX = "-----END ";
    private static final String LINE_SUFFIX = " PRIVATE KEY-----";
    private static final String KEY_SEPARATOR = ":";
    private static final String PROC_TYPE_FIELD = "Proc-Type" + KEY_SEPARATOR;
    private static final String DEK_INFO_FIELD = "DEK-Info" + KEY_SEPARATOR;
    private static final String DATA_SEPARATOR = ",";

    private final Base64 base64Coder = new Base64();
    private final CipherFactory cipherFactory = new CipherFactory();

    /**
     * Creates a new format instance.
     */
    protected OpenSSHFormat() {
        this.cipherFactory.add("DES-EDE3-CBC", "DESede/CBC/NoPadding", 24);
        this.cipherFactory.add("AES-128-CBC", "AES/CBC/NoPadding", 16);
    }

    /**
     * @see KeyPairFormat#restoreKeyPair(byte[], String)
     */
    public KeyPair restoreKeyPair(byte[] data, String passphrase)
            throws FormatException {
        Assert.notNull(data);
        Assert.notNull(passphrase);

        Record record = restoreRecord(data);

        try {
            if (record.isEncrypted()) {
                KeyFactory keyFactory = new KeyFactory(passphrase, record.getIv());
                record = record.decrypt(this.cipherFactory, keyFactory);
            }

            return restoreKeyPair(record).toKeyPair();
        } catch (IOException e) {
            throw new FormatException(e.getMessage());
        } catch (GeneralSecurityException e) {
            throw new FormatException(e.getMessage());
        }
    }

    /**
     * Returns supported key pair algorithm name.
     *
     * @return supported key pair algorithm name
     */
    protected abstract String getAlgorithm();

    /**
     * Restores key pair.
     *
     * @param record key pair record
     *
     * @return restored key pair
     *
     * @throws IOException if key pair cannot be restored
     */
    protected abstract KeyPairSpec restoreKeyPair(Record record)
            throws IOException;

    /**
     * Restores key pair record.
     *
     * @param data record data
     *
     * @return restored key pair record
     *
     * @throws IllegalArgumentException if record cannot be restored
     */
    private Record restoreRecord(byte[] data) {
        try {
            String[] recordLines = restoreRecordData(new String(data, StandardCharsets.US_ASCII));
            String[] procTypeParts = StringTools.split(getFieldValue(PROC_TYPE_FIELD, recordLines), DATA_SEPARATOR, false);
            String procType = procTypeParts.length >= 2 ? procTypeParts[0] : "";
            String[] dekInfoParts = StringTools.split(getFieldValue(DEK_INFO_FIELD, recordLines), DATA_SEPARATOR, false);
            String encryprionAlgorithm = dekInfoParts.length >= 2 ? dekInfoParts[0] : Record.NO_ENCRYPTION;
            byte[] iv = dekInfoParts.length >= 2 ? createIv(dekInfoParts[1].trim()) : new byte[0];
            byte[] keyData = restoreKeyPairData(recordLines);

            return new Record(getAlgorithm(), procType, encryprionAlgorithm, iv, keyData);
        } catch (Exception e) {
            throw new IllegalArgumentException("bad format");
        }
    }

    /**
     * Restores record data.
     *
     * @param data record data
     *
     * @return record lines
     *
     * @throws IOException if an I/O error occurs
     */
    private String[] restoreRecordData(String data)
            throws IOException {
        BufferedReader reader = new BufferedReader(new StringReader(data));

        String startLine = reader.readLine();
        Assert.isTrue(startLine != null && startLine.startsWith(START_LINE_PREFIX + getAlgorithm() + LINE_SUFFIX), "Bad format.");

        List lines = new ArrayList();

        for (String line = reader.readLine(); line != null; line = reader.readLine()) {
            lines.add(line);
        }

        String stopLine = (String) lines.remove(lines.size() - 1);
        Assert.isTrue(stopLine.startsWith(STOP_LINE_PREFIX + getAlgorithm() + LINE_SUFFIX), "bad format");

        return (String[]) lines.toArray(new String[0]);
    }

    /**
     * Restores field value.
     *
     * @param name field name
     * @param data record data
     *
     * @return restored field value
     */
    private String getFieldValue(String name, String[] data) {
        for (int i = 0; i < data.length; i++) {
            if (data[i].startsWith(name)) {
                return data[i].substring(name.length()).trim();
            }
        }

        return "";
    }

    /**
     * Creates cipher IV.
     *
     * @param str IV string
     *
     * @return cipher IV
     */
    private byte[] createIv(String str) {
        byte[] iv = new byte[str.length() / 2];

        for (int i = 0; i < str.length(); i += 2) {
            iv[i / 2] = (byte) Integer.parseInt(str.substring(i, i + 2), 16);
        }

        return iv;
    }

    /**
     * Restores key data.
     *
     * @param data record lines
     *
     * @return key data
     */
    private byte[] restoreKeyPairData(String[] data) {
        StringBuffer buf = new StringBuffer();

        int i = 0;

        // skip field lines
        for (; i < data.length && data[i].indexOf(KEY_SEPARATOR) >= 0; ++i) {
            // empty
        }

        for (; i < data.length; ++i) {
            buf.append(data[i]);
        }

        return this.base64Coder.decode(buf.toString().toCharArray());
    }

    /**
     * @see Object#equals(Object)
     */
    public boolean equals(Object obj) {
        if (this == obj) {
            return true;
        }

        if (obj == null) {
            return false;
        }

        return getClass().equals(obj.getClass());
    }

}
