/*
 * Copyright (C) 2019 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.tasks;

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.doThrow;
import static org.mockito.Mockito.when;
import static org.testng.Assert.assertThrows;

import android.os.ParcelFileDescriptor;

import com.android.server.backup.encryption.chunk.ChunkHash;
import com.android.server.backup.encryption.chunking.ChunkHasher;
import com.android.server.backup.testing.CryptoTestUtils;
import com.android.server.testing.shadows.DataEntity;
import com.android.server.testing.shadows.ShadowBackupDataOutput;

import com.google.protobuf.nano.MessageNano;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.util.Arrays;
import java.util.HashSet;
import java.util.Set;

import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;
import org.junit.runner.RunWith;
import org.mockito.Mock;
import org.mockito.MockitoAnnotations;
import org.robolectric.RobolectricTestRunner;
import org.robolectric.annotation.Config;

@Config(shadows = {ShadowBackupDataOutput.class})
@RunWith(RobolectricTestRunner.class)
public class EncryptedKvRestoreTaskTest {
    private static final String TEST_KEY_1 = "test_key_1";
    private static final String TEST_KEY_2 = "test_key_2";
    private static final String TEST_KEY_3 = "test_key_3";
    private static final byte[] TEST_VALUE_1 = {1, 2, 3};
    private static final byte[] TEST_VALUE_2 = {4, 5, 6};
    private static final byte[] TEST_VALUE_3 = {20, 25, 30, 35};

    @Rule public TemporaryFolder temporaryFolder = new TemporaryFolder();

    private File temporaryDirectory;

    @Mock private ParcelFileDescriptor mParcelFileDescriptor;
    @Mock private ChunkHasher mChunkHasher;
    @Mock private FullRestoreToFileTask mFullRestoreToFileTask;
    @Mock private BackupFileDecryptorTask mBackupFileDecryptorTask;

    private EncryptedKvRestoreTask task;

    @Before
    public void setUp() throws Exception {
        MockitoAnnotations.initMocks(this);

        when(mChunkHasher.computeHash(any()))
                .thenAnswer(invocation -> fakeHash(invocation.getArgument(0)));
        doAnswer(invocation -> writeTestPairsToFile(invocation.getArgument(0)))
                .when(mFullRestoreToFileTask)
                .restoreToFile(any());
        doAnswer(
                        invocation ->
                                readPairsFromFile(
                                        invocation.getArgument(0), invocation.getArgument(1)))
                .when(mBackupFileDecryptorTask)
                .decryptFile(any(), any());

        temporaryDirectory = temporaryFolder.newFolder();
        task =
                new EncryptedKvRestoreTask(
                        temporaryDirectory,
                        mChunkHasher,
                        mFullRestoreToFileTask,
                        mBackupFileDecryptorTask);
    }

    @Test
    public void testGetRestoreData_writesPairsToOutputInOrder() throws Exception {
        task.getRestoreData(mParcelFileDescriptor);

        assertThat(ShadowBackupDataOutput.getEntities())
                .containsExactly(
                        new DataEntity(TEST_KEY_1, TEST_VALUE_1),
                        new DataEntity(TEST_KEY_2, TEST_VALUE_2),
                        new DataEntity(TEST_KEY_3, TEST_VALUE_3))
                .inOrder();
    }

    @Test
    public void testGetRestoreData_exceptionDuringDecryption_throws() throws Exception {
        doThrow(IOException.class).when(mBackupFileDecryptorTask).decryptFile(any(), any());
        assertThrows(IOException.class, () -> task.getRestoreData(mParcelFileDescriptor));
    }

    @Test
    public void testGetRestoreData_exceptionDuringDownload_throws() throws Exception {
        doThrow(IOException.class).when(mFullRestoreToFileTask).restoreToFile(any());
        assertThrows(IOException.class, () -> task.getRestoreData(mParcelFileDescriptor));
    }

    @Test
    public void testGetRestoreData_exceptionDuringDecryption_deletesTemporaryFiles() throws Exception {
        doThrow(InvalidKeyException.class).when(mBackupFileDecryptorTask).decryptFile(any(), any());
        assertThrows(InvalidKeyException.class, () -> task.getRestoreData(mParcelFileDescriptor));
        assertThat(temporaryDirectory.listFiles()).isEmpty();
    }

    @Test
    public void testGetRestoreData_exceptionDuringDownload_deletesTemporaryFiles() throws Exception {
        doThrow(IOException.class).when(mFullRestoreToFileTask).restoreToFile(any());
        assertThrows(IOException.class, () -> task.getRestoreData(mParcelFileDescriptor));
        assertThat(temporaryDirectory.listFiles()).isEmpty();
    }

    private static Void writeTestPairsToFile(File file) throws IOException {
        // Write the pairs out of order to check the task sorts them.
        Set<byte[]> pairs =
                new HashSet<>(
                        Arrays.asList(
                                createPair(TEST_KEY_1, TEST_VALUE_1),
                                createPair(TEST_KEY_3, TEST_VALUE_3),
                                createPair(TEST_KEY_2, TEST_VALUE_2)));

        try (ObjectOutputStream oos = new ObjectOutputStream(new FileOutputStream(file))) {
            oos.writeObject(pairs);
        }
        return null;
    }

    private static Void readPairsFromFile(File file, DecryptedChunkOutput decryptedChunkOutput)
            throws IOException, ClassNotFoundException, InvalidKeyException,
                    NoSuchAlgorithmException {
        try (ObjectInputStream ois = new ObjectInputStream(new FileInputStream(file));
                DecryptedChunkOutput output = decryptedChunkOutput.open()) {
            Set<byte[]> pairs = readPairs(ois);
            for (byte[] pair : pairs) {
                output.processChunk(pair, pair.length);
            }
        }

        return null;
    }

    private static byte[] createPair(String key, byte[] value) {
        return MessageNano.toByteArray(CryptoTestUtils.newPair(key, value));
    }

    @SuppressWarnings("unchecked") // deserialization.
    private static Set<byte[]> readPairs(ObjectInputStream ois)
            throws IOException, ClassNotFoundException {
        return (Set<byte[]>) ois.readObject();
    }

    private static ChunkHash fakeHash(byte[] data) {
        return new ChunkHash(Arrays.copyOf(data, ChunkHash.HASH_LENGTH_BYTES));
    }
}