/*
 * 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.nifi.security.util.crypto;

import org.apache.commons.codec.binary.Base64;
import org.apache.commons.codec.binary.Hex;
import org.apache.nifi.security.util.EncryptionMethod;
import org.apache.nifi.util.StringUtils;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;

import javax.crypto.Cipher;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.nio.charset.StandardCharsets;
import java.security.Security;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

import static org.junit.jupiter.api.Assertions.assertArrayEquals;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;

public class Argon2CipherProviderTest {
    private static final String PLAINTEXT = "ExactBlockSizeRequiredForProcess";
    private static final String BAD_PASSWORD = "thisIsABadPassword";
    private static final String SHORT_PASSWORD = "shortPassword";
    private static final int DEFAULT_KEY_LENGTH = 128;
    private final String SALT_HEX = "0123456789ABCDEFFEDCBA9876543210";
    private final List<Integer> FULL_SALT_LENGTH_RANGE = Arrays.asList(49, 50, 51, 52, 53);
    private static List<Integer> VALID_KEY_LENGTHS;

    private static List<EncryptionMethod> strongKDFEncryptionMethods;
    private RandomIVPBECipherProvider cipherProvider;


    @BeforeAll
    static void setUpOnce() throws Exception {
        Security.addProvider(new BouncyCastleProvider());

        strongKDFEncryptionMethods = Arrays.stream(EncryptionMethod.values())
                .filter(EncryptionMethod::isCompatibleWithStrongKDFs)
                .collect(Collectors.toList());

        VALID_KEY_LENGTHS = Arrays.asList(128, 192, 256);
    }

    @BeforeEach
    void setUp() {
        // Very fast parameters to test for correctness rather than production values
        cipherProvider = new Argon2CipherProvider(1024, 1, 3);
    }

    @Test
    void testGetCipherShouldBeInternallyConsistent() throws Exception {
        // Arrange
        final byte[] SALT = cipherProvider.generateSalt();

        // Act
        for (EncryptionMethod em : strongKDFEncryptionMethods) {
            // Initialize a cipher for encryption
            Cipher cipher = cipherProvider.getCipher(em, SHORT_PASSWORD, SALT, DEFAULT_KEY_LENGTH, true);
            byte[] iv = cipher.getIV();

            byte[] cipherBytes = cipher.doFinal(PLAINTEXT.getBytes("UTF-8"));

            cipher = cipherProvider.getCipher(em, SHORT_PASSWORD, SALT, iv, DEFAULT_KEY_LENGTH, false);
            byte[] recoveredBytes = cipher.doFinal(cipherBytes);
            String recovered = new String(recoveredBytes, "UTF-8");

            // Assert
            assertEquals(PLAINTEXT, recovered);
        }
    }

    @Test
    void testArgon2ShouldSupportExternalCompatibility() throws Exception {
        // Arrange

        // Default values are hashLength = 32, memory = 1024, parallelism = 1, iterations = 3, but the provided salt will contain the parameters used
        cipherProvider = new Argon2CipherProvider();

        final String PLAINTEXT = "This is a plaintext message.";
        final int hashLength = 256;

        // These values can be generated by running `$ ./openssl_argon2.rb` in the terminal
        final byte[] SALT = Hex.decodeHex("68d29a1d8021f45954333767358a2492".toCharArray());
        final byte[] IV = Hex.decodeHex("808590f35f9fba14dbda9c2bb2b76a79".toCharArray());

        final String CIPHER_TEXT = "d672412857916880c79d573aa4f9d4971b85f07438d6f62f38a0e31314caa2e5";
        byte[] cipherBytes = Hex.decodeHex(CIPHER_TEXT.toCharArray());
        EncryptionMethod encryptionMethod = EncryptionMethod.AES_CBC;

        // Sanity check
        String rubyKeyHex = "8caf581795886d38f0c605e3d674f4961c658ee3625a8e8868be36c902d234ef";
        Cipher rubyCipher = Cipher.getInstance(encryptionMethod.getAlgorithm(), "BC");
        SecretKeySpec rubyKey = new SecretKeySpec(Hex.decodeHex(rubyKeyHex.toCharArray()), "AES");
        IvParameterSpec ivSpec = new IvParameterSpec(IV);
        rubyCipher.init(Cipher.ENCRYPT_MODE, rubyKey, ivSpec);
        byte[] rubyCipherBytes = rubyCipher.doFinal(PLAINTEXT.getBytes());
        rubyCipher.init(Cipher.DECRYPT_MODE, rubyKey, ivSpec);
        assertArrayEquals(PLAINTEXT.getBytes(), rubyCipher.doFinal(rubyCipherBytes));
        assertArrayEquals(PLAINTEXT.getBytes(), rubyCipher.doFinal(cipherBytes));

        // $argon2id$v=19$m=memory,t=iterations,p=parallelism$saltB64$hashB64
        final String FULL_HASH = "$argon2id$v=19$m=256,t=3,p=1$aNKaHYAh9FlUMzdnNYokkg$jK9YF5WIbTjwxgXj1nT0lhxljuNiWo6IaL42yQLSNO8";

        final String FULL_SALT = FULL_HASH.substring(0, FULL_HASH.lastIndexOf("$"));

        final String[] hashComponents = FULL_HASH.split("\\$");
        final String saltB64 = hashComponents[4];
        byte[] salt = Base64.decodeBase64(saltB64);
        assertArrayEquals(SALT, salt);

        // Act
        Cipher cipher = cipherProvider.getCipher(encryptionMethod, BAD_PASSWORD, FULL_SALT.getBytes(), IV, hashLength, false);
        byte[] recoveredBytes = cipher.doFinal(cipherBytes);
        String recovered = new String(recoveredBytes, "UTF-8");

        // Assert
        assertEquals(PLAINTEXT, recovered);
    }

    @Test
    void testGetCipherShouldRejectInvalidIV() throws Exception {
        // Arrange
        final byte[] SALT = Hex.decodeHex(SALT_HEX.toCharArray());
        final int MAX_LENGTH = 15;
        final List<byte[]> INVALID_IVS = new ArrayList<>();
        for (int length = 0; length <= MAX_LENGTH; length++) {
            INVALID_IVS.add(new byte[length]);
        }

        EncryptionMethod encryptionMethod = EncryptionMethod.AES_CBC;

        // Act
        for (final byte[] badIV: INVALID_IVS) {
            // Encrypt should print a warning about the bad IV but overwrite it
            Cipher cipher = cipherProvider.getCipher(encryptionMethod, SHORT_PASSWORD, SALT, badIV, DEFAULT_KEY_LENGTH, true);

            // Decrypt should fail
            IllegalArgumentException iae = assertThrows(IllegalArgumentException.class,
                    () -> cipherProvider.getCipher(encryptionMethod, SHORT_PASSWORD, SALT, badIV, DEFAULT_KEY_LENGTH, false));

            // Assert
            assertTrue(iae.getMessage().contains("Cannot decrypt without a valid IV"));
        }
    }

    @Test
    void testGetCipherWithExternalIVShouldBeInternallyConsistent() throws Exception {
        // Arrange
        final byte[] SALT = cipherProvider.generateSalt();
        final byte[] IV = Hex.decodeHex(StringUtils.repeat("01", 16).toCharArray());

        // Act
        for (EncryptionMethod em : strongKDFEncryptionMethods) {
            // Initialize a cipher for encryption
            Cipher cipher = cipherProvider.getCipher(em, SHORT_PASSWORD, SALT, IV, DEFAULT_KEY_LENGTH, true);

            byte[] cipherBytes = cipher.doFinal(PLAINTEXT.getBytes("UTF-8"));

            cipher = cipherProvider.getCipher(em, SHORT_PASSWORD, SALT, IV, DEFAULT_KEY_LENGTH, false);
            byte[] recoveredBytes = cipher.doFinal(cipherBytes);
            String recovered = new String(recoveredBytes, "UTF-8");

            // Assert
            assertEquals(PLAINTEXT, recovered);
        }
    }

    @Test
    void testGetCipherWithUnlimitedStrengthShouldBeInternallyConsistent() throws Exception {
        // Arrange
        final byte[] SALT = cipherProvider.generateSalt();

        final int LONG_KEY_LENGTH = 256;

        // Act
        for (EncryptionMethod em : strongKDFEncryptionMethods) {
            // Initialize a cipher for encryption
            Cipher cipher = cipherProvider.getCipher(em, SHORT_PASSWORD, SALT, LONG_KEY_LENGTH, true);
            byte[] iv = cipher.getIV();

            byte[] cipherBytes = cipher.doFinal(PLAINTEXT.getBytes("UTF-8"));

            cipher = cipherProvider.getCipher(em, SHORT_PASSWORD, SALT, iv, LONG_KEY_LENGTH, false);
            byte[] recoveredBytes = cipher.doFinal(cipherBytes);
            String recovered = new String(recoveredBytes, "UTF-8");

            // Assert
            assertEquals(PLAINTEXT, recovered);
        }
    }

    @Test
    void testGetCipherShouldNotAcceptInvalidSalts() throws Exception {
        // Arrange
        final List<String> INVALID_SALTS = Arrays.asList("argon2", "$3a$11$", "x", "$2a$10$");
        final String LENGTH_MESSAGE = "The raw salt must be greater than or equal to 8 bytes";

        EncryptionMethod encryptionMethod = EncryptionMethod.AES_CBC;

        // Act
        for (final String salt : INVALID_SALTS) {
            IllegalArgumentException iae = assertThrows(IllegalArgumentException.class,
                    () -> cipherProvider.getCipher(encryptionMethod, BAD_PASSWORD, salt.getBytes(), DEFAULT_KEY_LENGTH, true));

            // Assert
            assertTrue(iae.getMessage().contains(LENGTH_MESSAGE));
        }
    }

    @Test
    void testGetCipherShouldHandleUnformattedSalts() throws Exception {
        // Arrange
        final List<String> RECOVERABLE_SALTS = Arrays.asList("$ab$00$acbdefghijklmnopqrstuv", "$4$1$1$0123456789abcdef", "$400$1$1$abcdefghijklmnopqrstuv");

        EncryptionMethod encryptionMethod = EncryptionMethod.AES_CBC;

        // Act
        for (final String salt : RECOVERABLE_SALTS) {
            Cipher cipher = cipherProvider.getCipher(encryptionMethod, BAD_PASSWORD, salt.getBytes(), DEFAULT_KEY_LENGTH, true);

            // Assert
            assertNotNull(cipher);
        }
    }

    @Test
    void testGetCipherShouldRejectEmptySalt() throws Exception {
        // Arrange
        EncryptionMethod encryptionMethod = EncryptionMethod.AES_CBC;

        // Act
        IllegalArgumentException iae = assertThrows(IllegalArgumentException.class,
                () -> cipherProvider.getCipher(encryptionMethod, BAD_PASSWORD, new byte[0], DEFAULT_KEY_LENGTH, true));

        // Assert
        assertTrue(iae.getMessage().contains("The salt cannot be empty. To generate a salt, use Argon2CipherProvider#generateSalt()"));
    }

    @Test
    void testGenerateSaltShouldProvideValidSalt() throws Exception {
        // Arrange
        Argon2CipherProvider cipherProvider = new Argon2CipherProvider();

        // Act
        byte[] saltBytes = cipherProvider.generateSalt();

        String fullSalt = new String(saltBytes, StandardCharsets.UTF_8);

        final Matcher matcher = Pattern.compile("\\$([\\w\\+\\/]+)\\$?$").matcher(fullSalt);
        matcher.find();
        final String rawSaltB64 = matcher.group(1);

        byte[] rawSaltBytes = Base64.decodeBase64(rawSaltB64);

        // Assert
        boolean isValidFormattedSalt = cipherProvider.isArgon2FormattedSalt(fullSalt);
        assertTrue(isValidFormattedSalt);

        boolean fullSaltIsValidLength = FULL_SALT_LENGTH_RANGE.contains(saltBytes.length);
        assertTrue(fullSaltIsValidLength);

        byte[] notExpected = new byte[16];
        Arrays.fill(notExpected, (byte) '\0');
        assertFalse(Arrays.equals(notExpected, rawSaltBytes));
    }

    @Test
    void testGetCipherForDecryptShouldRequireIV() throws Exception {
        // Arrange
        final byte[] SALT = cipherProvider.generateSalt();
        final byte[] IV = Hex.decodeHex(StringUtils.repeat("00", 16).toCharArray());

        // Act
        for (EncryptionMethod em : strongKDFEncryptionMethods) {
            // Initialize a cipher for encryption
            Cipher cipher = cipherProvider.getCipher(em, SHORT_PASSWORD, SALT, IV, DEFAULT_KEY_LENGTH, true);

            byte[] cipherBytes = cipher.doFinal(PLAINTEXT.getBytes("UTF-8"));

            IllegalArgumentException iae = assertThrows(IllegalArgumentException.class,
                    () -> cipherProvider.getCipher(em, SHORT_PASSWORD, SALT, DEFAULT_KEY_LENGTH, false));

            // Assert
            assertTrue(iae.getMessage().contains("Cannot decrypt without a valid IV"));
        }
    }

    @Test
    void testGetCipherShouldAcceptValidKeyLengths() throws Exception {
        // Arrange
        final byte[] SALT = cipherProvider.generateSalt();
        final byte[] IV = Hex.decodeHex(StringUtils.repeat("01", 16).toCharArray());

        EncryptionMethod encryptionMethod = EncryptionMethod.AES_CBC;

        // Act
        for (final int keyLength : VALID_KEY_LENGTHS) {
            // Initialize a cipher for encryption
            Cipher cipher = cipherProvider.getCipher(encryptionMethod, SHORT_PASSWORD, SALT, IV, keyLength, true);

            byte[] cipherBytes = cipher.doFinal(PLAINTEXT.getBytes("UTF-8"));

            cipher = cipherProvider.getCipher(encryptionMethod, SHORT_PASSWORD, SALT, IV, keyLength, false);
            byte[] recoveredBytes = cipher.doFinal(cipherBytes);
            String recovered = new String(recoveredBytes, "UTF-8");

            // Assert
            assertEquals(PLAINTEXT, recovered);
        }
    }

    @Test
    void testGetCipherShouldNotAcceptInvalidKeyLengths() throws Exception {
        // Arrange
        final byte[] SALT = cipherProvider.generateSalt();
        final byte[] IV = Hex.decodeHex(StringUtils.repeat("00", 16).toCharArray());

        final List<Integer> INVALID_KEY_LENGTHS = Arrays.asList(-1, 40, 64, 112, 512);
        EncryptionMethod encryptionMethod = EncryptionMethod.AES_CBC;

        // Act
        for (final int keyLength : INVALID_KEY_LENGTHS) {
            // Initialize a cipher for
            IllegalArgumentException iae = assertThrows(IllegalArgumentException.class,
                    () -> cipherProvider.getCipher(encryptionMethod, SHORT_PASSWORD, SALT, IV, keyLength, true));

            // Assert
            assertTrue(iae.getMessage().contains(keyLength + " is not a valid key length for AES"));
        }
    }

    @Test
    void testArgon2ShouldNotAcceptInvalidPassword() {
        // Arrange
        String badPassword = "";
        byte[] salt = new byte[16];
        Arrays.fill(salt, (byte) 0x01);

        EncryptionMethod encryptionMethod = EncryptionMethod.AES_CBC;

        // Act
        IllegalArgumentException iae = assertThrows(IllegalArgumentException.class,
                () -> cipherProvider.getCipher(encryptionMethod, badPassword, salt, DEFAULT_KEY_LENGTH, true));

        // Assert
        assertTrue(iae.getMessage().contains("Encryption with an empty password is not supported"));
    }

    @Test
    void testShouldRejectInvalidSalt() throws Exception {
        // Arrange
        Argon2CipherProvider cipherProvider = new Argon2CipherProvider();

        final String FULL_HASH = "$argon2id$v=19$m=1024,t=4,p=1$hiKyaQbZyQBmCmD1zGcyMw$rc+ec+/hQeBcwzjH+OEmUtaTUqhZYKN4ZKJtWzFZYjQ";

        // Act
        boolean isValid = cipherProvider.isArgon2FormattedSalt(FULL_HASH);

        // Assert
        assertFalse(isValid);
    }

    @Test
    void testShouldExtractSalt() throws Exception {
        // Arrange
        Argon2CipherProvider cipherProvider = new Argon2CipherProvider();

        final byte[] EXPECTED_RAW_SALT = Hex.decodeHex("8622b26906d9c900660a60f5cc673233".toCharArray());

        final String FORMATTED_SALT = "$argon2id$v=19$m=1024,t=4,p=1$hiKyaQbZyQBmCmD1zGcyMw";

        byte[] rawSalt;

        // Act
        rawSalt = cipherProvider.extractRawSaltFromArgon2Salt(FORMATTED_SALT);

        // Assert
        assertArrayEquals(EXPECTED_RAW_SALT, rawSalt);
    }
}
