/*
 * Copyright 2012 Google Inc. All Rights Reserved.
 */

#include "media/stagefright/foundation/ABase.h"
#include "utils/UniquePtr.h"
#include "WVCryptoPlugin.h"
#include "OEMCryptoDASH.h"
#include "gtest/gtest.h"

using namespace wvclearkey;

using android::status_t;

bool oemCryptoSessionOpened = false;
OEMCrypto_SESSION openedCryptoSession = 0;
bool oemCryptoSessionClosed = false;
OEMCrypto_SESSION closedCryptoSession = 0;

extern "C" {
    OEMCryptoResult OEMCrypto_OpenSession(OEMCrypto_SESSION *session) {
        oemCryptoSessionOpened = true;
        openedCryptoSession = 5;
        *session = openedCryptoSession;
        return OEMCrypto_SUCCESS;
    }

    OEMCryptoResult OEMCrypto_CloseSession(OEMCrypto_SESSION session) {
        oemCryptoSessionClosed = true;
        closedCryptoSession = session;
        return OEMCrypto_SUCCESS;
    }

}

TEST(WVCryptoPluginTest, ManagesASession) {
    oemCryptoSessionOpened = false;
    openedCryptoSession = 0;
    oemCryptoSessionClosed = false;
    closedCryptoSession = 0;

    UniquePtr<WVCryptoPlugin> plugin(new WVCryptoPlugin());

    EXPECT_TRUE(oemCryptoSessionOpened) <<
        "WVCryptoPlugin did not call OEMCrypto_OpenSession()";

    plugin.reset();

    EXPECT_TRUE(oemCryptoSessionClosed) <<
        "WVCryptoPlugin did not call OEMCrypto_CloseSession()";
    EXPECT_EQ(openedCryptoSession, closedCryptoSession) <<
        "WVCryptoPlugin closed the wrong session";
}

TEST(WVCryptoPluginTest, CorrectlyReportsSecureBuffers) {
    UniquePtr<WVCryptoPlugin> plugin(new WVCryptoPlugin());

    EXPECT_FALSE(plugin->requiresSecureDecoderComponent("video/mp4")) <<
        "WVCryptoPlugin incorrectly expects a secure video decoder";
    EXPECT_FALSE(plugin->requiresSecureDecoderComponent("audio/aac")) <<
        "WVCryptoPlugin incorrectly expects a secure audio decoder";
}

typedef const uint8_t* constChars;

bool selectKeyCalled;
constChars selectedKey;
size_t selectedKeyLength;

uint32_t decryptCallCount;
constChars decryptInPointers[3];
size_t decryptInLengths[3];
bool decryptIsEncryptedValues[3];
constChars decryptIvs[3];
size_t decryptOffsets[3];
OEMCrypto_DestBufferDesc decryptOutBuffers[3];

extern "C" {
    OEMCryptoResult OEMCrypto_SelectKey(
            const OEMCrypto_SESSION session,
            const uint8_t* key_id,
            const size_t key_id_length) {
        selectKeyCalled = true;
        selectedKey = key_id;
        selectedKeyLength = key_id_length;
        return OEMCrypto_SUCCESS;
    }

    OEMCryptoResult OEMCrypto_DecryptCTR(
            OEMCrypto_SESSION session,
            const uint8_t *data_addr,
            size_t data_length,
            bool is_encrypted,
            const uint8_t *iv,
            size_t offset,
            const OEMCrypto_DestBufferDesc* out_buffer) {
        if (decryptCallCount < 3) {
            decryptInPointers[decryptCallCount] = data_addr;
            decryptInLengths[decryptCallCount] = data_length;
            decryptIsEncryptedValues[decryptCallCount] = is_encrypted;
            decryptIvs[decryptCallCount] = iv;
            decryptOffsets[decryptCallCount] = offset;
            decryptOutBuffers[decryptCallCount] = *out_buffer;
        }

        decryptCallCount++;
        return OEMCrypto_SUCCESS;
    }
}

