/*---------------------------------------------------------------------------*
  Project:  Horizon
  File:     main.cpp

  Copyright (C)2009-2012 Nintendo Co., Ltd.  All rights reserved.

  These coded instructions, statements, and computer programs contain
  proprietary information of Nintendo of America Inc. and/or Nintendo
  Company Ltd., and are protected by Federal copyright law.  They may
  not be disclosed to third parties or copied or duplicated in any form,
  in whole or in part, without the prior written consent of Nintendo.

  $Rev:$
 *---------------------------------------------------------------------------*/

#include <memory>
#include <nn.h>


namespace
{
    // Data targeted for encryption
    const char TEST_DATA[] =
        "abcdefghijklmnopqrstuvwxyz ABCDEFGHIJKLMNOPQRSTUVWXYZ";

    // Key used for encryption
    const bit8 TEST_KEY[nn::crypto::Aes128::KEY_SIZE] =
        {
            0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77,
            0x88, 0x99, 0xaa, 0xbb, 0xcc, 0xdd, 0xee, 0xff
        };


    //-----------------------------------------------------------------------------

    void PrintBin(const void* pv, size_t size, const char* indent)
    {
        const bit8* p = reinterpret_cast<const bit8*>(pv);

        for( size_t i = 0; i < size; ++i )
        {
            if( (i % 16) == 0 )
            {
                NN_LOG("%s%s", ((i > 0) ? "\n": ""), indent);
            }

            NN_LOG("%02x ", p[i]);
        }

        NN_LOG("\n");
    }

    //-----------------------------------------------------------------------------

    //-------------------------------------------------------------------------------
    // 1. Perform everything with a single function
    //
    class Type1
    {
    public:
        size_t GetCipherSize1( size_t plainSize )
        {
            return plainSize + nn::crypto::ENCRYPT_HEADER_SIZE;
        }
        size_t GetCipherSize2( size_t plainSize )
        {
            return plainSize + nn::crypto::GENERATE_HEADER_SIZE;
        }
        void Encrypt(void* pDst, size_t dstSize, const void* pSrc, size_t srcSize, const void* pKey, size_t keySize)
        {
            // dstSize and keySize are not referenced
            NN_UNUSED_VAR(dstSize);
            NN_UNUSED_VAR(keySize);
            NN_EQUAL_ASSERT(dstSize, GetCipherSize1(srcSize));

            // Perform everything with a single function
            nn::crypto::EncryptAes128Ctr(pDst, pSrc, srcSize, pKey);
        }
        void Decrypt(void* pDst, size_t dstSize, const void* pSrc, size_t srcSize, const void* pKey, size_t keySize)
        {
            // dstSize and keySize are not referenced
            NN_UNUSED_VAR(dstSize);
            NN_UNUSED_VAR(keySize);
            NN_EQUAL_ASSERT(GetCipherSize1(dstSize), srcSize);

            // Perform everything with a single function
            nn::crypto::DecryptAes128Ctr(pDst, pSrc, srcSize, pKey);
        }
        void EncryptAndGenerate(void* pDst, size_t dstSize, const void* pSrc, size_t srcSize, const void* pKey, size_t keySize)
        {
            // dstSize and keySize are not referenced
            NN_UNUSED_VAR(dstSize);
            NN_UNUSED_VAR(keySize);
            NN_EQUAL_ASSERT(dstSize, GetCipherSize2(srcSize));

            // Perform everything with a single function
            nn::crypto::EncryptAndGenerateAes128Ccm(pDst, pSrc, srcSize, pKey);
        }
        bool DecryptAndVerify(void* pDst, size_t dstSize, const void* pSrc, size_t srcSize, const void* pKey, size_t keySize)
        {
            // dstSize and keySize are not referenced
            NN_UNUSED_VAR(dstSize);
            NN_UNUSED_VAR(keySize);
            NN_EQUAL_ASSERT(GetCipherSize2(dstSize), srcSize);

            // Perform everything with a single function
            return nn::crypto::DecryptAndVerifyAes128Ccm(pDst, pSrc, srcSize, pKey);
        }
    };

