package com.atlassian.bitbucket.mesh.util;

import com.google.common.collect.ImmutableSet;
import com.google.common.hash.HashFunction;
import com.google.common.hash.Hashing;
import com.google.common.io.BaseEncoding;
import org.bouncycastle.asn1.pkcs.PrivateKeyInfo;
import org.bouncycastle.asn1.x509.SubjectPublicKeyInfo;
import org.bouncycastle.jcajce.provider.asymmetric.ec.BCECPublicKey;
import org.bouncycastle.jce.ECNamedCurveTable;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.bouncycastle.jce.spec.ECParameterSpec;
import org.bouncycastle.openssl.PEMKeyPair;
import org.bouncycastle.openssl.PEMParser;
import org.bouncycastle.openssl.jcajce.JcaPEMKeyConverter;
import org.bouncycastle.openssl.jcajce.JcaPEMWriter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import java.io.*;
import java.nio.ByteBuffer;
import java.nio.file.Files;
import java.nio.file.NoSuchFileException;
import java.nio.file.Path;
import java.nio.file.attribute.PosixFilePermission;
import java.nio.file.attribute.PosixFilePermissions;
import java.security.KeyPair;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.Security;
import java.security.interfaces.ECPublicKey;
import java.security.interfaces.RSAPublicKey;
import java.util.EnumSet;
import java.util.Locale;
import java.util.Set;

import static java.nio.charset.StandardCharsets.US_ASCII;
import static java.nio.file.StandardOpenOption.*;
import static java.util.Objects.requireNonNull;

/**
 * Utilities for working with {@link KeyPair}s and {@link PublicKey}s.
 * <p>
 * This is part of {@code mesh-shared-common} to allow reuse in several different places:
 * <ul>
 *     <li>{@code mesh-grpc}: Allows integrating with gRPC to facilitate inbound and outbound authentication</li>
 *     <li>{@code mesh-test-cluster}: Allows the test control plane and nodes to run with authentication enabled
 *     and use keys the way they'll be used in production</li>
 *     <li>Bitbucket Server: Allows use by the control plane when managing its own {@link KeyPair} and public
 *     keys from the connected nodes</li>
 * </ul>
 */
public class KeyUtils {

    private static final String ALGORITHM_EC = "EC";
    private static final String ALGORITHM_ECDSA = "ECDSA";
    private static final String ALGORITHM_RSA = "RSA";

    private static final JcaPEMKeyConverter CONVERTER;

    private static final String EC_CURVE_P256 = "nistp256";
    private static final String EC_CURVE_P384 = "nistp384";
    private static final String EC_CURVE_P521 = "nistp521";
    private static final ECParameterSpec EC_SPEC_P256 = ECNamedCurveTable.getParameterSpec("secp256r1");
    private static final ECParameterSpec EC_SPEC_P384 = ECNamedCurveTable.getParameterSpec("secp384r1");
    private static final ECParameterSpec EC_SPEC_P521 = ECNamedCurveTable.getParameterSpec("secp521r1");
    private static final String EC_TYPE_P256 = "ecdsa-sha2-" + EC_CURVE_P256;
    private static final String EC_TYPE_P384 = "ecdsa-sha2-" + EC_CURVE_P384;
    private static final String EC_TYPE_P512 = "ecdsa-sha2-" + EC_CURVE_P521;

    private static final BaseEncoding ENCODING = BaseEncoding.base64().omitPadding();
    private static final HashFunction HASHER = Hashing.sha256();

    private static final String RSA_TYPE = "ssh-rsa";

    private static final Set<PosixFilePermission> TOO_OPEN;
    private static final Set<PosixFilePermission> OWNER_RW;

    private static final Logger log = LoggerFactory.getLogger(KeyUtils.class);

