package com.atlassian.bitbucket.mesh.util;

import com.atlassian.bitbucket.mesh.io.MoreFiles;
import com.atlassian.bitbucket.mesh.junit.LogListener;
import com.atlassian.bitbucket.mesh.junit.TemporaryDirectory;
import io.jsonwebtoken.Jwts;
import io.jsonwebtoken.SignatureAlgorithm;
import io.jsonwebtoken.impl.crypto.EllipticCurveProvider;
import io.jsonwebtoken.security.Keys;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.junit.Rule;
import org.junit.Test;

import java.io.IOException;
import java.nio.file.Files;
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.SecureRandom;
import java.util.Locale;
import java.util.Set;

import static org.junit.Assert.*;

public class KeyUtilsTest {

    private static final SecureRandom RANDOM;

    static {
        RANDOM = new SecureRandom();
        RANDOM.nextBytes(new byte[64]);
    }

    @Rule
    public final LogListener logListener = new LogListener(KeyUtils.class);
    @Rule
    public final TemporaryDirectory tempDir = new TemporaryDirectory();

    //For the DER tests, the text used is from writePublicKey with the BEGIN and END delimiters and any
    //padding (trailing =s) removed

    @Test
    public void testParsePublicKeyFromDerForEcKey() {
        PublicKey publicKey = KeyUtils.parsePublicKeyFromDer("MFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEaucvQzbD4ZfCN+h" +
                "8xuyCxFDa9tzRU2wikaQcuXS8gR/bIyzejKGQ3ylARv1oV4Pc5JWhOh4i+LB040JALAmZEw");
        assertEquals("Vo9tVC71o+0Fc6Re9dp096OqjzsWSmGy9dsN593XDuY", KeyUtils.getFingerprint(publicKey));
    }

    @Test
    public void testParsePublicKeyFromDerForRsaKey() {
        PublicKey publicKey = KeyUtils.parsePublicKeyFromDer("MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAspADKDS" +
                "2sRRXyJaZ7l10aiftkPOqYj2Qu/R7fU851GrSs0egLd0w50hNnWiV+5FOt8XScI7HbisyNE+AiPVb3oEozjGMxzU2EMH9vR6" +
                "hFH3y20taEcGrMarc8K/0+s12GXaDcn2DVSBZ9VMv3qdyorFFXfT4QwvtInhydpbHsO+VtkwC+7/0zZ46XljzyZhSz7oK34k" +
                "0W39JO40GBm8fJe2TLiBg7te+8Qsl0apsZK14TSZidM9PSI6q71HSWpU2A43Iaoue5X/eH6RFxtBByfpe/iAKZDkeMyHuAM6" +
                "b475oi8allRTF19gXXaufbRtyRb5BaEoYKNkKu4w65bp57QIDAQAB");
        assertEquals("ystB64A4bjThuk7satYvlvijD+XWUOHlo2bkH69/sXM", KeyUtils.getFingerprint(publicKey));
    }

    //For the PKCS8 tests, the text used is directly from a generated public key

    @Test
    public void testParsePublicKeyFromPkcs8ForEcKey() {
        PublicKey publicKey = KeyUtils.parsePublicKeyFromPkcs8("-----BEGIN PUBLIC KEY-----\n" +
                "MFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAE35kq/Peh/gRq3lR18wawfEkFYrD+\n" +
                "cGPHsIa1TX1v+XVMEjtm0QNJzElta/pbzAqqXN4sC7E3QuUlW0C7miioow==\n" +
                "-----END PUBLIC KEY-----");
        assertEquals("hYagSx8RN7Y2zrXhyp/ppUjK3E+gO+X6jwerCMwghI8", KeyUtils.getFingerprint(publicKey));
    }

