package org.apache.sshd.common.config.keys.loader.openssh;

import java.io.IOException;
import java.io.InputStream;
import java.io.StreamCorruptedException;
import java.security.GeneralSecurityException;
import java.security.InvalidKeyException;
import java.security.KeyFactory;
import java.security.KeyPairGenerator;
import java.security.interfaces.EdECPrivateKey;
import java.security.interfaces.EdECPublicKey;
import java.security.spec.EdECPoint;
import java.security.spec.EdECPrivateKeySpec;
import java.security.spec.EdECPublicKeySpec;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.NamedParameterSpec;
import java.util.Arrays;
import java.util.Collections;
import org.apache.sshd.common.config.keys.FilePasswordProvider;
import org.apache.sshd.common.config.keys.impl.AbstractPrivateKeyEntryDecoder;
import org.apache.sshd.common.keyprovider.KeyPairProvider;
import org.apache.sshd.common.session.SessionContext;
import org.apache.sshd.common.util.io.SecureByteArrayOutputStream;
import org.apache.sshd.common.util.io.der.ASN1Object;
import org.apache.sshd.common.util.io.der.ASN1Type;
import org.apache.sshd.common.util.io.der.DERParser;

public class OpenSSHEd25519PrivateKeyEntryDecoder extends AbstractPrivateKeyEntryDecoder<EdECPublicKey, EdECPrivateKey> {

    public static final OpenSSHEd25519PrivateKeyEntryDecoder INSTANCE = new OpenSSHEd25519PrivateKeyEntryDecoder();

    public OpenSSHEd25519PrivateKeyEntryDecoder() {
        super(EdECPublicKey.class, EdECPrivateKey.class, Collections.singletonList(KeyPairProvider.SSH_ED25519));
    }

    @Override
    public EdECPrivateKey decodePrivateKey(SessionContext session,
                                           String keyType,
                                           FilePasswordProvider passwordProvider,
                                           InputStream keyData)
            throws IOException, GeneralSecurityException {
        if (!"ssh-ed25519".equals(keyType)) {
            throw new InvalidKeyException("Unsupported key type: " + keyType);
        }
        byte[] b = keyData.readAllBytes();
        OpenSsh openSsh = decodeFromOpenSsh(b);
        NamedParameterSpec params = NamedParameterSpec.ED25519;
        EdECPrivateKeySpec spec = new EdECPrivateKeySpec(params, openSsh.privatekey);
        KeyFactory factory = KeyFactory.getInstance("EdDSA");
        return (EdECPrivateKey) factory.generatePrivate(spec);
    }

    @Override
    public String encodePrivateKey(SecureByteArrayOutputStream s,
                                   EdECPrivateKey key,
                                   EdECPublicKey pubKey) throws IOException {
        throw new UnsupportedOperationException();
    }

    @Override
    public EdECPublicKey recoverPublicKey(EdECPrivateKey key) throws GeneralSecurityException {
        throw new UnsupportedOperationException();
    }

    @Override
    public EdECPublicKey clonePublicKey(EdECPublicKey key) throws GeneralSecurityException {
        if (key == null) {
            return null;
        } else {
            NamedParameterSpec namedParameterSpec = NamedParameterSpec.ED25519;
            EdECPoint edECPoint = key.getPoint();
            EdECPublicKeySpec spec = new EdECPublicKeySpec(namedParameterSpec, edECPoint);
            return generatePublicKey(spec);
        }
    }

    @Override
    public EdECPrivateKey clonePrivateKey(EdECPrivateKey key) throws GeneralSecurityException {
        if (key == null) {
            return null;
        } else {
            try (DERParser parser = new DERParser(key.getEncoded())) {
                ASN1Object obj = parser.readObject();
                if (obj == null) {
                    throw new StreamCorruptedException("Missing key data container");
                }
                ASN1Type objType = obj.getObjType();
                if (objType != ASN1Type.OCTET_STRING) {
                    throw new StreamCorruptedException("Mismatched key data container type: " + objType);
                }
                NamedParameterSpec spec = NamedParameterSpec.ED25519;
                EdECPrivateKeySpec keySpec = new EdECPrivateKeySpec(spec, obj.getValue());
                KeyFactory factory = KeyFactory.getInstance("EdDSA");
                return (EdECPrivateKey) factory.generatePrivate(keySpec);
            } catch (IOException e) {
                throw new GeneralSecurityException(e);
            }
        }
    }

    @Override
    public KeyPairGenerator getKeyPairGenerator() throws GeneralSecurityException {
        return KeyPairGenerator.getInstance("Ed25519");
    }

    @Override
    public KeyFactory getKeyFactoryInstance() throws GeneralSecurityException {
        return KeyFactory.getInstance("EdDSA");
    }

    private static OpenSsh decodeFromOpenSsh(byte[] rawKey) throws InvalidKeySpecException {
        OpenSsh openSsh = new OpenSsh();
        openSsh.verify = Arrays.copyOfRange(rawKey, 0, OPENSSH_KEY_V1.length());
        if (!new String(openSsh.verify).equals(OPENSSH_KEY_V1)) {
            throw new InvalidKeySpecException("invalid OpenSSH key file");
        }
        boolean occurred = false;
        int index = 0;
        for (int i = 0; i < rawKey.length; i++) {
            if (rawKey[i] == 's'
                    && rawKey[i + 1] == 's'
                    && rawKey[i + 2] == 'h'
                    && rawKey[i + 3] == '-'
                    && rawKey[i + 4] == 'e'
                    && rawKey[i + 5] == 'd'
                    && rawKey[i + 6] == '2'
                    && rawKey[i + 7] == '5'
                    && rawKey[i + 8] == '5'
                    && rawKey[i + 9] == '1'
                    && rawKey[i + 10] == '9'
                    && rawKey[i + 11] == 0x00
                    && rawKey[i + 12] == 0x00
                    && rawKey[i + 13] == 0x00
                    && rawKey[i + 14] == ' ') {
                index = i + 15;
                if (occurred) {
                    break;
                }
                occurred = true;
            }
        }
        openSsh.publickey = Arrays.copyOfRange(rawKey, index, index + 32);
        index += 32;
        for (int i = index; i < rawKey.length; i++) {
            if (rawKey[i] == 0x00
                    && rawKey[i + 1] == 0x00
                    && rawKey[i + 2] == 0x00
                    && rawKey[i + 3] == '@') {
                index = i + 4;
                break;
            }
        }
        openSsh.privatekey = Arrays.copyOfRange(rawKey, index, index + 32);
        return openSsh;
    }

    private static final String OPENSSH_KEY_V1 = "openssh-key-v1";

    private static class OpenSsh {
        byte[] verify;
        byte[] publickey;
        byte[] privatekey;
    }
}
