/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements. See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership. The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License. You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied. See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
package org.apache.sshd.common.config.keys.loader.openssh;

import java.io.IOException;
import java.math.BigInteger;
import java.security.GeneralSecurityException;
import java.security.Key;
import java.security.KeyFactory;
import java.security.KeyPair;
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.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.TreeMap;

import org.apache.sshd.common.NamedResource;
import org.apache.sshd.common.config.keys.FilePasswordProvider;
import org.apache.sshd.common.config.keys.PrivateKeyEntryDecoder;
import org.apache.sshd.common.config.keys.loader.AbstractKeyPairResourceParser;
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.ValidateUtils;

/**
 * Basic support for <A HREF=
 * "http://cvsweb.openbsd.org/cgi-bin/cvsweb/src/usr.bin/ssh/PROTOCOL.key?rev=1.1&content-type=text/x-cvsweb-markup">OpenSSH
 * key file(s)</A>
 * 
 * @author <a href="mailto:dev@mina.apache.org">Apache MINA SSHD Project</a>
 */
public class OpenSSHKeyPairResourceParser extends AbstractKeyPairResourceParser {

    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);

    public static final String AUTH_MAGIC = "openssh-key-v1";
    public static final OpenSSHKeyPairResourceParser INSTANCE = new OpenSSHKeyPairResourceParser();

    private static final Map<String, PrivateKeyEntryDecoder<?, ?>> BY_KEY_TYPE_DECODERS_MAP
            = new TreeMap<>(String.CASE_INSENSITIVE_ORDER);

    private static final Map<Class<?>, PrivateKeyEntryDecoder<?, ?>> BY_KEY_CLASS_DECODERS_MAP = new HashMap<>();

    static {
        registerPrivateKeyEntryDecoder(OpenSSHRSAPrivateKeyEntryDecoder.INSTANCE);
        registerPrivateKeyEntryDecoder(OpenSSHDSSPrivateKeyEntryDecoder.INSTANCE);
        registerPrivateKeyEntryDecoder(OpenSSHECDSAPrivateKeyEntryDecoder.INSTANCE);
        registerPrivateKeyEntryDecoder(OpenSSHEd25519PrivateKeyEntryDecoder.INSTANCE);
    }

    public OpenSSHKeyPairResourceParser() {
        super(BEGINNERS, ENDERS);
    }

    /**
     * Ed25519 only. No password.
     *
     * @param session                  The {@link SessionContext} for invoking this load command - may be {@code null}
     *                                  if not invoked within a session context (e.g., offline tool or session unknown).
     * @param resourceKey              A hint as to the origin of the text lines
     * @param beginMarker              The line containing the begin marker
     * @param endMarker                The line containing the end marker
     * @param passwordProvider         The {@link FilePasswordProvider} to use in case the data is encrypted - may be
     *                                  {@code null} if no encrypted
     * @param rawKey                    The bytes from the lines containing the data
     * @param headers                  Any headers that may have been available when data was read
     * @return
     * @throws IOException
     * @throws GeneralSecurityException
     */
    @Override
    public Collection<KeyPair> extractKeyPairs(SessionContext session,
                                               NamedResource resourceKey,
                                               String beginMarker,
                                               String endMarker,
                                               FilePasswordProvider passwordProvider,
                                               byte[] rawKey,
                                               Map<String, String> headers)
            throws IOException, GeneralSecurityException {
        OpenSsh openSsh = decodeFromOpenSsh(rawKey);
        NamedParameterSpec params = NamedParameterSpec.ED25519;
        EdECPrivateKeySpec edECPrivateKeySpec = new EdECPrivateKeySpec(params, openSsh.privatekey);
        byte[] arr = openSsh.publickey;
        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);
        EdECPoint edECPoint = new EdECPoint(xOdd, y);
        EdECPublicKeySpec edECPublicKeySpec = new EdECPublicKeySpec(params, edECPoint);
        KeyFactory keyFactory = KeyFactory.getInstance("EdDSA");
        KeyPair keyPair = new KeyPair(keyFactory.generatePublic(edECPublicKeySpec), keyFactory.generatePrivate(edECPrivateKeySpec));
        return List.of(keyPair);
    }

    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;
    }

    /**
     * @param  decoder                  The decoder to register
     * @throws IllegalArgumentException if no decoder or not key type or no supported names for the decoder
     * @see                             PrivateKeyEntryDecoder#getPublicKeyType()
     * @see                             PrivateKeyEntryDecoder#getSupportedKeyTypes()
     */
    public static void registerPrivateKeyEntryDecoder(PrivateKeyEntryDecoder<?, ?> decoder) {
        Objects.requireNonNull(decoder, "No decoder specified");

        Class<?> pubType = Objects.requireNonNull(decoder.getPublicKeyType(), "No public key type declared");
        Class<?> prvType = Objects.requireNonNull(decoder.getPrivateKeyType(), "No private key type declared");
        synchronized (BY_KEY_CLASS_DECODERS_MAP) {
            BY_KEY_CLASS_DECODERS_MAP.put(pubType, decoder);
            BY_KEY_CLASS_DECODERS_MAP.put(prvType, decoder);
        }

        Collection<String> names
                = ValidateUtils.checkNotNullAndNotEmpty(decoder.getSupportedKeyTypes(), "No supported key type");
        synchronized (BY_KEY_TYPE_DECODERS_MAP) {
            for (String n : names) {
                PrivateKeyEntryDecoder<?, ?> prev = BY_KEY_TYPE_DECODERS_MAP.put(n, decoder);
                if (prev != null) {
                    continue; // debug breakpoint
                }
            }
        }
    }

    /**
     * @param  keyType The {@code OpenSSH} key type string - e.g., {@code ssh-rsa, ssh-dss} - ignored if
     *                 {@code null}/empty
     * @return         The registered {@link PrivateKeyEntryDecoder} or {code null} if not found
     */
    public static PrivateKeyEntryDecoder<?, ?> getPrivateKeyEntryDecoder(String keyType) {
        if (GenericUtils.isEmpty(keyType)) {
            return null;
        }

        synchronized (BY_KEY_TYPE_DECODERS_MAP) {
            return BY_KEY_TYPE_DECODERS_MAP.get(keyType);
        }
    }

    /**
     * @param  kp The {@link KeyPair} to examine - ignored if {@code null}
     * @return    The matching {@link PrivateKeyEntryDecoder} provided <U>both</U> the public and private keys have the
     *            same decoder - {@code null} if no match found
     * @see       #getPrivateKeyEntryDecoder(Key)
     */
    public static PrivateKeyEntryDecoder<?, ?> getPrivateKeyEntryDecoder(KeyPair kp) {
        if (kp == null) {
            return null;
        }

        PrivateKeyEntryDecoder<?, ?> d1 = getPrivateKeyEntryDecoder(kp.getPublic());
        PrivateKeyEntryDecoder<?, ?> d2 = getPrivateKeyEntryDecoder(kp.getPrivate());
        if (d1 == d2) {
            return d1;
        } else {
            return null; // some kind of mixed keys...
        }
    }

    /**
     * @param  key The {@link Key} (public or private) - ignored if {@code null}
     * @return     The registered {@link PrivateKeyEntryDecoder} for this key or {code null} if no match found
     * @see        #getPrivateKeyEntryDecoder(Class)
     */
    public static PrivateKeyEntryDecoder<?, ?> getPrivateKeyEntryDecoder(Key key) {
        if (key == null) {
            return null;
        } else {
            return getPrivateKeyEntryDecoder(key.getClass());
        }
    }

    /**
     * @param  keyType The key {@link Class} - ignored if {@code null} or not a {@link Key} compatible type
     * @return         The registered {@link PrivateKeyEntryDecoder} or {code null} if no match found
     */
    public static PrivateKeyEntryDecoder<?, ?> getPrivateKeyEntryDecoder(Class<?> keyType) {
        if ((keyType == null) || (!Key.class.isAssignableFrom(keyType))) {
            return null;
        }
        synchronized (BY_KEY_TYPE_DECODERS_MAP) {
            PrivateKeyEntryDecoder<?, ?> decoder = BY_KEY_CLASS_DECODERS_MAP.get(keyType);
            if (decoder != null) {
                return decoder;
            }

            // in case it is a derived class
            for (PrivateKeyEntryDecoder<?, ?> dec : BY_KEY_CLASS_DECODERS_MAP.values()) {
                Class<?> pubType = dec.getPublicKeyType();
                Class<?> prvType = dec.getPrivateKeyType();
                if (pubType.isAssignableFrom(keyType) || prvType.isAssignableFrom(keyType)) {
                    return dec;
                }
            }
        }
        return null;
    }
}
