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

  Copyright 2009 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.h>
#include <nn/version.h>
#include <nn/ptm/CTR/ptm_ApiSysmenu.h>
#include <nn/mcu.h>
#include <nn/cfg/CTR/cfg_Api.h>
#include <nn/cfg/CTR/cfg_ApiInit.h>
#include <nn/cfg/CTR/cfg_ApiSys.h>
#include <nn/ps.h>
#include <nn/nwm.h>
#include <nn/ac.h>
#include <nn/friends.h>
#include <nn/friends/CTR/friends_ApiPrivate.h>
#include <nn/fs/CTR/fs_ArchiveTypesForSystem.h>
#include <nn/fs/CTR/MPCore/fs_FileSystemBasePrivate.h>
#include <nn/fs/fs_ApiSysSaveData.h>
#include <nn/cfg/CTR/cfg_ApiNor.h>
#include <nn/os/os_SharedInfo.h>
#include <nn/pl/CTR/pl_Version.h>
#include <nn/fs/fs_ApiSharedExtSaveData.h>
#include <nn/nim.h>
#include <nn/ndm.h>

#include <nn/crypto/crypto_AesCmac.h>
#include <nn/crypto/crypto_SwAesCtrContext.h>
#include <nn/crypto/crypto_Sha256.h>
#include <nn/crypto/crypto_SwAesCmac.h>

#include "Aes_define.h"
#include "FileTransfer.h"
#include "CommonLogger.h"

#include "demo.h"

#include <vector>
#include <string>

#include "DrawSystemState.h"
#include "FileName.h"
#include "SimplePlayer.h"
#include "CommonLogger.h"
#include "SDMountManager.h"
#include "HeapManager.h"
#include "common_Types.h"
#include "VersionDetect.h"
#include "Util.h"
#include "CommonLogger.h"
#include "SdReaderWriter.h"

namespace
{

// グラフィックスに割り当てるメモリ
const size_t s_GxHeapSize = 0x800000;

const u32 CONSOLE_WIDTH = 38;
const u32 CONSOLE_HEIGHT = 24;
const u32 CONSOLE_MAX_LINE = 1000;

const size_t DECRYPT_THREAD_STACK_SIZE = 0x4000;
nn::os::Thread                                  s_DecryptThread;
nn::os::StackBuffer<DECRYPT_THREAD_STACK_SIZE>    s_DecryptThreadStack;

u32 s_VerifySuccess = 0;
u32 s_VerifyFail = 0;

const wchar_t* const DECRYPT_ROOT_DIRECTORY_PATH = L"sdmc:/CTR_Console_Repair_Decrypt";
const wchar_t* const SD_SAVEDATA_DECRYPT_ROOT_NAME = L"CTR_Console_Repair_Decrypt/CTRBackup/";
const wchar_t* const SD_SAVEDATA_DECRYPT_TWL_PHOTO_ROOT_NAME = L"CTR_Console_Repair_Decrypt/TWLPhotoBackup/";
const wchar_t* const SD_SAVEDATA_DECRYPT_TWL_SOUND_ROOT_NAME = L"CTR_Console_Repair_Decrypt/TWLSoundBackup/";
const wchar_t* const SD_SAVEDATA_DECRYPT_TWL_ROOT_NAME = L"CTR_Console_Repair_Decrypt/TWLBackup/";


}