    static {
        BouncyCastleProvider provider = new BouncyCastleProvider();
        Security.addProvider(provider);

        CONVERTER = new JcaPEMKeyConverter().setProvider(provider);

        EnumSet<PosixFilePermission> rw = EnumSet.of(PosixFilePermission.OWNER_READ, PosixFilePermission.OWNER_WRITE);
        OWNER_RW = ImmutableSet.copyOf(rw);

        EnumSet<PosixFilePermission> other = EnumSet.complementOf(rw);
        TOO_OPEN = ImmutableSet.copyOf(other);
    }

    private KeyUtils() {
        throw new UnsupportedOperationException(getClass() + " is not intended to be instantiated");
    }

    @Nonnull
    public static String getFingerprint(@Nonnull PublicKey publicKey) {
        requireNonNull(publicKey, "publicKey");

        byte[] bytes;
        switch (publicKey.getAlgorithm().toUpperCase(Locale.ROOT)) {
            case ALGORITHM_EC:
            case ALGORITHM_ECDSA:
                bytes = encodeEc(publicKey);
                break;
            case ALGORITHM_RSA:
                bytes = encodeRsa(publicKey);
                break;
            default:
                throw new UnsupportedOperationException("Fingerprints are not supported for " +
                        publicKey.getAlgorithm());
        }

        return ENCODING.encode(HASHER.hashBytes(bytes).asBytes());
    }

    @Nonnull
    public static KeyPair parseKeyPairFromPkcs8(@Nonnull String pkcs8) {
        try (StringReader reader = new StringReader(requireNonNull(pkcs8, "pkcs8"));
             PEMParser parser = new PEMParser(reader)) {
            Object keyPair = parser.readObject();
            if (keyPair instanceof PEMKeyPair) {
                return CONVERTER.getKeyPair((PEMKeyPair) keyPair);
            }
            if (keyPair == null) {
                throw new IllegalArgumentException("The provided PKCS8 data did not contain an object");
            }

            throw new IllegalArgumentException("The provided PKCS8 data contained a " + keyPair.getClass());
        } catch (IOException e) {
            throw new IllegalArgumentException("The provided PKCS8 data could not be parsed", e);
        }
    }

    @Nonnull
    public static PrivateKey parsePrivateKeyFromDer(@Nonnull String derBase64) {
        return parsePrivateKeyFromDer(ENCODING.decode(requireNonNull(derBase64, "derBase64")));
    }

    @Nonnull
    public static PrivateKey parsePrivateKeyFromDer(@Nonnull byte[] der) {
        PrivateKeyInfo info = PrivateKeyInfo.getInstance(requireNonNull(der, "der"));

        try {
            return CONVERTER.getPrivateKey(info);
        } catch (IOException e) {
            throw new IllegalArgumentException("The provided DER data could not be parsed");
        }
    }

    @Nonnull
    public static PrivateKey parsePrivateKeyFromPkcs8(@Nonnull String pkcs8) {
        try (StringReader reader = new StringReader(requireNonNull(pkcs8, "pkcs8"));
             PEMParser parser = new PEMParser(reader)) {
            Object privateKey = parser.readObject();
            if (privateKey instanceof PEMKeyPair) {
                //For some key types, the public key on the PEMKeyPair may be null. That causes
                //CONVERTER.getKeyPair to throw. Since we don't need the public key anyway, here,
                //we cast to PEMKeyPair and then explicitly just convert the private key
                return CONVERTER.getPrivateKey(((PEMKeyPair) privateKey).getPrivateKeyInfo());
            }
            if (privateKey instanceof PrivateKeyInfo) {
                return CONVERTER.getPrivateKey((PrivateKeyInfo) privateKey);
            }
            if (privateKey == null) {
                throw new IllegalArgumentException("The provided PKCS8 data did not contain an object");
            }

            throw new IllegalArgumentException("The provided PKCS8 data contained a " + privateKey.getClass());
        } catch (IOException e) {
            throw new IllegalArgumentException("The provided PKCS8 data could not be parsed");
        }
    }