    //-------------------------------------------------------------------------------
    // 2. Perform a batch process combining block cipher and encryption use mode
    //
    class Type2
    {
    private:
        static const size_t CCM_IV_SIZE = 12;
        NN_COMPILER_ASSERT( CCM_IV_SIZE <= nn::crypto::CcmEncryptor::IV_SIZE );

    private:
        bit8    m_Iv[nn::crypto::CtrEncryptor128::IV_SIZE];
        bit8    m_Mac[nn::crypto::CcmEncryptor::MAC_SIZE];

    public:
        size_t GetCipherSize1( size_t plainSize ) { return plainSize; }
        size_t GetCipherSize2( size_t plainSize ) { return plainSize; }

        void Encrypt(void* pDst, size_t dstSize, const void* pSrc, size_t srcSize, const void* pKey, size_t keySize)
        {
            // Randomly generate IV.
            // IV must use a different value for each encryption
            // Furthermore, it must be saved separately because it is needed during decryption.
            nn::crypto::GenerateRandomBytes(m_Iv, sizeof(m_Iv));

            // Select 128 bit AES for the block cipher used for encryption.
            // The block cipher and key used during encryption and during decryption must be the same.
            nn::crypto::Aes128 aes;
            aes.SetKey(pKey, keySize);

            // Performs encryption with CTR mode using 128 bit AES 
            nn::crypto::CtrEncryptor128::Encrypt(pDst, dstSize, pSrc, srcSize, m_Iv, sizeof(m_Iv), aes);
        }
        void Decrypt(void* pDst, size_t dstSize, const void* pSrc, size_t srcSize, const void* pKey, size_t keySize)
        {
            // Select 128 bit AES for the block cipher used for decryption.
            // The block cipher and key used during encryption and during decryption must be the same.
            nn::crypto::Aes128 aes;
            aes.SetKey(pKey, keySize);

            // Performs decryption with CTR mode using 128 bit AES 
            // For IV, the same one used for encryption of the data indicated by pSrc must be specified.
            nn::crypto::CtrDecryptor128::Decrypt(pDst, dstSize, pSrc, srcSize, m_Iv, sizeof(m_Iv), aes);
        }
        void EncryptAndGenerate(void* pDst, size_t dstSize, const void* pSrc, size_t srcSize, const void* pKey, size_t keySize)
        {
            // Randomly generate IV.
            // IV must use a different value for each encryption
            // Furthermore, it must be saved separately because it is needed during decryption.
            NN_COMPILER_ASSERT( CCM_IV_SIZE <= sizeof(m_Iv) );
            nn::crypto::GenerateRandomBytes(m_Iv, CCM_IV_SIZE);

            // Select 128 bit AES for the block cipher used for encryption.
            // The block cipher and key used during encryption and during decryption must be the same.
            nn::crypto::Aes128 aes;
            aes.SetKey(pKey, keySize);

            // Performs encryption with CCM mode using 128 bit AES 
            // This demo does not use Adata.
            nn::crypto::CcmEncryptor::EncryptAndGenerate(
                m_Mac, pDst, dstSize, NULL, 0, pSrc, srcSize, m_Iv, CCM_IV_SIZE, sizeof(m_Mac), aes);
        }
        bool DecryptAndVerify(void* pDst, size_t dstSize, const void* pSrc, size_t srcSize, const void* pKey, size_t keySize)
        {
            // Select 128 bit AES for the block cipher used for encryption.
            // The block cipher and key used during encryption and during decryption must be the same.
            nn::crypto::Aes128 aes;
            aes.SetKey(pKey, keySize);

            // Performs decryption with CCM mode using 128 bit AES 
            // This demo does not use Adata.
            // For IV, the same one used for encryption of the data indicated by pSrc must be specified.
            // MAC specifies the one generated when the data indicated by pSrc is encrypted.
            return nn::crypto::CcmDecryptor::DecryptAndVerify(
                pDst, dstSize, NULL, 0, pSrc, srcSize, m_Iv, CCM_IV_SIZE, m_Mac, sizeof(m_Mac), aes);
        }
    };


    //-------------------------------------------------------------------------------
    // 3. Processes with streaming by combining block cipher and encryption use mode
    //
    class Type3
    {
    private:
        static const size_t CCM_IV_SIZE = 12;
        NN_COMPILER_ASSERT( CCM_IV_SIZE <= nn::crypto::CcmEncryptor::IV_SIZE );

