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

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.StreamCorruptedException;
import java.math.BigInteger;
import java.security.GeneralSecurityException;
import java.security.KeyFactory;
import java.security.KeyPair;
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.NamedParameterSpec;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import org.apache.sshd.common.NamedResource;
import org.apache.sshd.common.config.keys.FilePasswordProvider;
import org.apache.sshd.common.session.SessionContext;
import org.apache.sshd.common.util.ArrayUtil;
import org.apache.sshd.common.util.GenericUtils;
import org.apache.sshd.common.util.io.NoCloseInputStream;
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 Ed25519PEMResourceKeyParser extends AbstractPEMResourceKeyPairParser {

    public static final String BEGIN_MARKER = "BEGIN OPENSSH PRIVATE KEY";

    public static final List<String> BEGINNERS = Collections.singletonList(BEGIN_MARKER);

    public static final String END_MARKER = "END OPENSSH PRIVATE KEY";
    public static final List<String> ENDERS = Collections.singletonList(END_MARKER);

    /**
     * @see <A HREF="https://tools.ietf.org/html/rfc8410#section-3">RFC8412 section 3</A>
     */
    public static final String ED25519_OID = "1.3.101.112";

    public static final Ed25519PEMResourceKeyParser INSTANCE = new Ed25519PEMResourceKeyParser();

    public Ed25519PEMResourceKeyParser() {
        super("Ed25519", ED25519_OID, BEGINNERS, ENDERS);
    }

    @Override
    public Collection<KeyPair> extractKeyPairs(SessionContext session,
                                               NamedResource resourceKey,
                                               String beginMarker, String endMarker,
                                               FilePasswordProvider passwordProvider,
                                               byte[] rawKey,
                                               Map<String, String> headers)
            throws IOException, GeneralSecurityException {
        try (InputStream stream = new ByteArrayInputStream(rawKey)) {
            KeyPair kp = parseEd25519KeyPair(stream, false);
            return Collections.singletonList(kp);
        }
    }

    public static KeyPair parseEd25519KeyPair(InputStream inputStream,
                                               boolean okToClose) throws IOException, GeneralSecurityException {
        try (DERParser parser = new DERParser(NoCloseInputStream.resolveInputStream(inputStream, okToClose))) {
            return parseKeyPair(parser);
        }
    }

    /*
     * See https://tools.ietf.org/html/rfc8410#section-7
     *
     * SEQUENCE {
     *      INTEGER 0x00 (0 decimal)
     *      SEQUENCE {
     *          OBJECTIDENTIFIER 1.3.101.112
     *      }
     *      OCTETSTRING keyData
     * }
     *
     * NOTE: there is another variant that also has some extra parameters
     * but it has the same "prefix" structure so we don't care
     */
    private static KeyPair parseKeyPair(DERParser parser) throws IOException, GeneralSecurityException {
        ASN1Object obj = parser.readObject();
        if (obj == null) {
            throw new StreamCorruptedException("Missing version value");
        }
        BigInteger version = obj.asInteger();
        if (!BigInteger.ZERO.equals(version)) {
            throw new StreamCorruptedException("Invalid version: " + version);
        }
        obj = parser.readObject();
        if (obj == null) {
            throw new StreamCorruptedException("Missing OID container");
        }
        ASN1Type objType = obj.getObjType();
        if (objType != ASN1Type.SEQUENCE) {
            throw new StreamCorruptedException("Unexpected OID object type: " + objType);
        }
        List<Integer> curveOid;
        try (DERParser oidParser = obj.createParser()) {
            obj = oidParser.readObject();
            if (obj == null) {
                throw new StreamCorruptedException("Missing OID value");
            }
            curveOid = obj.asOID();
        }
        String oid = GenericUtils.join(curveOid, '.');
        // TODO modify if more curves supported
        if (!ED25519_OID.equals(oid)) {
            throw new StreamCorruptedException("Unsupported curve OID: " + oid);
        }
        obj = parser.readObject();
        if (obj == null) {
            throw new StreamCorruptedException("Missing key data");
        }
        return decodeKeyPair(obj.getValue());
    }

    public static KeyPair decodeKeyPair(byte[] keyData) throws IOException, GeneralSecurityException {
        EdECPrivateKey privateKey = getPrivateKey(keyData);
        EdECPublicKey publicKey = getPublicKey(keyData);
        return new KeyPair(publicKey, privateKey);
    }

    private static EdECPrivateKey getPrivateKey(byte[] keyData) throws IOException, GeneralSecurityException {
        try (DERParser parser = new DERParser(keyData)) {
            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);
            }
            return generatePrivateKey(obj.getValue());
        }
    }

    private static EdECPrivateKey generatePrivateKey(byte[] seed) throws GeneralSecurityException {
        NamedParameterSpec spec = NamedParameterSpec.ED25519;
        EdECPrivateKeySpec keySpec = new EdECPrivateKeySpec(spec, seed);
        KeyFactory factory = KeyFactory.getInstance("EdDSA");
        return (EdECPrivateKey) factory.generatePrivate(keySpec);
    }

    private static EdECPublicKey getPublicKey(byte[] arr) throws GeneralSecurityException {
        byte msb = arr[arr.length - 1];
        boolean xOdd = (msb & 0x80) != 0;
        arr[arr.length - 1] &= (byte) 0x7F;
        ArrayUtil.reverse(arr);
        BigInteger y = new BigInteger(1, arr);
        NamedParameterSpec namedParameterSpec = NamedParameterSpec.ED25519;
        EdECPoint edECPoint = new EdECPoint(xOdd, y);
        EdECPublicKeySpec publicKeySpec = new EdECPublicKeySpec(namedParameterSpec, edECPoint);
        KeyFactory factory = KeyFactory.getInstance("EdDSA");
        return (EdECPublicKey) factory.generatePublic(publicKeySpec);
    }
}