namespace tools
{
namespace ExportedDataDecrypter
{

void GenerateNandPath(wchar_t* toPath, const wchar_t* fromPath)
{
    // 切り詰める
    std::string tmp(common::GetCharStr(fromPath));
    std::string twlRoot(common::GetCharStr(common::SD_SAVEDATA_TWL_ROOT_NAME));
    std::string twlPhotoRoot(common::GetCharStr(common::SD_SAVEDATA_TWL_PHOTO_ROOT_NAME));
    std::string twlSoundRoot(common::GetCharStr(common::SD_SAVEDATA_TWL_SOUND_ROOT_NAME));
    std::string ctrRoot(common::GetCharStr(common::SD_SAVEDATA_ROOT_NAME));

    std::string output;

    std::string::size_type size;
    size = tmp.find(twlPhotoRoot.c_str());
    if(size == std::string::npos)
    {
        size = tmp.find(twlSoundRoot.c_str());
        if(size == std::string::npos)
        {
            size = tmp.find(ctrRoot.c_str());
            if(size == std::string::npos)
            {
                size = tmp.find(twlRoot.c_str());
                if(size == std::string::npos)
                {
                    // 想定外のパスへの出力のためreturn
                    return;
                }
                else
                {
                    output += std::string("twln:/title/");
                    output += tmp.substr(size + twlRoot.size());
                }
            }
            else
            {
                output += std::string("nand:/data/");
                output += tmp.substr(size + ctrRoot.size());
            }
        }
        else
        {
            output += std::string("twls:/");
            output += tmp.substr(size + twlSoundRoot.size());
        }
    }
    else
    {
        output += std::string("twlp:/");
        output += tmp.substr(size + twlPhotoRoot.size());
    }

    s32 length = std::mbstowcs(toPath, output.c_str(), nn::fs::MAX_FILE_PATH_LENGTH);
    if(length == -1)
    {
        NN_PANIC("failed mbstowcs");
    }
}

bool CalculateAndCompareCmac(nn::crypto::Sha256Context* context, bit8* sdCmac)
{
    nn::Result result;
    bit8 sha256Hash[nn::crypto::Sha256Context::HASH_SIZE];
    bit8 cmac[nn::crypto::AES_CMAC_MAC_SIZE];

    context->GetHash(sha256Hash);
    context->Finalize();

    result = nn::crypto::CalculateAesCmacSw(cmac, sha256Hash, sizeof(sha256Hash), common::cmacKey);
    COMMON_LOGGER_RETURN_FALSE_IF_FAILED(result);

    if(std::memcmp(cmac, sdCmac, sizeof(cmac)) != 0)
    {
        NN_LOG("Faild. Expected CMAC:\n");
        for (u32 i = 0; i < sizeof(cmac); i++)
        {
            NN_LOG("%02X ", cmac[i]);
        }
        NN_LOG("\n");
    }

    return std::memcmp(cmac, sdCmac, sizeof(cmac)) == 0;
}

// ディレクトリ間のコピー
// アーカイブ越しのコピーが可能
// アーカイブにマウントした状態で呼び出す必要あり
// 書き込み先のディレクトリはあらかじめ消去しておくこと。
// 引数はスラッシュ付き
// TODO:分割して短くする
bool DecryptDirectory(const wchar_t * from_path, const wchar_t* to_path, void* buf,
        const size_t bufSize)
{
    using namespace common;

    nn::fs::Directory from_dir;
    nn::fs::DirectoryEntry entry;
    s32 numread = 0;
    std::wostringstream target_from;
    std::wostringstream target_to;
    bool ret_value = true;

    nn::Result result = from_dir.TryInitialize(from_path);
    COMMON_LOGGER_RETURN_FALSE_IF_FAILED(result);

    while (1)
    {
        result = from_dir.TryRead(&numread, &entry, 1);
        if (result.IsFailure() || numread != 1)
        {
            break;
        }

        if (std::wcscmp(entry.entryName, L".") == 0 || std::wcscmp(entry.entryName, L"..") == 0)
        {
            continue;
        }

        target_from.str(L"");
        target_from.clear(std::stringstream::goodbit);
        target_from << from_path << entry.entryName;

        target_to.str(L"");
        target_to.clear(std::stringstream::goodbit);
        target_to << to_path << entry.entryName;


        // ディレクトリの場合
        if (entry.attributes.isDirectory)
        {
            // ディレクトリ作成
            NN_LOG("Create Directory %ls\n", target_to.str().c_str());
            result = nn::fs::TryCreateDirectory(target_to.str().c_str());
            if (result.IsSuccess() || result.IsFailure() && result <= nn::fs::ResultAlreadyExists())
            {
                target_from << L"/";
                target_to << L"/";

                // 再帰処理
                if (!DecryptDirectory(target_from.str().c_str(), target_to.str().c_str(), buf, bufSize))
                {
                    ret_value = false;
                }
            }
            else
            {
                COMMON_LOGGER_RETURN_FALSE_IF_FAILED(result);
            }
        }
        // ファイルの場合
        // SDカード上のファイルのCMACを検証する
        else
        {
            nn::fs::FileInputStream sdFile;
            nn::fs::FileOutputStream sdOutFile;
            s64 sdFileSize;

            result = sdFile.TryInitialize(target_from.str().c_str());
            if(result.IsFailure())
            {
                ret_value = false;
                break;
            }


            result = sdFile.TryGetSize(&sdFileSize);
            COMMON_LOGGER_RETURN_FALSE_IF_FAILED(result);

            {

                bit8 sdCmac[nn::crypto::AES_CMAC_MAC_SIZE];

                // ハッシュが付加されていないとエラー
                if (sdFileSize < nn::crypto::AES_CMAC_MAC_SIZE)
                {
                    return false;
                }

                s32 readSize;
                // ハッシュを取得する
                result = sdFile.TrySetPosition(sdFileSize - nn::crypto::AES_CMAC_MAC_SIZE);
                COMMON_LOGGER_RETURN_FALSE_IF_FAILED(result);

                result = sdFile.TryRead(&readSize, sdCmac, sizeof(sdCmac));
                COMMON_LOGGER_RETURN_FALSE_IF_FAILED(result);

                sdFile.SetPosition(0);


                // 復号化しながらハッシュを計算する
                nn::crypto::SwAesCtrContext swAesCtrContext;
                swAesCtrContext.Initialize(iv, common::key, sizeof(key));

                nn::crypto::Sha256Context context;
                context.Initialize();

                wchar_t nandPath[nn::fs::MAX_FILE_PATH_LENGTH];
                // sdパスからnandパスを生成する
                GenerateNandPath(nandPath, target_from.str().c_str());

                // NAND上のフルパスをハッシュに含めている
                context.Update(nandPath, std::wcslen(nandPath) * sizeof(wchar_t));

                size_t totalReadSize = 0;

                BackupDataHeader enc;
                BackupDataHeader dec;
                std::memset(&enc, 0, sizeof(enc));
                std::memset(&dec, 0, sizeof(dec));
                result = sdFile.TryRead(&readSize, &enc, sizeof(enc));
                COMMON_LOGGER_RETURN_FALSE_IF_FAILED(result);
                totalReadSize += readSize;
                context.Update(&enc, sizeof(enc));
                result = swAesCtrContext.Decrypt(&dec, &enc, sizeof(enc));
                COMMON_LOGGER_RETURN_FALSE_IF_FAILED(result);

                // 書き込み対象ファイル作成
                s32 writeSize;

                result = nn::fs::TryCreateFile(target_to.str().c_str(), dec.size);
                COMMON_LOGGER_RETURN_FALSE_IF_FAILED(result);

                result = sdOutFile.TryInitialize(target_to.str().c_str(), true);
                COMMON_LOGGER_RESULT_IF_FAILED(result);

                while (1)
                {
                    result = sdFile.TryRead(&readSize, buf, bufSize / 2);
                    COMMON_LOGGER_RETURN_FALSE_IF_FAILED(result);
                    totalReadSize += readSize;

                    {
                        if (readSize == 0)
                        {
                            ret_value = CalculateAndCompareCmac(&context, sdCmac);
                            if(!ret_value)
                            {
                                COMMON_LOGGER("********** Verification Failed ********** %s\n", common::GetCharStr(target_from.str().c_str()));
                                s_VerifyFail++;
                            }
                            else
                            {
                                COMMON_LOGGER("Success %s\n", common::GetCharStr(target_from.str().c_str()));
                                s_VerifySuccess++;
                            }
                            break;
                        }
                        else
                        {
                            bool readDone = false;;

                            // 復号化
                            result = swAesCtrContext.Decrypt(reinterpret_cast<bit8*>(buf) + bufSize / 2, buf, readSize);
                            COMMON_LOGGER_RETURN_FALSE_IF_FAILED(result);

                            // SHA256Hash確認のためCMACまで読んだかどうか調べる
                            s32 readSizeforCmac = readSize;
                            if (sdFileSize - nn::crypto::AES_CMAC_MAC_SIZE < totalReadSize)
                            {
                                // 最大CMACぶんのサイズを減らす
                                readSizeforCmac -= totalReadSize - (sdFileSize - nn::crypto::AES_CMAC_MAC_SIZE);
                                readDone = true;
                            }

                            if (readSizeforCmac != 0)
                            {
                                context.Update(buf, readSizeforCmac);
                            }
                            else
                            {
                                ret_value = CalculateAndCompareCmac(&context, sdCmac);
                                if (!ret_value)
                                {
                                    COMMON_LOGGER(
                                            "********** Verification Failed ********** %s\n", common::GetCharStr(target_from.str().c_str()));
                                    s_VerifyFail++;
                                }
                                else
                                {
                                    COMMON_LOGGER("Success %s\n", common::GetCharStr(target_from.str().c_str()));
                                    s_VerifySuccess++;
                                }
                                break;
                            }

                            s32 sdWriteSize = readSize;
                            // パディングまで読んでいたら書き込みサイズを減らす
                            if (dec.size + sizeof(dec) < totalReadSize)
                            {
                                sdWriteSize -= totalReadSize - (dec.size + sizeof(dec));
                            }

                            // 書き込み
                            result = sdOutFile.TryWrite(&writeSize, reinterpret_cast<bit8*>(buf) + bufSize / 2, sdWriteSize,
                                    true);
                            COMMON_LOGGER_RETURN_FALSE_IF_FAILED(result);

                            if(readDone)
                            {
                                ret_value = CalculateAndCompareCmac(&context, sdCmac);
                                if(!ret_value)
                                {
                                    COMMON_LOGGER("********** Verification Failed ********** %s\n", common::GetCharStr(target_from.str().c_str()));
                                    s_VerifyFail++;
                                }
                                else
                                {
                                    COMMON_LOGGER("Success %s\n", common::GetCharStr(target_from.str().c_str()));
                                    s_VerifySuccess++;
                                }
                                break;
                            }
                        }
                    }
                }
                sdFile.Finalize();
                sdOutFile.Finalize();
            }
        }
    }
    from_dir.Finalize();
    return ret_value;
}


void DecryptThreadFunc()
{
    nn::Result result;

    COMMON_LOGGER("DecryptThreadFunc Start\n");
    s_VerifyFail = 0;
    s_VerifySuccess = 0;

    result = common::SdMountManager::Mount();

    size_t bufSize = common::GetAllocatableSize(AES_BLOCK_SIZE * 2);
    if(bufSize > common::FILE_COPY_HEAP_SIZE)
    {
        bufSize = common::FILE_COPY_HEAP_SIZE;
    }

    common::SdReaderWriter sdWriter;
    common::HeapManager heap(bufSize, AES_BLOCK_SIZE * 2);
    void* buf = heap.GetAddr();
    if (buf != NULL)
    {
        result = nn::fs::TryDeleteDirectoryRecursively(
                DECRYPT_ROOT_DIRECTORY_PATH);

        // ディレクトリが無ければ作る
        nn::fs::Directory dir;
        result = dir.TryInitialize(DECRYPT_ROOT_DIRECTORY_PATH);
        if (result.IsFailure())
        {
            result = nn::fs::TryCreateDirectory(DECRYPT_ROOT_DIRECTORY_PATH);
        }
        dir.Finalize();

        result = sdWriter.CreateDirectory((::std::wstring(common::SDMC_ROOT_DIRECTORY_PATH) +
                std::wstring(SD_SAVEDATA_DECRYPT_TWL_ROOT_NAME)).c_str());
        DecryptDirectory(
                (std::wstring(common::SDMC_ROOT_DIRECTORY_PATH) + ::std::wstring(common::SD_SAVEDATA_TWL_ROOT_NAME)).c_str(),
                (std::wstring(common::SDMC_ROOT_DIRECTORY_PATH) + ::std::wstring(SD_SAVEDATA_DECRYPT_TWL_ROOT_NAME)).c_str(),
                buf, bufSize);

        result = sdWriter.CreateDirectory((::std::wstring(common::SDMC_ROOT_DIRECTORY_PATH) +
                std::wstring(SD_SAVEDATA_DECRYPT_TWL_SOUND_ROOT_NAME)).c_str());
        DecryptDirectory(
                (std::wstring(common::SDMC_ROOT_DIRECTORY_PATH)
                        + ::std::wstring(common::SD_SAVEDATA_TWL_SOUND_ROOT_NAME)).c_str(),
                (std::wstring(common::SDMC_ROOT_DIRECTORY_PATH)
                        + ::std::wstring(SD_SAVEDATA_DECRYPT_TWL_SOUND_ROOT_NAME)).c_str(),
                        buf, bufSize);

        result = sdWriter.CreateDirectory((::std::wstring(common::SDMC_ROOT_DIRECTORY_PATH) +
                std::wstring(SD_SAVEDATA_DECRYPT_TWL_PHOTO_ROOT_NAME)).c_str());
        DecryptDirectory(
                (std::wstring(common::SDMC_ROOT_DIRECTORY_PATH)
                        + ::std::wstring(common::SD_SAVEDATA_TWL_PHOTO_ROOT_NAME)).c_str(),
                (std::wstring(common::SDMC_ROOT_DIRECTORY_PATH)
                        + ::std::wstring(SD_SAVEDATA_DECRYPT_TWL_PHOTO_ROOT_NAME)).c_str(),
                        buf, bufSize);

        result = sdWriter.CreateDirectory((::std::wstring(common::SDMC_ROOT_DIRECTORY_PATH) +
                std::wstring(SD_SAVEDATA_DECRYPT_ROOT_NAME)).c_str());
        DecryptDirectory(
                (std::wstring(common::SDMC_ROOT_DIRECTORY_PATH) + ::std::wstring(common::SD_SAVEDATA_ROOT_NAME)).c_str(),
                (std::wstring(common::SDMC_ROOT_DIRECTORY_PATH) + ::std::wstring(SD_SAVEDATA_DECRYPT_ROOT_NAME)).c_str(),
                buf, bufSize);
    }

    common::SdMountManager::Unmount();

    COMMON_LOGGER("Verify Thread Finalize\n");

    COMMON_LOGGER("\n\n");
    COMMON_LOGGER("Verify Finished, success = %d, fail = %d\n", s_VerifySuccess, s_VerifyFail);
}


} // namespace ExportedDataDecrypter
} // namespace tools