TEST(WVCryptoPluginTest, AttemptsToDecrypt) {
    selectKeyCalled = false;
    selectedKey = NULL;
    selectedKeyLength = 0;
    decryptCallCount = 0;
    memset(decryptInPointers, 0, sizeof(decryptInPointers));
    memset(decryptInLengths, 0, sizeof(decryptInLengths));
    memset(decryptIsEncryptedValues, 0, sizeof(decryptIsEncryptedValues));
    memset(decryptIvs, 0, sizeof(decryptIvs));
    memset(decryptOffsets, 0, sizeof(decryptOffsets));
    memset(decryptOutBuffers, 0, sizeof(decryptOutBuffers));

    UniquePtr<WVCryptoPlugin> plugin(new WVCryptoPlugin());

    android::CryptoPlugin::SubSample subSamples[3];
    subSamples[0].mNumBytesOfEncryptedData = 16;
    subSamples[1].mNumBytesOfEncryptedData = 24;
    subSamples[2].mNumBytesOfEncryptedData = 8;

    uint8_t key[16] = {};
    uint8_t iv[16] = {};

    uint8_t in[48] = {};
    uint8_t out[48] = {};

    ssize_t decrypted = plugin->decrypt(
            false,
            key,
            iv,
            android::CryptoPlugin::kMode_AES_CTR,
            in,
            subSamples, 3,
            out,
            NULL);

    EXPECT_EQ(48, decrypted) <<
        "WVCryptoPlugin decrypted the wrong number of bytes";
    ASSERT_EQ(3u, decryptCallCount) <<
        "WVCryptoPlugin called OEMCrypto_DecryptCTR the wrong number of times";

    // Test the 1st call
    EXPECT_EQ(in, decryptInPointers[0]) <<
        "1st OEMCrypto_DecryptCTR call targetted the wrong input location";
    EXPECT_EQ(16u, decryptInLengths[0]) <<
        "1st OEMCrypto_DecryptCTR call targetted the wrong input length";
    EXPECT_TRUE(decryptIsEncryptedValues[0]) <<
        "1st OEMCrypto_DecryptCTR call thought data was unencrypted";
    EXPECT_EQ(iv, decryptIvs[0]) <<
        "1st OEMCrypto_DecryptCTR call had the wrong iv";
    EXPECT_EQ(0u, decryptOffsets[0]) <<
        "1st OEMCrypto_DecryptCTR call had the wrong offset";
    EXPECT_EQ(out, decryptOutBuffers[0].buffer.clear.address) <<
        "1st OEMCrypto_DecryptCTR call targetted the wrong output location";

    // Test the 2nd call
    EXPECT_EQ(in + 16, decryptInPointers[0]) <<
        "2nd OEMCrypto_DecryptCTR call targetted the wrong input location";
    EXPECT_EQ(24u, decryptInLengths[0]) <<
        "2nd OEMCrypto_DecryptCTR call targetted the wrong input length";
    EXPECT_TRUE(decryptIsEncryptedValues[0]) <<
        "2nd OEMCrypto_DecryptCTR call thought data was unencrypted";
    EXPECT_EQ(iv, decryptIvs[0]) <<
        "2nd OEMCrypto_DecryptCTR call had the wrong iv";
    EXPECT_EQ(0u, decryptOffsets[0]) <<
        "2nd OEMCrypto_DecryptCTR call had the wrong offset";
    EXPECT_EQ(out + 16, decryptOutBuffers[0].buffer.clear.address) <<
        "2nd OEMCrypto_DecryptCTR call targetted the wrong output location";

    // Test the 3rd call
    EXPECT_EQ(in + 40, decryptInPointers[0]) <<
        "3rd OEMCrypto_DecryptCTR call targetted the wrong input location";
    EXPECT_EQ(8u, decryptInLengths[0]) <<
        "3rd OEMCrypto_DecryptCTR call targetted the wrong input length";
    EXPECT_TRUE(decryptIsEncryptedValues[0]) <<
        "3rd OEMCrypto_DecryptCTR call thought data was unencrypted";
    EXPECT_EQ(iv, decryptIvs[0]) <<
        "3rd OEMCrypto_DecryptCTR call had the wrong iv";
    EXPECT_EQ(8u, decryptOffsets[0]) <<
        "3rd OEMCrypto_DecryptCTR call had the wrong offset";
    EXPECT_EQ(out + 40, decryptOutBuffers[0].buffer.clear.address) <<
        "3rd OEMCrypto_DecryptCTR call targetted the wrong output location";
}