/*
 * 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.putty;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.StreamCorruptedException;
import java.security.GeneralSecurityException;
import java.security.KeyPair;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.util.Base64;
import java.util.Base64.Decoder;
import java.util.Collection;
import java.util.Collections;
import java.util.List;

import org.apache.sshd.common.config.keys.FilePasswordProvider;
import org.apache.sshd.common.config.keys.impl.AbstractIdentityResourceLoader;
import org.apache.sshd.common.config.keys.loader.KeyPairResourceParser;
import org.apache.sshd.common.util.GenericUtils;
import org.apache.sshd.common.util.ValidateUtils;

/**
 * @param <PUB> Generic public key type
 * @param <PRV> Generic private key type
 * @author <a href="mailto:dev@mina.apache.org">Apache MINA SSHD Project</a>
 */
public abstract class AbstractPuttyKeyDecoder<PUB extends PublicKey, PRV extends PrivateKey>
                extends AbstractIdentityResourceLoader<PUB, PRV>
                implements PuttyKeyPairResourceParser<PUB, PRV> {
    public static final String ENCRYPTION_HEADER = "Encryption";

    protected AbstractPuttyKeyDecoder(Class<PUB> pubType, Class<PRV> prvType, Collection<String> names) {
        super(pubType, prvType, names);
    }

    @Override
    public boolean canExtractKeyPairs(String resourceKey, List<String> lines)
            throws IOException, GeneralSecurityException {
        if (!PuttyKeyPairResourceParser.super.canExtractKeyPairs(resourceKey, lines)) {
            return false;
        }

        for (String l : lines) {
            l = GenericUtils.trimToEmpty(l);
            if (!l.startsWith(KEY_FILE_HEADER_PREFIX)) {
                continue;
            }

            int pos = l.indexOf(':');
            if ((pos <= 0) || (pos >= (l.length() - 1))) {
                return false;
            }

            Collection<String> supported = getSupportedTypeNames();
            String typeValue = l.substring(pos + 1).trim();
            return supported.contains(typeValue);
        }

        return false;
    }

    @Override
    public Collection<KeyPair> loadKeyPairs(
            String resourceKey, FilePasswordProvider passwordProvider, List<String> lines)
                    throws IOException, GeneralSecurityException {
        List<String> pubLines = Collections.emptyList();
        List<String> prvLines = Collections.emptyList();
        String prvEncryption = null;
        for (int index = 0, numLines = lines.size(); index < numLines; index++) {
            String l = lines.get(index);
            l = GenericUtils.trimToEmpty(l);
            int pos = l.indexOf(':');
            if ((pos <= 0) || (pos >= (l.length() - 1))) {
                continue;
            }

            String hdrName = l.substring(0, pos).trim();
            String hdrValue = l.substring(pos + 1).trim();
            switch (hdrName) {
                case ENCRYPTION_HEADER:
                    if (prvEncryption != null) {
                        throw new StreamCorruptedException("Duplicate " + hdrName + " in" + resourceKey);
                    }
                    prvEncryption = hdrValue;
                    break;
                case PUBLIC_LINES_HEADER:
                    pubLines = extractDataLines(resourceKey, lines, index + 1, hdrName, hdrValue, pubLines);
                    index += pubLines.size();
                    break;
                case PRIVATE_LINES_HEADER:
                    prvLines = extractDataLines(resourceKey, lines, index + 1, hdrName, hdrValue, prvLines);
                    index += prvLines.size();
                    break;
                default:    // ignored
            }
        }

        return loadKeyPairs(resourceKey, pubLines, prvLines, prvEncryption, passwordProvider);
    }

    public static List<String> extractDataLines(
            String resourceKey, List<String> lines, int startIndex, String hdrName, String hdrValue, List<String> curLines)
                throws IOException {
        if (GenericUtils.size(curLines) > 0) {
            throw new StreamCorruptedException("Duplicate " + hdrName + " in " + resourceKey);
        }

        int numLines;
        try {
            numLines = Integer.parseInt(hdrValue);
        } catch (NumberFormatException e) {
            throw new StreamCorruptedException("Bad " + hdrName + " value (" + hdrValue + ") in " + resourceKey);
        }

        int endIndex = startIndex + numLines;
        int totalLines = lines.size();
        if (endIndex > totalLines) {
            throw new StreamCorruptedException("Excessive " + hdrName + " value (" + hdrValue + ") in " + resourceKey);
        }

        return lines.subList(startIndex, endIndex);
    }

    public Collection<KeyPair> loadKeyPairs(
            String resourceKey, List<String> pubLines, List<String> prvLines, String prvEncryption, FilePasswordProvider passwordProvider)
                throws IOException, GeneralSecurityException {
        return loadKeyPairs(resourceKey,
                KeyPairResourceParser.joinDataLines(pubLines), KeyPairResourceParser.joinDataLines(prvLines),
                prvEncryption, passwordProvider);
    }

    public Collection<KeyPair> loadKeyPairs(
            String resourceKey, String pubData, String prvData, String prvEncryption, FilePasswordProvider passwordProvider)
                throws IOException, GeneralSecurityException {
        Decoder b64Decoder = Base64.getDecoder();
        byte[] pubBytes = b64Decoder.decode(pubData);
        byte[] prvBytes = b64Decoder.decode(prvData);
        String password = null;
        if ((GenericUtils.length(prvEncryption) > 0)
                && (!NO_PRIVATE_KEY_ENCRYPTION_VALUE.equalsIgnoreCase(prvEncryption))) {
            password = passwordProvider.getPassword(resourceKey);
        }

        if (GenericUtils.isEmpty(prvEncryption)
                || NO_PRIVATE_KEY_ENCRYPTION_VALUE.equalsIgnoreCase(prvEncryption)
                || GenericUtils.isEmpty(password)) {
            return loadKeyPairs(resourceKey, pubBytes, prvBytes);
        }

        // format is "<cipher><bits>-<mode>" - e.g., "aes256-cbc"
        int pos = prvEncryption.indexOf('-');
        if (pos <= 0) {
            throw new StreamCorruptedException("Missing private key encryption mode in " + prvEncryption);
        }

        String mode = prvEncryption.substring(pos + 1).toUpperCase();
        String algName = null;
        int numBits = 0;
        for (int index = 0; index < pos; index++) {
            char ch = prvEncryption.charAt(index);
            if ((ch >= '0') && (ch <= '9')) {
                algName = prvEncryption.substring(0, index).toUpperCase();
                numBits = Integer.parseInt(prvEncryption.substring(index, pos));
                break;
            }
        }

        if (GenericUtils.isEmpty(algName) || (numBits <= 0)) {
            throw new StreamCorruptedException("Missing private key encryption algorithm details in " + prvEncryption);
        }

        prvBytes = PuttyKeyPairResourceParser.decodePrivateKeyBytes(prvBytes, algName, numBits, mode, password);
        return loadKeyPairs(resourceKey, pubBytes, prvBytes);
    }

    public Collection<KeyPair> loadKeyPairs(String resourceKey, byte[] pubData, byte[] prvData)
            throws IOException, GeneralSecurityException {
        ValidateUtils.checkNotNullAndNotEmpty(pubData, "No public key data in %s", resourceKey);
        ValidateUtils.checkNotNullAndNotEmpty(prvData, "No private key data in %s", resourceKey);
        try (InputStream pubStream = new ByteArrayInputStream(pubData);
             InputStream prvStream = new ByteArrayInputStream(prvData)) {
            return loadKeyPairs(resourceKey, pubStream, prvStream);
        }
    }

    public Collection<KeyPair> loadKeyPairs(String resourceKey, InputStream pubData, InputStream prvData)
            throws IOException, GeneralSecurityException {
        try (PuttyKeyReader pubReader =
                new PuttyKeyReader(ValidateUtils.checkNotNull(pubData, "No public key data in %s", resourceKey));
             PuttyKeyReader prvReader =
                new PuttyKeyReader(ValidateUtils.checkNotNull(prvData, "No private key data in %s", resourceKey))) {
            return loadKeyPairs(resourceKey, pubReader, prvReader);
        }
    }

    public abstract Collection<KeyPair> loadKeyPairs(String resourceKey, PuttyKeyReader pubReader, PuttyKeyReader prvReader)
            throws IOException, GeneralSecurityException;
}
