/*
 * Copyright (C) 2018 The Android Open Source Project
 *
 * Licensed 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 com.android.server.backup.encryption.chunking;

import static com.android.server.backup.testing.CryptoTestUtils.generateAesKey;

import static com.google.common.truth.Truth.assertThat;

import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.mock;

import static java.nio.charset.StandardCharsets.UTF_8;

import android.platform.test.annotations.Presubmit;

import com.android.server.backup.encryption.chunk.ChunkHash;

import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.Mock;
import org.mockito.invocation.InvocationOnMock;
import org.mockito.stubbing.Answer;
import org.robolectric.RobolectricTestRunner;

import java.security.SecureRandom;

import javax.crypto.Cipher;
import javax.crypto.Mac;
import javax.crypto.SecretKey;
import javax.crypto.spec.GCMParameterSpec;

@RunWith(RobolectricTestRunner.class)
@Presubmit
public class ChunkEncryptorTest {
    private static final String MAC_ALGORITHM = "HmacSHA256";
    private static final String CIPHER_ALGORITHM = "AES/GCM/NoPadding";
    private static final int GCM_NONCE_LENGTH_BYTES = 12;
    private static final int GCM_TAG_LENGTH_BYTES = 16;
    private static final String CHUNK_PLAINTEXT =
            "A little Learning is a dang'rous Thing;\n"
                    + "Drink deep, or taste not the Pierian Spring:\n"
                    + "There shallow Draughts intoxicate the Brain,\n"
                    + "And drinking largely sobers us again.";
    private static final byte[] PLAINTEXT_BYTES = CHUNK_PLAINTEXT.getBytes(UTF_8);
    private static final byte[] NONCE_1 = "0123456789abc".getBytes(UTF_8);
    private static final byte[] NONCE_2 = "123456789abcd".getBytes(UTF_8);

    private static final byte[][] NONCES = new byte[][] {NONCE_1, NONCE_2};

    @Mock private SecureRandom mSecureRandomMock;
    private SecretKey mSecretKey;
    private ChunkHash mPlaintextHash;
    private ChunkEncryptor mChunkEncryptor;

    @Before
    public void setUp() throws Exception {
        mSecretKey = generateAesKey();
        ChunkHasher chunkHasher = new ChunkHasher(mSecretKey);
        mPlaintextHash = chunkHasher.computeHash(PLAINTEXT_BYTES);
        mSecureRandomMock = mock(SecureRandom.class);
        mChunkEncryptor = new ChunkEncryptor(mSecretKey, mSecureRandomMock);

        // Return NONCE_1, then NONCE_2 for invocations of mSecureRandomMock.nextBytes().
        doAnswer(
                        new Answer<Void>() {
                            private int mInvocation = 0;

                            @Override
                            public Void answer(InvocationOnMock invocation) {
                                byte[] nonceDestination = invocation.getArgument(0);
                                System.arraycopy(
                                        NONCES[this.mInvocation],
                                        0,
                                        nonceDestination,
                                        0,
                                        GCM_NONCE_LENGTH_BYTES);
                                this.mInvocation++;
                                return null;
                            }
                        })
                .when(mSecureRandomMock)
                .nextBytes(any(byte[].class));
    }

    @Test
    public void encrypt_withHash_resultContainsHashAsKey() throws Exception {
        EncryptedChunk chunk = mChunkEncryptor.encrypt(mPlaintextHash, PLAINTEXT_BYTES);

        assertThat(chunk.key()).isEqualTo(mPlaintextHash);
    }

    @Test
    public void encrypt_generatesHmacOfPlaintext() throws Exception {
        EncryptedChunk chunk = mChunkEncryptor.encrypt(mPlaintextHash, PLAINTEXT_BYTES);

        byte[] generatedHash = chunk.key().getHash();
        Mac mac = Mac.getInstance(MAC_ALGORITHM);
        mac.init(mSecretKey);
        byte[] plaintextHmac = mac.doFinal(PLAINTEXT_BYTES);
        assertThat(generatedHash).isEqualTo(plaintextHmac);
    }

    @Test
    public void encrypt_whenInvokedAgain_generatesNewNonce() throws Exception {
        EncryptedChunk chunk1 = mChunkEncryptor.encrypt(mPlaintextHash, PLAINTEXT_BYTES);

        EncryptedChunk chunk2 = mChunkEncryptor.encrypt(mPlaintextHash, PLAINTEXT_BYTES);

        assertThat(chunk1.nonce()).isNotEqualTo(chunk2.nonce());
    }

    @Test
    public void encrypt_whenInvokedAgain_generatesNewCiphertext() throws Exception {
        EncryptedChunk chunk1 = mChunkEncryptor.encrypt(mPlaintextHash, PLAINTEXT_BYTES);

        EncryptedChunk chunk2 = mChunkEncryptor.encrypt(mPlaintextHash, PLAINTEXT_BYTES);

        assertThat(chunk1.encryptedBytes()).isNotEqualTo(chunk2.encryptedBytes());
    }

    @Test
    public void encrypt_generates12ByteNonce() throws Exception {
        EncryptedChunk encryptedChunk = mChunkEncryptor.encrypt(mPlaintextHash, PLAINTEXT_BYTES);

        byte[] nonce = encryptedChunk.nonce();
        assertThat(nonce).hasLength(GCM_NONCE_LENGTH_BYTES);
    }

    @Test
    public void encrypt_decryptedResultCorrespondsToPlaintext() throws Exception {
        EncryptedChunk chunk = mChunkEncryptor.encrypt(mPlaintextHash, PLAINTEXT_BYTES);

        Cipher cipher = Cipher.getInstance(CIPHER_ALGORITHM);
        cipher.init(
                Cipher.DECRYPT_MODE,
                mSecretKey,
                new GCMParameterSpec(GCM_TAG_LENGTH_BYTES * 8, chunk.nonce()));
        byte[] decrypted = cipher.doFinal(chunk.encryptedBytes());
        assertThat(decrypted).isEqualTo(PLAINTEXT_BYTES);
    }
}