    private:
        bit8    m_Iv[nn::crypto::CtrEncryptor128::IV_SIZE];
        bit8    m_Mac[nn::crypto::CcmEncryptor::MAC_SIZE];

    public:
        size_t GetCipherSize1( size_t plainSize ) { return plainSize; }
        size_t GetCipherSize2( size_t plainSize ) { return plainSize; }

        void Encrypt(void* pDst, size_t dstSize, const void* pSrc, size_t srcSize, const void* pKey, size_t keySize)
        {
            // Randomly generate IV.
            // IV must use a different value for each encryption
            // Furthermore, it must be saved separately because it is needed during decryption.
            nn::crypto::GenerateRandomBytes(m_Iv, sizeof(m_Iv));

            // Select 128 bit AES for the block cipher used for encryption.
            // The block cipher and key used during encryption and during decryption must be the same.
            nn::crypto::Aes128 aes;
            aes.SetKey(pKey, keySize);

            // Uses CTR mode.
            nn::crypto::CtrEncryptor128 ctr;
            ctr.Initialize(aes, m_Iv, sizeof(m_Iv));

            // The input can be passed in small pieces.
            // In this demo, the input is passed 1 byte at a time.
            const size_t packetSize = 1;
            const bit8* pS = reinterpret_cast<const bit8*>(pSrc);
            bit8* pD = reinterpret_cast<bit8*>(pDst);
            size_t offset = 0;
            for( size_t i = 0; i < srcSize; ++i )
            {
                const size_t out = ctr.Update(pD + offset, dstSize - offset, &pS[i], packetSize);
                NN_MAX_ASSERT(out, dstSize - offset);

                offset += out;
            }
            {
                const size_t out = ctr.UpdateFinal(pD + offset, dstSize - offset);
                NN_MAX_ASSERT(out, dstSize - offset);

                offset += out;
                NN_EQUAL_ASSERT(dstSize, offset);
            }

            ctr.Finalize();
        }
        void Decrypt(void* pDst, size_t dstSize, const void* pSrc, size_t srcSize, const void* pKey, size_t keySize)
        {
            // Select 128 bit AES for the block cipher used for decryption.
            // The block cipher and key used during encryption and during decryption must be the same.
            nn::crypto::Aes128 aes;
            aes.SetKey(pKey, keySize);

            // Uses CTR mode.
            // For IV, the same one used for encryption of the data indicated by pSrc must be specified.
            nn::crypto::CtrDecryptor128 ctr;
            ctr.Initialize(aes, m_Iv, sizeof(m_Iv));
            
            // The input can be passed in small pieces.
            // In this demo, the input is passed 1 byte at a time.
            const size_t packetSize = 1;
            const bit8* pS = reinterpret_cast<const bit8*>(pSrc);
            bit8* pD = reinterpret_cast<bit8*>(pDst);
            size_t offset = 0;
            for( size_t i = 0; i < srcSize; ++i )
            {
                const size_t out = ctr.Update(pD + offset, dstSize - offset, &pS[i], packetSize);
                NN_MAX_ASSERT(out, dstSize - offset);

                offset += out;
            }
            {
                const size_t out = ctr.UpdateFinal(pD + offset, dstSize - offset);
                NN_MAX_ASSERT(out, dstSize - offset);

                offset += out;
                NN_EQUAL_ASSERT(dstSize, offset);
            }

            ctr.Finalize();
        }
        void EncryptAndGenerate(void* pDst, size_t dstSize, const void* pSrc, size_t srcSize, const void* pKey, size_t keySize)
        {
            // Randomly generate IV.
            // IV must use a different value for each encryption
            // Furthermore, it must be saved separately because it is needed during decryption.
            NN_COMPILER_ASSERT( CCM_IV_SIZE <= sizeof(m_Iv) );
            nn::crypto::GenerateRandomBytes(m_Iv, CCM_IV_SIZE);

            // Select 128 bit AES for the block cipher used for encryption.
            // The block cipher and key used during encryption and during decryption must be the same.
            nn::crypto::Aes128 aes;
            aes.SetKey(pKey, keySize);

            // Uses CCM mode.
            nn::crypto::CcmEncryptor ccm;
            ccm.Initialize(aes, m_Iv, CCM_IV_SIZE, 0, srcSize, sizeof(m_Mac));

            // This demo does not use Adata.
            ccm.UpdateAdataFinal();

            // The input can be passed in small pieces.
            // In this demo, the input is passed 1 byte at a time.
            const size_t packetSize = 1;
            const bit8* pS = reinterpret_cast<const bit8*>(pSrc);
            bit8* pD = reinterpret_cast<bit8*>(pDst);
            size_t offset = 0;
            for( size_t i = 0; i < srcSize; ++i )
            {
                const size_t out = ccm.UpdatePdata(pD + offset, dstSize - offset, &pS[i], packetSize);
                NN_MAX_ASSERT(out, dstSize - offset);

                offset += out;
            }
            {
                const size_t out = ccm.UpdatePdataFinal(pD + offset, dstSize - offset);
                NN_MAX_ASSERT(out, dstSize - offset);

                offset += out;
                NN_EQUAL_ASSERT(dstSize, offset);
            }

            // Generates a MAC for verifying validity.
            ccm.GenerateMac(m_Mac, sizeof(m_Mac));

            ccm.Finalize();
        }
        bool DecryptAndVerify(void* pDst, size_t dstSize, const void* pSrc, size_t srcSize, const void* pKey, size_t keySize)
        {
            // Select 128 bit AES for the block cipher used for encryption.
            // The block cipher and key used during encryption and during decryption must be the same.
            nn::crypto::Aes128 aes;
            aes.SetKey(pKey, keySize);

            // Uses CCM mode.
            // For the MAC size and IV, the same one used for encryption of the data indicated by pSrc must be specified.
            nn::crypto::CcmDecryptor ccm;
            ccm.Initialize(aes, m_Iv, CCM_IV_SIZE, 0, srcSize, sizeof(m_Mac));

            // This demo does not use Adata.
            ccm.UpdateAdataFinal();

            // The input can be passed in small pieces.
            // In this demo, the input is passed 1 byte at a time.
            const size_t packetSize = 1;
            const bit8* pS = reinterpret_cast<const bit8*>(pSrc);
            bit8* pD = reinterpret_cast<bit8*>(pDst);
            size_t offset = 0;
            for( size_t i = 0; i < srcSize; ++i )
            {
                const size_t out = ccm.UpdateCdata(pD + offset, dstSize - offset, &pS[i], packetSize);
                NN_MAX_ASSERT(out, dstSize - offset);

                offset += out;
            }
            {
                const size_t out = ccm.UpdateCdataFinal(pD + offset, dstSize - offset);
                NN_MAX_ASSERT(out, dstSize - offset);

                offset += out;
                NN_EQUAL_ASSERT(dstSize, offset);
            }

            // Generates a MAC for verifying validity.
            bit8 testMac[nn::crypto::CcmEncryptor::MAC_SIZE];
            ccm.GenerateMac(testMac, sizeof(testMac));

            // Compares the MAC generated during encryption to the MAC generated during decryption.
            // If they do not match, this indicates that the data has been altered.
            return std::memcmp(testMac, m_Mac, sizeof(testMac)) == 0;
        }
    };