    @Nonnull
    public static PublicKey parsePublicKeyFromDer(@Nonnull String derBase64) {
        return parsePublicKeyFromDer(ENCODING.decode(requireNonNull(derBase64, "derBase64")));
    }

    @Nonnull
    public static PublicKey parsePublicKeyFromDer(@Nonnull byte[] der) {
        SubjectPublicKeyInfo info = SubjectPublicKeyInfo.getInstance(requireNonNull(der, "der"));

        try {
            return CONVERTER.getPublicKey(info);
        } catch (IOException e) {
            throw new IllegalArgumentException("The provided DER data could not be parsed");
        }
    }

    @Nonnull
    public static PublicKey parsePublicKeyFromPkcs8(@Nonnull String pkcs8) {
        try (StringReader reader = new StringReader(requireNonNull(pkcs8, "pkcs8"));
             PEMParser parser = new PEMParser(reader)) {
            Object publicKey = parser.readObject();
            if (publicKey instanceof PEMKeyPair) {
                return CONVERTER.getPublicKey(((PEMKeyPair) publicKey).getPublicKeyInfo());
            }
            if (publicKey instanceof SubjectPublicKeyInfo) {
                return CONVERTER.getPublicKey((SubjectPublicKeyInfo) publicKey);
            }
            if (publicKey == null) {
                throw new IllegalArgumentException("The provided PKCS8 data did not contain an object");
            }

            throw new IllegalArgumentException("The provided PKCS8 data contained a " + publicKey.getClass());
        } catch (IOException e) {
            throw new IllegalArgumentException("The provided PKCS8 data could not be parsed");
        }
    }

    @Nullable
    public static KeyPair readKeyPair(@Nonnull Path path) throws IOException {
        Object keyPair = read(path);
        if (keyPair instanceof PEMKeyPair) {
            return CONVERTER.getKeyPair((PEMKeyPair) keyPair);
        }
        if (keyPair == null) {
            return null;
        }

        throw new IllegalStateException(path.getFileName() + " contains unexpected data: " + keyPair.getClass());
    }

    @Nullable
    public static PublicKey readPublicKey(@Nonnull Path path) throws IOException {
        Object publicKey = read(path);
        if (publicKey instanceof PEMKeyPair) {
            return CONVERTER.getKeyPair((PEMKeyPair) publicKey).getPublic();
        }
        if (publicKey instanceof SubjectPublicKeyInfo) {
            return CONVERTER.getPublicKey((SubjectPublicKeyInfo) publicKey);
        }
        if (publicKey == null) {
            return null;
        }

        throw new IllegalStateException(path.getFileName() + " contains unexpected data: " + publicKey.getClass());
    }

    @Nonnull
    public static String toPkcs8(@Nonnull KeyPair keyPair) throws IOException {
        return toPkcs8(requireNonNull(keyPair, "keyPair").getPrivate());
    }

    @Nonnull
    public static String toPkcs8(@Nonnull PublicKey publicKey) throws IOException {
        return toPkcs8((Object) requireNonNull(publicKey, "publicKey"));
    }

    public static void writeKeyPair(@Nonnull Path path, @Nonnull KeyPair keyPair) throws IOException {
        write(path, requireNonNull(keyPair, "keyPair").getPrivate());
    }

    public static void writePublicKey(@Nonnull Path path, @Nonnull PublicKey publicKey) throws IOException {
        write(path, requireNonNull(publicKey, "publicKey"));
    }