extern "C" void nninitSetupDaemons(void)
{
}


extern "C" void nnMain(void)
{
    nn::Result result;

    // os の初期化
    nn::os::Initialize();

    // fs の初期化
    nn::fs::Initialize();

    // appletの初期化
    nn::applet::Enable( false );

    // hid の初期化
    result = nn::hid::Initialize();
    NN_ERR_THROW_FATAL_IF_FATAL_ONLY(result);

    // ヒープの確保
    common::InitializeHeap();

    // RenderSystem の準備
    common::HeapManager gxHeap(s_GxHeapSize);
    uptr heapForGx = reinterpret_cast<uptr>(gxHeap.GetAddr());
    demo::RenderSystemDrawing s_RenderSystem;
    s_RenderSystem.Initialize(heapForGx, s_GxHeapSize);

    // ログ描画の初期化
    common::Logger::GetLoggerInstance()->Initialize(CONSOLE_WIDTH, CONSOLE_HEIGHT, CONSOLE_MAX_LINE, &s_RenderSystem);

    // RenderSystemを作ってからログが出せる
    common::Logger::InitializeEjectThread();

    COMMON_LOGGER("Decrypt Start\n");

    // ボタン入力
    nn::hid::PadReader s_PadReader;
    nn::hid::PadStatus padStatus;

    for(;;)
    {
        s_PadReader.ReadLatest(&padStatus);

        if(padStatus.trigger & nn::hid::BUTTON_A)
        {
            // SDにコピーするためのスレッドの作成
            if(s_DecryptThread.IsValid() && !s_DecryptThread.IsAlive())
            {
                s_DecryptThread.Join();
                s_DecryptThread.Finalize();
            }
            s_DecryptThread.Start(tools::ExportedDataDecrypter::DecryptThreadFunc, s_DecryptThreadStack);
        }

        // コンソールスクロール
        if(padStatus.hold & nn::hid::BUTTON_UP)
        {
            common::Logger::GetLoggerInstance()->ScrollUp();
        }

        // コンソールスクロール
        if(padStatus.hold & nn::hid::BUTTON_DOWN)
        {
            common::Logger::GetLoggerInstance()->ScrollDown();
        }

        if(padStatus.hold & nn::hid::BUTTON_LEFT)
        {
            common::Logger::GetLoggerInstance()->ScrollToBegin();
        }

        if(padStatus.hold & nn::hid::BUTTON_RIGHT)
        {
            common::Logger::GetLoggerInstance()->ScrollToEnd();
        }

        s_RenderSystem.SetRenderTarget(NN_GX_DISPLAY0);
        s_RenderSystem.Clear();
        s_RenderSystem.SetColor(1.f, 1.f, 1.f);

        common::Logger::GetLoggerInstance()->DrawConsole();
        s_RenderSystem.SwapBuffers();

        s_RenderSystem.WaitVsync(NN_GX_DISPLAY_BOTH);
    }
}