    @Test
    public void testParsePublicKeyFromPkcs8ForRsaKey() {
        PublicKey publicKey = KeyUtils.parsePublicKeyFromPkcs8("-----BEGIN PUBLIC KEY-----\n" +
                "MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEA5YXMVMoayvrbs61lGFgj\n" +
                "DCnzCIr/OIryjoJhPQkkXZlPTlzl0legvCLDYWRaWDfoZYSNTDPGWsHKLZFNKpQw\n" +
                "rdp5VPNWH2zJEznvtgdw08AlsljR4zw2HVDajsDw6UKQsp6VcKO/Qq1BGn8EAaYF\n" +
                "sXsw30sMnvP54qST8xQkpit5BPEL5jOfgcC0tA1eJ2xdbgPOVXHNfJvHT2nmCjMx\n" +
                "VEqxkESEv3PK6yZzbzs8kRVHaL8zJ95OGruBS9glDQV250s0cLRMrZfeale1oJD8\n" +
                "89HgvF7H3NFXpafYUvgvRgrLp5n1HweGKREKVqvuA9cQxYtRmAE9xtnuFyogtAQy\n" +
                "wQIDAQAB\n" +
                "-----END PUBLIC KEY-----");
        assertEquals("G8rcD6pfoGqE0BghRdKGZIce2SLldchn9nVT8iEm8Vo", KeyUtils.getFingerprint(publicKey));
    }

    @Test
    public void testParsePrivateKeyFromDerAndPkcs8ForEcKey() throws IOException {
        KeyPair keyPair = generateEcKeyPair(SignatureAlgorithm.ES256);

        byte[] privateDer = keyPair.getPrivate().getEncoded();
        PrivateKey privateKey = KeyUtils.parsePrivateKeyFromDer(privateDer);
        assertEquals(keyPair.getPrivate(), privateKey);

        String pkcs8 = KeyUtils.toPkcs8(keyPair);
        privateKey = KeyUtils.parsePrivateKeyFromPkcs8(pkcs8);
        assertEquals(keyPair.getPrivate(), privateKey);
    }

    @Test
    public void testParsePrivateKeyFromDerAndPkcs8ForRsaKey() throws IOException {
        KeyPair keyPair = Keys.keyPairFor(SignatureAlgorithm.RS256);

        byte[] privateDer = keyPair.getPrivate().getEncoded();
        PrivateKey privateKey = KeyUtils.parsePrivateKeyFromDer(privateDer);
        assertEquals(keyPair.getPrivate(), privateKey);

        String pkcs8 = KeyUtils.toPkcs8(keyPair);
        privateKey = KeyUtils.parsePrivateKeyFromPkcs8(pkcs8);
        assertEquals(keyPair.getPrivate(), privateKey);
    }

    @Test
    public void testReadKeyPairIgnoresFilesWithLoosePermissions() throws IOException {
        Path file = tempDir.newFile("open-kp.pem");

        MoreFiles.write(file, "This isn't a real key pair");
        Files.setPosixFilePermissions(file, PosixFilePermissions.fromString("rw-r--r--")); //644

        assertNull(KeyUtils.readKeyPair(file));

        logListener.assertWarn("permissions are too loose");
    }

    @Test
    public void testReadKeyPairNonexistent() throws IOException {
        assertNull(KeyUtils.readKeyPair(tempDir.getRoot().resolve("nonexistent-kp.pem")));
    }

    @Test
    public void testReadPublicKeyIgnoresFilesWithLoosePermissions() throws IOException {
        Path file = tempDir.newFile("open-pk.pem");

        MoreFiles.write(file, "This isn't a real public key");
        Files.setPosixFilePermissions(file, PosixFilePermissions.fromString("rw-r--r--")); //644

        assertNull(KeyUtils.readPublicKey(file));

        logListener.assertWarn("permissions are too loose");
    }

    @Test
    public void testReadPublicKeyNonexistent() throws IOException {
        assertNull(KeyUtils.readPublicKey(tempDir.getRoot().resolve("nonexistent-pk.pem")));
    }

    @Test
    public void testToAndFromPkcs8ForKeyPair() throws IOException {
        for (SignatureAlgorithm algorithm : SignatureAlgorithm.values()) {
            KeyPair keyPair = generateKeyPair(algorithm);
            if (keyPair == null) {
                continue;
            }
            String pkcs8 = KeyUtils.toPkcs8(keyPair);

            KeyPair parsed = KeyUtils.parseKeyPairFromPkcs8(pkcs8);
            assertEquals("Public key mismatch for " + algorithm, keyPair.getPublic(), parsed.getPublic());
            assertEquals("Private key mismatch for " + algorithm, keyPair.getPrivate(), parsed.getPrivate());
        }
    }

