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

 Copyright 2009-2011 Nintendo.  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 <nn/crypto/crypto_SwAesCtrContext.h>

#include "Aes_define.h"
#include "HeapChecker.h"
#include "SdReaderWriter.h"
#include "CommonLogger.h"

namespace ConsoleBackup
{

namespace
{
HeapChecker s_HeapChecker;
}

nn::Result CheckHeap(std::wstring saveRoot, void* buf, size_t& bufSize)
{
    return s_HeapChecker.Check(saveRoot, buf, bufSize);
}

size_t GetCheckedHeapSize()
{
    return s_HeapChecker.GetCheckedSize();
}


HeapChecker::HeapChecker() :
        m_CheckedSize(0), m_IsAlreadyChecked(false)
{

}

HeapChecker::~HeapChecker()
{
    // TODO Auto-generated destructor stub
}

nn::Result HeapChecker::Check(std::wstring saveRoot, void* buf, size_t& bufSize)
{
    // 2回目のチェックは行わない
    if(m_IsAlreadyChecked)
    {
        NN_PANIC("HeapChecker Already checked");
    }

    nn::Result result;
    common::SdReaderWriter sdReaderWriter;

    // cecdセーブデータが存在しない場合は値を変更せずreturnする
    {
        nn::fs::FileInputStream nandFile;
        result = nandFile.TryInitialize(std::wstring(saveRoot + std::wstring(L"sysdata/00010026/00000000")).c_str());
        if (result <= nn::fs::ResultNotFound())
        {
            COMMON_LOGGER("No Reference Data. Use Default Memory Size.\n");
            m_CheckedSize = bufSize;
            m_IsAlreadyChecked = true;
            return nn::ResultSuccess();
        }
    }

    for(; HEAP_SIZE_MIN < bufSize; bufSize /= 2)
    {
        result = nn::fs::TryDeleteFile(HEAP_CHECKER_FILE);
        if(!(result <= nn::fs::ResultNotFound()))
        {
            COMMON_LOGGER_RETURN_RESULT_IF_FAILED(result);
        }

        // NANDから読み込む
        nn::fs::FileInputStream nandFile;
        result = nandFile.TryInitialize(std::wstring(saveRoot + std::wstring(L"sysdata/00010026/00000000")).c_str());
        COMMON_LOGGER_RETURN_RESULT_IF_FAILED(result);
        nn::fs::FileStream sdFile;
        result = sdFile.TryInitialize(HEAP_CHECKER_FILE,
                nn::fs::OPEN_MODE_READ | nn::fs::OPEN_MODE_WRITE | nn::fs::OPEN_MODE_CREATE);
        COMMON_LOGGER_RETURN_RESULT_IF_FAILED(result);

        s32 readSize = 0;
        nn::crypto::Sha256Context writeContext;
        writeContext.Initialize();
        nn::crypto::SwAesCtrContext aes;
        aes.Initialize(common::checkiv, common::checkkey, sizeof(common::checkkey));
        for(;;)
        {
            result = nandFile.TryRead(&readSize, buf, bufSize);
            COMMON_LOGGER_RETURN_RESULT_IF_FAILED(result);

            if(readSize == 0)
            {
                break;
            }

            aes.Encrypt(buf, buf, readSize);
            writeContext.Update(buf, readSize);

            s32 writeSize;
            result = sdFile.TryWrite(&writeSize, buf, readSize, false);
            COMMON_LOGGER_RETURN_RESULT_IF_FAILED(result);

            result = sdFile.TryFlush();
            COMMON_LOGGER_RETURN_RESULT_IF_FAILED(result);
        }

        bit8 sha256WriteContext[nn::crypto::Sha256Context::HASH_SIZE];
        writeContext.GetHash(sha256WriteContext);

        nn::crypto::Sha256Context readContext;
        readContext.Initialize();
        sdFile.SetPosition(0);
        for(;;)
        {
            result = sdFile.TryRead(&readSize, buf, bufSize);
            COMMON_LOGGER_RETURN_RESULT_IF_FAILED(result);

            if(readSize == 0)
            {
                break;
            }

            readContext.Update(buf, readSize);
        }

        bit8 sha256ReadContext[nn::crypto::Sha256Context::HASH_SIZE];
        readContext.GetHash(sha256ReadContext);

        if(std::memcmp(sha256WriteContext, sha256ReadContext, sizeof(sha256ReadContext)) != 0)
        {
            COMMON_LOGGER("MemoryChecker Failure!! MemrySize: %d\n", bufSize);
        }
        else
        {
            // チェックOK
            break;
        }
    }

    nn::fs::TryDeleteFile(HEAP_CHECKER_FILE);

    if(bufSize == HEAP_SIZE_MIN)
    {
        return nn::Result(nn::Result::LEVEL_FATAL, nn::Result::SUMMARY_OUT_OF_RESOURCE, nn::Result::MODULE_APPLICATION,
                nn::Result::DESCRIPTION_OUT_OF_MEMORY);
    }
    else
    {
        m_CheckedSize = bufSize;
        m_IsAlreadyChecked = true;
        return nn::ResultSuccess();
    }
}

size_t HeapChecker::GetCheckedSize()
{
    if(m_IsAlreadyChecked)
    {
        return m_CheckedSize;
    }
    else
    {
        return 0;
    }
}

} /* namespace ConsoleBackup */