/*---------------------------------------------------------------------------*
  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 <vector>

#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 "common_Types.h"

namespace common
{

namespace
{

u64 s_TotalFileSize;
u64 s_FinishedFileSize = 0;
u64 s_Progress = 0;

}

bool VerifyMac(nn::fs::FileInputStream* file, s64 filesize, void* buf, size_t bufSize);

const char* GetCharStr(const wchar_t* path)
{
    static char filename[256];
    std::memset(filename, 0, sizeof(filename));
    std::wcstombs(filename, path, sizeof(filename));
    filename[sizeof(filename) - 1] = '\0';
    return filename;
}

nn::Result CalculateFileNum(std::wstring currentDirectory, u32& fileNum, u32& fileSize)
{
    nn::fs::FileInputStream fis;
    nn::fs::Directory dir;
    nn::Result result;
    std::vector<nn::fs::DirectoryEntry> entryList; //カレントディレクトリのエントリ一覧を格納
    std::vector<nn::fs::DirectoryEntry>::iterator entryIndex;

    result = dir.TryInitialize(currentDirectory.c_str());

    if(result.IsSuccess())
    {
        nn::fs::DirectoryEntry entry;
        s32 numEntry;
        for(;;)
        {
            result = dir.TryRead(&numEntry, &entry, 1);
            if(result.IsFailure())
            {
                dir.Finalize();
                return result;
            }
            if(numEntry == 0)
            {
                // カレントディレクトリを閉じる
                dir.Finalize();

                // カレントディレクトリの子を開く
                for(entryIndex = entryList.begin(); entryIndex != entryList.end(); entryIndex++)
                {
                    if(entryIndex->attributes.isDirectory)
                    {
                        CalculateFileNum(currentDirectory + std::wstring(entryIndex->entryName) + std::wstring(L"/"), fileNum, fileSize);
                    }
                }

                return nn::ResultSuccess();
            }

            entryList.push_back(entry);
            fileNum++;
            fileSize+= entry.entrySize;
        }
    }
    else
    {
        NN_LOG("failed initialize directory\n");
        dir.Finalize();
        return result;
    }

}

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

    if (result.IsFailure())
    {
        nn::dbg::PrintResult(result);
        return false;
    }

    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.IsFailure())
            {
                nn::dbg::PrintResult(result);
                ret_value = false;
            }
            else
            {
                target_from << L"/";
                target_to << L"/";
                // 再帰処理
                if (!CopyDirectory(target_from.str().c_str(), target_to.str().c_str(), buf, bufSize, encode))
                {
                    ret_value = false;
                }
            }
        }
        // ファイルの場合
        else
        {
            // ファイル作成
            nn::fs::FileInputStream from_file;
            nn::fs::FileOutputStream to_file;
            s64 filesize;
            s64 fileSizeWithoutHeaderAndFooter;
            s32 readsize;
            s32 writesize;

            NN_LOG("Copy File %ls\n", target_from.str().c_str());

            // 読み込み対象ファイル開く
            result = from_file.TryInitialize(target_from.str().c_str());

            if (result.IsFailure())
            {
                nn::dbg::PrintResult(result);
                ret_value = false;
            }
            else
            {
                // 読み込み対象ファイルのサイズ取得
                result = from_file.TryGetSize(&filesize);
                if (result.IsFailure())
                {
                    nn::dbg::PrintResult(result);
                    ret_value = false;
                }
                else
                {
                    // 復号するならまず検証する
                    if(!encode)
                    {
                        NN_LOG("Verify CMAC %ls\n", target_from.str().c_str());
                        if(!VerifyMac(&from_file, filesize, buf, bufSize))
                        {
                            COMMON_LOGGER("Verification Failed %s\n", GetCharStr(target_from.str().c_str()));
                            ret_value = false;
                            break;
                        }
                    }

                    nn::crypto::SwAesCtrContext swAesCtrContext;
                    swAesCtrContext.Initialize(iv, common::key, sizeof(key));

                    size_t totalReadSize = 0;
                    nn::crypto::Sha256Context context;
                    context.Initialize();

                    // ファイルサイズをヘッダに書いておく
                    if (encode)
                    {
                        // 書き込み対象ファイル作成
                        result = nn::fs::TryCreateFile(target_to.str().c_str(), filesize);
                        result = to_file.TryInitialize(target_to.str().c_str(), false);
                        if (result.IsFailure())
                        {
                            nn::dbg::PrintResult(result);
                            ret_value = false;
                        }

                        BackupDataHeader header;
                        BackupDataHeader enc;
                        header.size = filesize;
                        result = swAesCtrContext.Encrypt(&enc, &header, sizeof(header));
                        COMMON_LOGGER_RESULT_IF_FAILED_WITH_LINE(result);
                        context.Update(&enc, sizeof(enc));
                        s32 writeSize;
                        result = to_file.TryWrite(&writeSize, &enc, sizeof(enc), false);
                        if (result.IsFailure())
                        {
                            ret_value = false;
                            continue;
                        }
                    }
                    else
                    {
                        // ヘッダを読む
                        // ハッシュの計算は終わっているので復号化のみ
                        BackupDataHeader header;
                        BackupDataHeader dec;
                        s32 readSize;
                        result = from_file.TryRead(&readSize, &header, sizeof(header));
                        if (result.IsFailure())
                        {
                            COMMON_LOGGER_RESULT_IF_FAILED_WITH_LINE(result);
                            ret_value = false;
                            continue;
                        }
                        s_FinishedFileSize += sizeof(header);
                        swAesCtrContext.Decrypt(&dec, &header, sizeof(header));
                        COMMON_LOGGER_RESULT_IF_FAILED_WITH_LINE(result);
                        fileSizeWithoutHeaderAndFooter = dec.size;

                        // 書き込み対象ファイル作成
                        result = nn::fs::TryCreateFile(target_to.str().c_str(), fileSizeWithoutHeaderAndFooter);
                        result = to_file.TryInitialize(target_to.str().c_str(), false);
                        if (result.IsFailure())
                        {
                            nn::dbg::PrintResult(result);
                            ret_value = false;
                        }
                    }

                    while (1)
                    {
                        // バッファの後半半分を暗号・復号用に使う
                        result = from_file.TryRead(&readsize, buf, bufSize / 2);
                        totalReadSize += readsize;

                        if (result.IsFailure())
                        {
                            nn::dbg::PrintResult(result);
                            ret_value = false;
                            break;
                        }
                        else
                        {
                            if (readsize == 0)
                            {
                                if (encode)
                                {
                                    NN_LOG("Add CMAC %ls\n", target_from.str().c_str());
                                    // SHA256を計算してCMACを付加する
                                    bit8 sha256Hash[nn::crypto::Sha256Context::HASH_SIZE];
                                    context.GetHash(sha256Hash);

                                    bit8 cmac[nn::crypto::AES_CMAC_MAC_SIZE];
                                    result = nn::crypto::CalculateAesCmacSw(cmac, sha256Hash,
                                            nn::crypto::Sha256Context::HASH_SIZE, common::cmacKey);
                                    COMMON_LOGGER_RESULT_IF_FAILED_WITH_LINE(result);

                                    result = to_file.TryWrite(&writesize, cmac, sizeof(cmac));
                                    if (result.IsFailure())
                                    {
                                        COMMON_LOGGER_RESULT_IF_FAILED_WITH_LINE(result);
                                        ret_value = false;
                                    }
                                }

                                result = to_file.TryFlush();
                                if (result.IsFailure())
                                {
                                    nn::dbg::PrintResult(result);
                                    ret_value = false;
                                }
                                break;
                            }
                            else
                            {
                                if (encode)
                                {
                                    NN_LOG("EncryptSize = %d\n", readsize);

                                    u8 paddingSize = 0;
                                    // 最後の書き込み時にはPKCS5でパディングする
                                    if (readsize < bufSize / 2)
                                    {
                                        if ((readsize % AES_BLOCK_SIZE) != 0)
                                        {
                                            paddingSize = AES_BLOCK_SIZE - readsize % AES_BLOCK_SIZE;
                                            std::memset(reinterpret_cast<bit8*> (buf) + readsize, paddingSize,
                                                    paddingSize);
                                            readsize += paddingSize;
                                        }
                                    }

                                    // 暗号化後SHA256を計算しつつ書き込み
                                    result = swAesCtrContext.Encrypt(reinterpret_cast<bit8*> (buf) + bufSize / 2, buf,
                                            readsize);
                                    COMMON_LOGGER_RESULT_IF_FAILED_WITH_LINE(result);
                                    context.Update(reinterpret_cast<bit8*> (buf) + bufSize / 2, readsize);

                                    result = to_file.TryWrite(&writesize, reinterpret_cast<bit8*> (buf) + bufSize / 2,
                                            readsize, false);

                                    // 事前計算したファイルサイズに一致させるためパディング分減算
                                    readsize -= paddingSize;

                                    s_FinishedFileSize += readsize;
                                    s_Progress = s_FinishedFileSize * 100 / s_TotalFileSize;
                                    NN_LOG("finish = %lld, total = %lld, progress = %lld\n", s_FinishedFileSize, s_TotalFileSize, s_Progress);
                                }
                                else
                                {
                                    // ハッシュ検証は通っているので復号化しつつ書き込み
                                    // パディング以降は書き込まないよう書き込みサイズを変更する

                                    NN_LOG("DecryptSize = %d\n", readsize);
                                    result = swAesCtrContext.Decrypt(reinterpret_cast<bit8*> (buf) + bufSize / 2, buf,
                                            readsize);
                                    COMMON_LOGGER_RESULT_IF_FAILED_WITH_LINE(result);

                                    // パディングまで読んだかどうか
                                    bool readDone = false;
                                    // パディングまで読んでいたら書き込みサイズを減らす
                                    if (fileSizeWithoutHeaderAndFooter < totalReadSize)
                                    {
                                        readsize -= totalReadSize - fileSizeWithoutHeaderAndFooter;
                                        readDone = true;
                                        s_FinishedFileSize += totalReadSize - fileSizeWithoutHeaderAndFooter;
                                    }

                                    result = to_file.TryWrite(&writesize, reinterpret_cast<bit8*> (buf) + bufSize / 2,
                                            readsize, false);
                                    s_FinishedFileSize += readsize;
                                    s_Progress = s_FinishedFileSize * 100 / s_TotalFileSize;
                                    NN_LOG("finish = %lld, total = %lld, progress = %lld\n", s_FinishedFileSize, s_TotalFileSize, s_Progress);

                                    // 読みきったので次のファイルへ
                                    if (readDone)
                                    {
                                        result = to_file.TryFlush();
                                        if (result.IsFailure())
                                        {
                                            ret_value = false;
                                        }
                                        COMMON_LOGGER_RESULT_IF_FAILED_WITH_LINE(result);
                                        break;
                                    }
                                }

                                if (result.IsFailure())
                                {
                                    nn::dbg::PrintResult(result);
                                    ret_value = false;
                                }

                            }
                        }
                    }
                    to_file.Finalize();
                }
            }
            from_file.Finalize();
        }
    }

    from_dir.Finalize();
    return ret_value;
}

u32 GetProgress()
{
    return s_Progress;
}

void InitializeTransferProgress(u64 totalSize)
{
    s_TotalFileSize = totalSize;
    s_FinishedFileSize = 0;
}

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);
    if (result.IsFailure())
    {
        COMMON_LOGGER_RESULT_IF_FAILED_WITH_LINE(result);
        return false;
    }

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


bool VerifyMac(nn::fs::FileInputStream* file, s64 filesize, void* buf, size_t bufSize)
{
    nn::Result result;
    bit8 sdCmac[nn::crypto::AES_CMAC_MAC_SIZE];

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

    s32 readSize;
    // ハッシュを取得する
    nn::crypto::Initialize();
    result = file->TrySetPosition(filesize - nn::crypto::AES_CMAC_MAC_SIZE);
    if (result.IsSuccess())
    {
        result = file->TryRead(&readSize, sdCmac, sizeof(sdCmac));
        if(result.IsFailure())
        {
            COMMON_LOGGER_RESULT_IF_FAILED_WITH_LINE(result);
            return false;
        }
    }
    else
    {
        COMMON_LOGGER_RESULT_IF_FAILED_WITH_LINE(result);
        return false;
    }

    file->SetPosition(0);

    // ハッシュを計算する
    nn::crypto::Sha256Context context;
    context.Initialize();
    bool ret_value = false;

    size_t totalReadSize = 0;
    while (1)
    {
        result = file->TryRead(&readSize, buf, bufSize);
        totalReadSize  += readSize;

        if (result.IsFailure())
        {
            COMMON_LOGGER_RESULT_IF_FAILED_WITH_LINE(result);
            ret_value = false;
            break;
        }
        else
        {
            if (readSize == 0)
            {
                ret_value = CalculateAndCompareCmac(&context, sdCmac);
                break;
            }
            else
            {
                bool readDone = false;
                // ハッシュ部分まで読んでいたらハッシュ検証サイズを減らす
                size_t fileSizeWithoutCmac = filesize - nn::crypto::AES_CMAC_MAC_SIZE;
                if (fileSizeWithoutCmac < totalReadSize)
                {
                    readSize -= totalReadSize - fileSizeWithoutCmac;
                    readDone = true;
                }

                context.Update(buf, readSize);
                if(readDone)
                {
                    ret_value = CalculateAndCompareCmac(&context, sdCmac);
                    break;
                }

                if (result.IsFailure())
                {
                    COMMON_LOGGER_RESULT_IF_FAILED_WITH_LINE(result);
                    ret_value = false;
                }
            }
        }
    }
    nn::crypto::Finalize();

    file->SetPosition(0);

    return ret_value;
}

}