    @Test
    public void testToAndFromPkcs8ForPublicKey() throws IOException {
        for (SignatureAlgorithm algorithm : SignatureAlgorithm.values()) {
            KeyPair keyPair = generateKeyPair(algorithm);
            if (keyPair == null) {
                continue;
            }
            String pkcs8 = KeyUtils.toPkcs8(keyPair.getPublic());

            PublicKey parsed = KeyUtils.parsePublicKeyFromPkcs8(pkcs8);
            assertEquals(keyPair.getPublic(), parsed);
        }
    }

    @Test
    public void testWriteReadAndSignForEcKeyPair() throws IOException {
        runWriteReadAndSignTest(SignatureAlgorithm.ES256);
    }

    @Test
    public void testWriteReadAndSignForRsaKeyPair() throws IOException {
        runWriteReadAndSignTest(SignatureAlgorithm.RS256);
    }

    @Test
    public void testWriteAndReadForEcPublicKey() throws IOException {
        runWriteAndReadPublicKeyTest(SignatureAlgorithm.ES256);
    }

    @Test
    public void testWriteAndReadForRsaPublicKey() throws IOException {
        runWriteAndReadPublicKeyTest(SignatureAlgorithm.RS256);
    }

    private static void assertPermissions(Path keyFile) throws IOException {
        Set<PosixFilePermission> permissions = Files.getPosixFilePermissions(keyFile);
        assertEquals("Key permissions were not set", "rw-------", PosixFilePermissions.toString(permissions));
    }

    private static KeyPair generateEcKeyPair(SignatureAlgorithm algorithm) {
        return EllipticCurveProvider.generateKeyPair("EC", BouncyCastleProvider.PROVIDER_NAME, algorithm, RANDOM);
    }

    private static KeyPair generateKeyPair(SignatureAlgorithm algorithm) {
        if (algorithm.isEllipticCurve()) {
            //We can't use Keys.keyPairFor here because it doesn't use BC
            return generateEcKeyPair(algorithm);
        } else if (algorithm.isRsa()) {
            return Keys.keyPairFor(algorithm);
        }

        return null;
    }

    private Path getKeyFile(SignatureAlgorithm algorithm) throws IOException {
        return tempDir.newFile(algorithm.name().toLowerCase(Locale.ROOT) + ".pem");
    }

    private void runWriteReadAndSignTest(SignatureAlgorithm algorithm) throws IOException {
        KeyPair keyPair = generateKeyPair(algorithm);
        assertNotNull(algorithm + " isn't supported, so it should be tested", keyPair);

        Path keyFile = getKeyFile(algorithm);
        KeyUtils.writeKeyPair(keyFile, keyPair);
        assertPermissions(keyFile);

        KeyPair loaded = KeyUtils.readKeyPair(keyFile);
        assertNotNull("Failed to load " + algorithm + " key", loaded);
        assertEquals("Public key mismatch for " + algorithm, keyPair.getPublic(), loaded.getPublic());
        assertEquals("Private key mismatch for " + algorithm, keyPair.getPrivate(), loaded.getPrivate());

        //As a further layer of validation, ensure interoperability between the generated and loaded keys:
        //- The generated public key can parse a token signed with the loaded private key
        //- The loaded public key can parse a token signed with the generated private key
        KeyPair[] pairs = new KeyPair[]{keyPair, loaded};
        for (int s = 0, p = 1; s <= 1 && p >= 0; ++s, --p) {
            String token = Jwts.builder()
                    .setSubject("Hello, World!")
                    .signWith(pairs[s].getPrivate(), algorithm) //Force use of the algorithm under test
                    .compact();
            Jwts.parserBuilder()
                    .requireSubject("Hello, World!")
                    .setSigningKey(pairs[p].getPublic())
                    .build()
                    .parseClaimsJws(token);
        }
    }

    private void runWriteAndReadPublicKeyTest(SignatureAlgorithm algorithm) throws IOException {
        KeyPair keyPair = generateKeyPair(algorithm);
        assertNotNull(algorithm + " isn't supported, so it should be tested", keyPair);

        Path keyFile = getKeyFile(algorithm);
        KeyUtils.writePublicKey(keyFile, keyPair.getPublic());
        assertPermissions(keyFile);

        PublicKey loaded = KeyUtils.readPublicKey(keyFile);
        assertNotNull(loaded);
        assertEquals(keyPair.getPublic(), loaded);
    }
}