nifi-commons/nifi-security-utils/src/test/groovy/org/apache/nifi/security/repository/AbstractAESEncryptorTest.groovy

/*
 * 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.repository

import org.apache.nifi.security.kms.EncryptionException
import org.apache.nifi.security.util.EncryptionMethod
import org.bouncycastle.jce.provider.BouncyCastleProvider
import org.bouncycastle.util.encoders.Hex
import org.junit.After
import org.junit.AfterClass
import org.junit.Before
import org.junit.BeforeClass
import org.junit.Test
import org.junit.runner.RunWith
import org.junit.runners.JUnit4
import org.slf4j.Logger
import org.slf4j.LoggerFactory

import javax.crypto.Cipher
import java.security.Security

@RunWith(JUnit4.class)
class AbstractAESEncryptorTest extends GroovyTestCase {
    private static final Logger logger = LoggerFactory.getLogger(AbstractAESEncryptor.class)

    private static final String KEY_HEX_128 = "0123456789ABCDEFFEDCBA9876543210"
    private static final String KEY_HEX_256 = KEY_HEX_128 * 2
    private static final String KEY_HEX = isUnlimitedStrengthCryptoAvailable() ? KEY_HEX_256 : KEY_HEX_128

    private static final String KEY_ID = "K1"

    private static final String LOG_PACKAGE = "org.slf4j.simpleLogger.log.org.apache.nifi.security.repository"
    private static String ORIGINAL_LOG_LEVEL

    @BeforeClass
    static void setUpOnce() throws Exception {
        ORIGINAL_LOG_LEVEL = System.getProperty(LOG_PACKAGE)
        System.setProperty(LOG_PACKAGE, "DEBUG")

        Security.addProvider(new BouncyCastleProvider())

        logger.metaClass.methodMissing = { String name, args ->
            logger.info("[${name?.toUpperCase()}] ${(args as List).join(" ")}")
        }
    }

    @Before
    void setUp() throws Exception {

    }

    @After
    void tearDown() throws Exception {

    }

    @AfterClass
    static void tearDownOnce() throws Exception {
        if (ORIGINAL_LOG_LEVEL) {
            System.setProperty(LOG_PACKAGE, ORIGINAL_LOG_LEVEL)
        }
    }

    private static boolean isUnlimitedStrengthCryptoAvailable() {
        Cipher.getMaxAllowedKeyLength("AES") > 128
    }

    /**
     * This serialized input was generated by code available in commit 094fea6d1c1f798c4e54cfaf0998416e6c59e41a.
     *
     * @return a valid {@link InputStream} containing a {@link RepositoryObjectEncryptionMetadata} and ciphertext
     */
    private static InputStream formCiphertextStream() {
        byte[] encryptedBytes = Hex.decode("0000aced00057372003f6f72672e6170616368652e6e6966692e73656375726974792e7265706" +
                "f7369746f72792e53747265616d696e67456e6372797074696f6e4d6574616461746118466982894d442c020000787200466f726" +
                "72e6170616368652e6e6966692e73656375726974792e7265706f7369746f72792e5265706f7369746f72794f626a656374456e6" +
                "372797074696f6e4d657461646174619f4328584edfdf08020005490010636970686572427974654c656e6774684c0009616c676" +
                "f726974686d7400124c6a6176612f6c616e672f537472696e673b5b0007697642797465737400025b424c00056b6579496471007" +
                "e00024c000776657273696f6e71007e00027870ffffffff7400114145532f4354522f4e6f50616464696e67757200025b42acf31" +
                "7f8060854e00200007870000000109a796446562404a917b9b06479be0f2f7400024b3174000276312356e626790d1b188345a3f" +
                "4b5e52cfa4641ed18647caec833ff6a26")
        new ByteArrayInputStream(encryptedBytes)
    }

    /**
     * This serialized input was generated by code available in commit 094fea6d1c1f798c4e54cfaf0998416e6c59e41a.
     *
     * @return a valid {@code byte[]} containing a {@link RepositoryObjectEncryptionMetadata} and ciphertext
     */
    private static byte[] formCiphertextBytes() {
        byte[] encryptedBytes = Hex.decode("aced0005737200416f72672e6170616368652e6e6966692e73656375726974792e7265706f736" +
                "9746f72792e626c6f636b2e426c6f636b456e6372797074696f6e4d6574616461746136c69c49d597a81f020000787200466f726" +
                "72e6170616368652e6e6966692e73656375726974792e7265706f7369746f72792e5265706f7369746f72794f626a656374456e6" +
                "372797074696f6e4d657461646174619f4328584edfdf08020005490010636970686572427974654c656e6774684c0009616c676" +
                "f726974686d7400124c6a6176612f6c616e672f537472696e673b5b0007697642797465737400025b424c00056b6579496471007" +
                "e00024c000776657273696f6e71007e000278700000002c7400114145532f47434d2f4e6f50616464696e67757200025b42acf31" +
                "7f8060854e002000078700000001054d7d66e359a194854af6d8211def7a47400024b31740002763145c7cfddb413ad677c5e5e4" +
                "ba59993db96df90cfff386a62bd9db094c62f752386017d28e267a3eb903090ca")
        encryptedBytes
    }

    @Test
    void testShouldPrepareInputStreamForDecryption() {
        // Arrange
        InputStream ciphertextStream = formCiphertextStream()

        // Act
        RepositoryObjectEncryptionMetadata metadata = AbstractAESEncryptor.prepareObjectForDecryption(ciphertextStream, "S1", "streaming repository object", ["v1"])
        logger.info("Extracted ROEM: ${metadata}")

        // Assert
        assert metadata.keyId == KEY_ID
        assert metadata.algorithm == EncryptionMethod.AES_CTR.algorithm
        assert metadata.version == "v1"
    }

    @Test
    void testShouldPrepareByteArrayForDecryption() {
        // Arrange
        byte[] ciphertextBytes = formCiphertextBytes()

        // Act
        RepositoryObjectEncryptionMetadata metadata = AbstractAESEncryptor.prepareObjectForDecryption(ciphertextBytes, "R1", "block repository object", ["v1"])
        logger.info("Extracted ROEM: ${metadata}")

        // Assert
        assert metadata.keyId == KEY_ID
        assert metadata.algorithm == EncryptionMethod.AES_GCM.algorithm
        assert metadata.version == "v1"
    }

    @Test
    void testPrepareObjectForDecryptionShouldFailOnUnsupportedSourceType() {
        // Arrange
        String ciphertextString = "This is not a valid ciphertext source"

        // Act
        def msg = shouldFail(EncryptionException) {
            RepositoryObjectEncryptionMetadata metadata = AbstractAESEncryptor.prepareObjectForDecryption(ciphertextString, "X1", "unsupported repository object", ["v1"])
            logger.info("Extracted ROEM: ${metadata}")
        }

        // Assert
        assert msg == "The unsupported repository object with ID X1 was detected as String; this is not a supported source of ciphertext"
    }
}