    template <typename T>
    void Demo(T& type)
    {
        // Demo for processing only encryption 
        NN_LOG("  " "encryption\n");
        {
            const size_t plainSize = sizeof(TEST_DATA);
            const size_t cipherSize = type.GetCipherSize1(plainSize);

            std::auto_ptr<bit8> pEncrypted(new bit8[cipherSize]);
            std::auto_ptr<bit8> pDecrypted(new bit8[plainSize]);
            NN_POINTER_ASSERT(pEncrypted.get());
            NN_POINTER_ASSERT(pDecrypted.get());

            // Displays data before encryption
            NN_LOG("    " "plain text\n");
            NN_LOG("    " "----\n");
            PrintBin(TEST_DATA, plainSize, "    ");
            NN_LOG("    " "----\n");

            // Encryption
            type.Encrypt(pEncrypted.get(), cipherSize, TEST_DATA, plainSize, TEST_KEY, sizeof(TEST_KEY));

            // Displays data after encryption
            NN_LOG("    " "encrypted text\n");
            NN_LOG("    " "----\n");
            PrintBin(pEncrypted.get(), cipherSize, "    ");
            NN_LOG("    " "----\n");

            // Decryption
            type.Decrypt(pDecrypted.get(), plainSize, pEncrypted.get(), cipherSize, TEST_KEY, sizeof(TEST_KEY));

            // Displays data after decryption
            NN_LOG("    " "decrypted text\n");
            NN_LOG("    " "----\n");
            PrintBin(pDecrypted.get(), plainSize, "    ");
            NN_LOG("    " "----\n");

            // Compares the data after decryption to the data before encryption
            bool match = (std::memcmp(TEST_DATA, pDecrypted.get(), plainSize) == 0);

            // Displays the comparison results
            NN_LOG("    " "decrypted data %s src data\n", (match ? "==": "!="));
        }
        NN_LOG("\n");

        // Demo to process encryption with validity testing
        NN_LOG("  " "encryption and authentication\n");
        {
            const size_t plainSize = sizeof(TEST_DATA);
            const size_t cipherSize = type.GetCipherSize2(plainSize);

            std::auto_ptr<bit8> pEncrypted(new bit8[cipherSize]);
            std::auto_ptr<bit8> pDecrypted(new bit8[plainSize]);
            NN_POINTER_ASSERT(pEncrypted.get());
            NN_POINTER_ASSERT(pDecrypted.get());

            // Displays data before encryption
            NN_LOG("    " "plain text\n");
            NN_LOG("    " "----\n");
            PrintBin(TEST_DATA, plainSize, "    ");
            NN_LOG("    " "----\n");

            // Encryption
            type.EncryptAndGenerate(pEncrypted.get(), cipherSize, TEST_DATA, plainSize, TEST_KEY, sizeof(TEST_KEY));

            // Displays data after encryption
            NN_LOG("    " "encrypted text\n");
            NN_LOG("    " "----\n");
            PrintBin(pEncrypted.get(), cipherSize, "    ");
            NN_LOG("    " "----\n");

            // Decryption
            bool ok1 = type.DecryptAndVerify(pDecrypted.get(), plainSize, pEncrypted.get(), cipherSize, TEST_KEY, sizeof(TEST_KEY));

            // Displays the decryption verification results (should be good) and the data after decryption
            NN_LOG("    " "authentication result = %s\n", (ok1 ? "good": "bad"));
            NN_LOG("    " "decrypted text\n");
            NN_LOG("    " "----\n");
            PrintBin(pDecrypted.get(), plainSize, "    ");
            NN_LOG("    " "----\n");

            // Compares the data after decryption to the data before encryption
            bool match = (std::memcmp(TEST_DATA, pDecrypted.get(), plainSize) == 0);

            // Displays the comparison results
            NN_LOG("    " "decrypted data %s src data\n", (match ? "==": "!="));

            // Try decryption after changing 1 bit of encrypted data
            NN_LOG("    " "edit test\n");
            pEncrypted.get()[0] ^= 0x01;
            bool ok2 = type.DecryptAndVerify(pDecrypted.get(), plainSize, pEncrypted.get(), cipherSize, TEST_KEY, sizeof(TEST_KEY));

            // Display the decryption verification results (should be bad)
            NN_LOG("    " "authentication result = %s\n", (ok2 ? "good": "bad"));
        }
    }
}


void nnMain()
{
    // Call only the nn::applet::Enable function to also allow execution from the HOME Menu
    // HOME Menu transitions, POWER Button presses, and sleep are all unsupported
    nn::applet::Enable();

    NN_LOG("BlockCipher\n");
    NN_LOG("\n");

    Type1 type1;
    Type2 type2;
    Type3 type3;

    NN_LOG("Type 1\n");
    Demo(type1);
    NN_LOG("\n");

    NN_LOG("Type 2\n");
    Demo(type2);
    NN_LOG("\n");

    NN_LOG("Type 3\n");
    Demo(type3);
    NN_LOG("\n");

    NN_LOG("end.\n");
}