    private static byte[] encodeEc(PublicKey publicKey) {
        //We could also do the things below without any BC code, but it's much easier with BC. Specifically:
        //* Matching ECParameterSpec with the java.security type is more work because it doesn't have equals.
        //* Encoding the point (see RFC 5656 section 3.1) requires bit fiddling and depends on key size.
        BCECPublicKey bcEcKey;
        if (publicKey instanceof BCECPublicKey) {
            bcEcKey = (BCECPublicKey) publicKey;
        } else {
            bcEcKey = new BCECPublicKey((ECPublicKey) publicKey, BouncyCastleProvider.CONFIGURATION);
        }
        ECParameterSpec ecParameterSpec = bcEcKey.getParameters();

        String curve;
        String type;
        if (ecParameterSpec.equals(EC_SPEC_P256)) {
            curve = EC_CURVE_P256;
            type = EC_TYPE_P256;
        } else if (ecParameterSpec.equals(EC_SPEC_P384)) {
            curve = EC_CURVE_P384;
            type = EC_TYPE_P384;
        } else if (ecParameterSpec.equals(EC_SPEC_P521)) {
            curve = EC_CURVE_P521;
            type = EC_TYPE_P512;
        } else {
            throw new IllegalArgumentException("Unable to encode public key: Unsupported elliptic curve");
        }

        byte[] t = type.getBytes(US_ASCII);
        byte[] c = curve.getBytes(US_ASCII);
        byte[] q = bcEcKey.getQ().getEncoded(false);
        ByteBuffer buffer = ByteBuffer.allocate(t.length + c.length + q.length + 12); //+12 for 3 4-byte integers

        buffer.putInt(t.length);
        buffer.put(t);

        buffer.putInt(c.length);
        buffer.put(c);

        buffer.putInt(q.length);
        buffer.put(q);

        return buffer.array();
    }

    private static byte[] encodeRsa(PublicKey publicKey) {
        RSAPublicKey rsaKey = (RSAPublicKey) publicKey;

        byte[] t = RSA_TYPE.getBytes(US_ASCII);
        byte[] exp = rsaKey.getPublicExponent().toByteArray();
        byte[] mod = rsaKey.getModulus().toByteArray();
        boolean padExp = (exp[0] & 0x80) != 0;
        boolean padMod = (mod[0] & 0x80) != 0;

        int size = t.length + exp.length + mod.length + 12; //+12 for 3 4-byte integers
        if (padExp) {
            ++size;
        }
        if (padMod) {
            ++size;
        }
        ByteBuffer buffer = ByteBuffer.allocate(size);

        buffer.putInt(t.length);
        buffer.put(t);

        if (padExp) {
            buffer.putInt(exp.length + 1); //+1 for padding
            buffer.put((byte) 0);
        } else {
            buffer.putInt(exp.length);
        }
        buffer.put(exp);

        if (padMod) {
            buffer.putInt(mod.length + 1); //+1 for padding
            buffer.put((byte) 0);
        } else {
            buffer.putInt(mod.length);
        }
        buffer.put(mod);

        return buffer.array();
    }

    private static Object read(Path path) throws IOException {
        requireNonNull(path, "path");

        try {
            Set<PosixFilePermission> permissions = Files.getPosixFilePermissions(path);
            if (permissions.stream().anyMatch(TOO_OPEN::contains)) {
                log.warn("{} will not be read; its permissions are too loose ({})",
                        path, PosixFilePermissions.toString(permissions));
                return null;
            }

            try (Reader reader = Files.newBufferedReader(path, US_ASCII);
                 PEMParser pemParser = new PEMParser(reader)) {
                return pemParser.readObject();
            }
        } catch (FileNotFoundException | NoSuchFileException e) {
            return null;
        }
    }

    private static String toPkcs8(Object data) throws IOException {
        StringWriter writer = new StringWriter();
        write(writer, data);

        return writer.toString();
    }

    private static void write(Path path, Object data) throws IOException {
        requireNonNull(path, "path");

        try (Writer writer = Files.newBufferedWriter(path, US_ASCII, CREATE, TRUNCATE_EXISTING, WRITE)) {
            write(writer, data);
        }

        //Restrict access to the written key to the owner
        Files.setPosixFilePermissions(path, OWNER_RW);
    }

    private static void write(Writer writer, Object data) throws IOException {
        try (JcaPEMWriter pemWriter = new JcaPEMWriter(writer)) {
            pemWriter.writeObject(data, null);
        }
    }
}
