ctr_Repair/trunk/NetworkUpdater/sources/common/SaveDataMover.cpp

/*---------------------------------------------------------------------------*
 Project:  Horizon
 File:     SaveDataMover.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 <string>

#include <nn/fs/fs_ApiSysSaveData.h>
#include <nn/fs/fs_ApiSharedExtSaveData.h>
#include <nn/fs/fs_ApiDeviceMove.h>
#include <nn/fs/fs_ResultPrivate.h>
#include <nn/crypto/crypto_SwAesCtrContext.h>

#include <nn/ndm.h>

#include "Aes_define.h"
#include "SaveDataMover.h"
#include "CommonLogger.h"
#include "FileName.h"
#include "SdReaderWriter.h"
#include "FileTransfer.h"
#include "HeapManager.h"


namespace common
{

namespace
{
    const char* ARC_NAME = "sext:";

    const s32 MAX_SYSTEM_SAVE_DATA_ID_NUM = 256;
    nn::fs::SystemSaveDataId s_SystemSaveDataIdList[MAX_SYSTEM_SAVE_DATA_ID_NUM * sizeof(bit32)];
    const s32 MAX_SHARED_EXT_SAVE_DATA_ID_NUM = 10;
    nn::fs::SharedExtSaveDataId s_SharedExtSaveDataIdList[MAX_SHARED_EXT_SAVE_DATA_ID_NUM * sizeof(bit32)];

    bool IsNotMovedId(bit32 id)
    {
        return id == 0x0001002c;    // nim
    }

    bit32 HexStringToBit32(const char* hex, int maxLen)
    {
        bit32 val = 0;
        for(int i = 0; i < maxLen; ++i)
        {
            val *= 16;
            if(hex[i] >= '0' && hex[i] <= '9')
            {
                val += (hex[i] - '0');
            }
            else if(hex[i] >= 'a' && hex[i] <= 'f')
            {
                val += (hex[i] - 'a' + 10);
            }
            else if(hex[i] >= 'A' && hex[i] <= 'F')
            {
                val += (hex[i] - 'A' + 10);
            }
            else
            {
                break;
            }
        }
        return val;
    }
}


SaveDataMover::SaveDataMover() :
        m_Progress(0), m_FinishedSize(0), m_TotalSize(0), m_Result(nn::ResultSuccess())
{
}

SaveDataMover::~SaveDataMover()
{
}

void SaveDataMover::StartExport(void* buf, size_t bufSize, u64* progress)
{
    SetupExport();
    COMMON_LOGGER_RETURN_VOID_IF_FAILED(GetLastResult());
    CalculateExportFileSize();
    COMMON_LOGGER_RETURN_VOID_IF_FAILED(GetLastResult());
    ExportSystemSaveData(buf, bufSize, progress);
    COMMON_LOGGER_RETURN_VOID_IF_FAILED(GetLastResult());
    ExportSharedExtSaveData(buf, bufSize, progress);
}

void SaveDataMover::StartImport(void* buf, size_t bufSize, u64* progress)
{
    SetupImport(buf, bufSize);
    COMMON_LOGGER_RETURN_VOID_IF_FAILED(GetLastResult());
    CalculateImportFileSize();
    COMMON_LOGGER_RETURN_VOID_IF_FAILED(GetLastResult());
    ImportSystemSaveData(buf, bufSize, progress);
    COMMON_LOGGER_RETURN_VOID_IF_FAILED(GetLastResult());
    ImportSharedExtSaveData(buf, bufSize, progress);
    COMMON_LOGGER_RETURN_VOID_IF_FAILED(GetLastResult());
}

nn::Result SaveDataMover::GetLastResult()
{
    return m_Result;
}

void SaveDataMover::SetupExport()
{
    // 1. デーモンの停止
    m_Result = nn::ndm::Initialize();
    COMMON_LOGGER_RETURN_VOID_IF_FAILED(m_Result);
    m_Result = nn::ndm::SuspendScheduler();
    COMMON_LOGGER_RETURN_VOID_IF_FAILED(m_Result);

    // 2. StartDeviceMoveAsSource
    nn::fs::DeviceMoveContext moveContext;
    m_Result = nn::fs::StartDeviceMoveAsSource(&moveContext);
    COMMON_LOGGER_RETURN_VOID_IF_FAILED(m_Result);

    // コンテキストのSDへの出力
    common::HeapManager contextHeap(sizeof(moveContext));
    void* enc;
    enc = contextHeap.GetAddr();
    if(!enc)
    {
        m_Result = nn::Result(nn::Result::LEVEL_FATAL, nn::Result::SUMMARY_OUT_OF_RESOURCE, nn::Result::MODULE_APPLICATION,
                nn::Result::DESCRIPTION_OUT_OF_MEMORY);
        return;
    }

    // AES暗号化する
    nn::crypto::SwAesCtrContext swAesCtrContest;

    swAesCtrContest.Initialize(common::iv, common::key, sizeof(common::key));
    m_Result = swAesCtrContest.Encrypt(enc, &moveContext, sizeof(moveContext));
    COMMON_LOGGER_RETURN_VOID_IF_FAILED(m_Result);

    common::SdReaderWriter sdWriter;
    m_Result = sdWriter.WriteBufWithCmac(common::MOVE_CONTEXT_PATHNAME, enc, sizeof(moveContext));
    COMMON_LOGGER_RETURN_VOID_IF_FAILED(m_Result);

    // 3. 出力用ディレクトリの作成
    // 3.1 システムセーブデータ用ディレクトリ
    m_Result = sdWriter.CreateDirectory((::std::wstring(common::SDMC_ROOT_DIRECTORY_PATH) + ::std::wstring(
            common::SD_SAVEDATA_ROOT_NAME) + std::wstring(common::SD_SAEVDATA_SYS_NAME)).c_str());
    COMMON_LOGGER_RETURN_VOID_IF_FAILED(m_Result);

    // 3.2 共有拡張セーブデータ用ディレクトリ
    m_Result = sdWriter.CreateDirectory((::std::wstring(common::SDMC_ROOT_DIRECTORY_PATH) + ::std::wstring(
            common::SD_SAVEDATA_ROOT_NAME) + std::wstring(common::SD_SAEVDATA_EXT_NAME)).c_str());
    COMMON_LOGGER_RETURN_VOID_IF_FAILED(m_Result);
}

void SaveDataMover::CalculateExportFileSize()
{
    CalculateExportSystemSaveDataSize();
    CalculateExportSharedExtSaveDataSize();
}

void SaveDataMover::CalculateExportSystemSaveDataSize()
{
    s32 systemSaveDataIdNum;
    m_Result = nn::fs::EnumerateSystemSaveData(&systemSaveDataIdNum, s_SystemSaveDataIdList, MAX_SYSTEM_SAVE_DATA_ID_NUM);
    COMMON_LOGGER_RETURN_VOID_IF_FAILED(m_Result);

    // システムセーブデータのサイズを確認する
    for(s32 i = 0; i < systemSaveDataIdNum; ++i)
    {
        nn::fs::FileInputStream input;

        bit32 id = s_SystemSaveDataIdList[i];
        if(IsNotMovedId(id))
        {
            continue;
        }

        m_Result = nn::fs::OpenSystemSaveDataRawStorageFile(&input, id);
        if(m_Result.IsFailure())
        {
            COMMON_LOGGER_WARN("Can't Export: %08x\n", id);
            COMMON_LOGGER_RETURN_VOID_IF_FAILED(m_Result);
            continue;
        }

        m_TotalSize += input.GetSize();
    }
}

void SaveDataMover::CalculateExportSharedExtSaveDataSize()
{
    // 共有拡張セーブデータのサイズを確認する
    // 0. ID の列挙を行う。
    s32 sharedExtSaveDataIdNum;
    m_Result = nn::fs::EnumerateSharedExtSaveData(&sharedExtSaveDataIdNum, s_SharedExtSaveDataIdList,
            MAX_SHARED_EXT_SAVE_DATA_ID_NUM);
    for(s32 i = 0; i < sharedExtSaveDataIdNum; ++i)
    {
        bit32 id = s_SharedExtSaveDataIdList[i];

        // 1. fs::MountSharedExtSaveDataRawStorage を呼び
        //    ID に対応するアーカイブをマウントする。
        // 既存の共有拡張セーブデータ領域へのアーカイブをマウントします。
        m_Result = nn::fs::MountSharedExtSaveDataRawStorage(ARC_NAME, id);
        COMMON_LOGGER_RETURN_VOID_IF_FAILED(m_Result);

        // 2. アーカイブ内を走査して、含まれるファイルサイズを計算する
        m_Result = common::CalculateFileSizeRecursively(L"sext:/", m_TotalSize);
        COMMON_LOGGER_RETURN_VOID_IF_FAILED(m_Result);

        // 3. アンマウントする
        m_Result = nn::fs::Unmount(ARC_NAME);
        COMMON_LOGGER_RETURN_VOID_IF_FAILED(m_Result);
    }
    NN_LOG("calculatedSize = %lld\n", m_TotalSize);
}

void SaveDataMover::ExportSystemSaveData(void* buf, size_t bufSize, u64* progress)
{
    s64 totalFileSize = 0;

    // 0. ID の列挙を行う。
    s32 systemSaveDataIdNum;
    m_Result = nn::fs::EnumerateSystemSaveData(&systemSaveDataIdNum, s_SystemSaveDataIdList, MAX_SYSTEM_SAVE_DATA_ID_NUM);
    COMMON_LOGGER_RETURN_VOID_IF_FAILED(m_Result);

    for(s32 i = 0; i < systemSaveDataIdNum; ++i)
    {
        nn::fs::FileInputStream input;

        bit32 id = s_SystemSaveDataIdList[i];
        if(IsNotMovedId(id))
        {
            continue;
        }

        NN_LOG("id: %08x\n", id);

        // 1. fs::OpenSystemSaveDataRawStorageFile を呼び
        //    ID に対するセーブデータの FileInputStream を得る。
        m_Result = nn::fs::OpenSystemSaveDataRawStorageFile(&input, id);
        if(m_Result.IsFailure())
        {
            COMMON_LOGGER_WARN("Can't Export: %08x\n", id);
            COMMON_LOGGER_RETURN_VOID_IF_FAILED(m_Result);
            continue;
        }
        NN_LOG("fileSize = %lld\n", input.GetSize());
        totalFileSize += input.GetSize();

        // SD へコピー
        {
            char outPath[64];
            nn::nstd::TSNPrintf(outPath, 64, "%s/%08x", SD_SAVEDATA_SYS_ROOT_PATH, id);
            NN_LOG("outPath = %s\n", outPath);
            nn::fs::FileOutputStream output;
            m_Result = output.TryInitialize(outPath, true);
            if(m_Result.IsFailure())
            {
                COMMON_LOGGER_WARN("Can't Export: %08x\n", id);
                COMMON_LOGGER_RETURN_VOID_IF_FAILED(m_Result);
                continue;
            }

            // 2. 1 で得たストリームからデータを読み込む。
            //    本来はここで読み込んだデータを引っ越し先に転送するが
            //    ここでは SD カードにエクスポートする。
            m_Result = CopyFile(input, output, buf, bufSize, progress);
            if(m_Result.IsFailure())
            {
                COMMON_LOGGER_WARN("Can't Export: %08x\n", id);
                COMMON_LOGGER_RETURN_VOID_IF_FAILED(m_Result);
                continue;
            }
        }
    }
    NN_LOG("\n");

    NN_LOG("totalFileSize = %lld\n", totalFileSize);
}

void SaveDataMover::ExportSharedExtSaveData(void* buf, size_t bufSize, u64* progress)
{
    // 0. ID の列挙を行う。
    s32 sharedExtSaveDataIdNum;
    m_Result = nn::fs::EnumerateSharedExtSaveData(&sharedExtSaveDataIdNum, s_SharedExtSaveDataIdList,
            MAX_SHARED_EXT_SAVE_DATA_ID_NUM);
    for(s32 i = 0; i < sharedExtSaveDataIdNum; ++i)
    {
        bit32 id = s_SharedExtSaveDataIdList[i];
        NN_LOG("Export: %08x\n", id);

        // 1. fs::MountSharedExtSaveDataRawStorage を呼び
        //    ID に対応するアーカイブをマウントする。
        // 既存の共有拡張セーブデータ領域へのアーカイブをマウントします。
        m_Result = nn::fs::MountSharedExtSaveDataRawStorage(ARC_NAME, id);
        COMMON_LOGGER_RETURN_VOID_IF_FAILED(m_Result);

        m_Result = Export(buf, bufSize, id, progress);
        COMMON_LOGGER_RETURN_VOID_IF_FAILED(m_Result);

        // 3. アンマウントする
        m_Result = nn::fs::Unmount(ARC_NAME);
        COMMON_LOGGER_RETURN_VOID_IF_FAILED(m_Result);
    }

}

nn::Result SaveDataMover::Export(void* buf, size_t bufSize, bit32 id, u64* progress)
{
    const s32 MAX_PATH_LEN = 127;
    char destRoot[MAX_PATH_LEN + 1];

    nn::nstd::TSNPrintf(destRoot, MAX_PATH_LEN, "%s/%08x", SD_SAVEDATA_EXT_ROOT_PATH, id);
    NN_LOG("destRoot = %s\n", destRoot);;

    // 出力先ディレクトリを作成
    COMMON_LOGGER_RETURN_RESULT_IF_FAILED(nn::fs::TryCreateDirectory(destRoot));

    // 2. アーカイブ内を走査して、含まれるファイルと
    //    ディレクトリを、その構造を保ったまま
    //    エクスポートする。
    COMMON_LOGGER_RETURN_RESULT_IF_FAILED(CopyRecursivly(buf, bufSize, ARC_NAME, destRoot, progress));

    return nn::ResultSuccess();
}

void SaveDataMover::SetupImport(void* buf, size_t bufSize)
{
    // 引っ越しコンテクストを取得
    nn::fs::DeviceMoveContext moveContext;

    // コンテキストのSDからの入力
    size_t readSize;
    common::SdReaderWriter sdReader;
    m_Result = sdReader.ReadBufWithCmac(common::MOVE_CONTEXT_PATHNAME, buf, bufSize, &readSize);
    COMMON_LOGGER_RETURN_VOID_IF_FAILED(m_Result);

    // AES復号化する
    nn::crypto::SwAesCtrContext swAesCtrContest;

    swAesCtrContest.Initialize(common::iv, common::key, sizeof(common::key));
    m_Result = swAesCtrContest.Decrypt(&moveContext, buf, readSize);
    COMMON_LOGGER_RETURN_VOID_IF_FAILED(m_Result);

    // 1. StartDeviceMoveAsDestination
    m_Result = StartDeviceMoveAsDestination(moveContext);
    COMMON_LOGGER_RETURN_VOID_IF_FAILED(m_Result);
}

void SaveDataMover::ImportSystemSaveData(void* buf, size_t bufSize, u64* progress)
{
    // セーブデータが格納されているディレクトリを開く
    nn::fs::Directory root;
    m_Result = root.TryInitialize(SD_SAVEDATA_SYS_ROOT_PATH);
    COMMON_LOGGER_RETURN_VOID_IF_FAILED(m_Result);
    nn::fs::DirectoryEntry entry;
    s32 numOut = 0;

    // ディレクトリ内を見てセーブデータをインポートする。
    while(1)
    {
        m_Result = root.TryRead(&numOut, &entry, 1);
        COMMON_LOGGER_RETURN_VOID_IF_FAILED(m_Result);
        if(numOut == 0)
        {
            break;
        }

        bit32 id = HexStringToBit32(entry.shortName.body, 8);
        if(id == 0)
        {
            break;
        }

        NN_LOG("id: %08x\n", id);

        nn::fs::FileInputStream input;
        char name[64];
        nn::nstd::TSNPrintf(name, 64, "%s/%s", SD_SAVEDATA_SYS_ROOT_PATH, entry.shortName.body);
        m_Result = input.TryInitialize(name);
        if(m_Result.IsFailure())
        {
            COMMON_LOGGER_WARN("Can't Import: %08x\n", id);
            COMMON_LOGGER_RETURN_VOID_IF_FAILED(m_Result);
            continue;
        }

        // 1. fs::CreateAndOpenNewSystemSaveDataRawStorageFile を呼び
        //    ID に対する FileOutputStream を得る。
        nn::fs::FileOutputStream output;
        m_Result = nn::fs::CreateAndOpenNewSystemSaveDataRawStorageFile(&output, id, input.GetSize());
        if(m_Result.IsFailure())
        {
            COMMON_LOGGER_WARN("Can't Import: %08x\n", id);
            COMMON_LOGGER_RETURN_VOID_IF_FAILED(m_Result);
            continue;
        }


        // 2. 1 で得たストリームに対して書き込みを行う。
        //    本来は引っ越し元からデータを受け取り、書き込むが
        //    ここでは SD から読み込んだデータを書き込む。
        m_Result = CopyFile(input, output, buf, bufSize, progress);
        if(m_Result.IsFailure())
        {
            COMMON_LOGGER_WARN("Can't Import: %08x\n", id);
            COMMON_LOGGER_RETURN_VOID_IF_FAILED(m_Result);
            continue;
        }
    }
    NN_LOG("\n");
}

void SaveDataMover::ImportSharedExtSaveData(void* buf, size_t bufSize, u64* progress)
{
    nn::fs::Directory srcRoot;
    m_Result = srcRoot.TryInitialize(SD_SAVEDATA_EXT_ROOT_PATH);
    COMMON_LOGGER_RETURN_VOID_IF_FAILED(m_Result);

    nn::fs::DirectoryEntry entry;
    s32 numEntry;
    while(1)
    {
        m_Result = srcRoot.TryRead(&numEntry, &entry, 1);
        COMMON_LOGGER_RETURN_VOID_IF_FAILED(m_Result);
        if(numEntry == 0)
        {
            break;
        }
        bit32 id = HexStringToBit32(entry.shortName.body, 8);
        NN_LOG("Import: %08x\n", id);

        m_Result = Import(buf, bufSize, id, progress);
        if(m_Result.IsFailure())
        {
            NN_LOG("Failed to import. result = %08x\n", m_Result.GetPrintableBits());
        }

    }
    NN_LOG("\n");
}

nn::Result SaveDataMover::Import(void* buf, size_t bufSize, bit32 id, u64* progress)
{
    const char* ARC_NAME = "sext:";

    const int MAX_PATH_LEN = 127;
    char srcRoot[MAX_PATH_LEN + 1];

    COMMON_LOGGER_RETURN_RESULT_IF_FAILED(nn::fs::CreateNewSharedExtSaveDataRawStorage(id));

    // 2. fs::CreateNewSharedExtSaveDataRawStorage を呼び
    //    書き込み用の領域を作成します。
    nn::nstd::TSNPrintf(srcRoot, MAX_PATH_LEN, "%s/%08x", SD_SAVEDATA_EXT_ROOT_PATH, id);

    m_Result = nn::fs::MountNewSharedExtSaveDataRawStorage(ARC_NAME, id);
    COMMON_LOGGER_RETURN_RESULT_IF_FAILED(m_Result);

    // 4. 3 でマウントした領域に、引っ越し元のアーカイブと
    //    同じ構造でファイルとディレクトリを置く。
    m_Result = CopyRecursivly(buf, bufSize, srcRoot, ARC_NAME, progress);
    COMMON_LOGGER_RETURN_RESULT_IF_FAILED(m_Result);

    m_Result = nn::fs::Unmount(ARC_NAME);
    COMMON_LOGGER_RETURN_RESULT_IF_FAILED(m_Result);

    return nn::ResultSuccess();
}

void SaveDataMover::CalculateImportFileSize()
{
    m_Result = common::CalculateFileSizeRecursively(
            (::std::wstring(common::SDMC_ROOT_DIRECTORY_PATH) + ::std::wstring(common::SD_SAVEDATA_ROOT_NAME)).c_str(),
            m_TotalSize);
    COMMON_LOGGER_RETURN_VOID_IF_FAILED(m_Result);
}

// ファイルの読み書き
nn::Result SaveDataMover::CopyFile(nn::fs::FileInputStream& is, nn::fs::FileOutputStream& os, void* pBuffer,
        const s32 bufferSize, u64* progress)
{
    s64 restSize;
    COMMON_LOGGER_RETURN_RESULT_IF_FAILED(is.TryGetSize(&restSize));

    while (restSize > 0)
    {
        s32 readSize = restSize > bufferSize ? bufferSize : restSize;
        s32 read = 0;
        s32 write = 0;
        COMMON_LOGGER_RETURN_RESULT_IF_FAILED(is.TryRead(&read, pBuffer, readSize));
        COMMON_LOGGER_RETURN_RESULT_IF_FAILED(os.TryWrite(&write, pBuffer, read, true));

        restSize -= read;
        m_FinishedSize += read;
        *progress = m_FinishedSize * 100 / m_TotalSize;
        NN_LOG("FinishedSize = %lld\n", m_FinishedSize);
        NN_LOG("progress = %lld\n", *progress);
    }
    return nn::ResultSuccess();
}

// ディレクトリ内を再帰的にコピー
// INFO: 本当は再帰を使わないで書く
nn::Result SaveDataMover::CopyRecursivly(void* buf, size_t bufSize, const char* src, const char* dest, u64* progress)
{
    char srcPath[128];
    char destPath[128];
    char entryName[128];

    nn::nstd::TSNPrintf(srcPath, sizeof(srcPath), "%s/", src);
    nn::fs::Directory srcDir;
    COMMON_LOGGER_RETURN_RESULT_IF_FAILED(srcDir.TryInitialize(srcPath));

    while (1)
    {
        nn::fs::DirectoryEntry entry;
        s32 numRead;
        COMMON_LOGGER_RETURN_RESULT_IF_FAILED(srcDir.TryRead(&numRead, &entry, 1));
        if (numRead == 0)
        {
            break;
        }

        std::wcstombs(entryName, entry.entryName, sizeof(entryName) - 1);
        entryName[sizeof(entryName) - 1] = '\0';
        nn::nstd::TSNPrintf(srcPath, sizeof(srcPath), "%s/%s", src, entryName);
        nn::nstd::TSNPrintf(destPath, sizeof(destPath), "%s/%s", dest, entryName);

        if (entry.attributes.isDirectory)
        {
            m_Result = nn::fs::TryCreateDirectory(destPath);
            if (m_Result.IsFailure())
            {
                if (m_Result.GetDescription() == nn::fs::DESCRIPTION_FAT_ALREADY_EXISTS
                        && m_Result.GetModule() == nn::Result::MODULE_NN_FS)
                {
                }
                else
                {
                    return m_Result;
                }
            }
            m_Result = CopyRecursivly(buf, bufSize, srcPath, destPath, progress);
            COMMON_LOGGER_RETURN_RESULT_IF_FAILED(m_Result);
        }
        else
        {
            nn::fs::FileInputStream input;
            COMMON_LOGGER_RETURN_RESULT_IF_FAILED(input.TryInitialize(srcPath));

            nn::fs::FileOutputStream output;
            COMMON_LOGGER_RETURN_RESULT_IF_FAILED(output.TryInitialize(destPath, true));

            CopyFile(input, output, buf, bufSize, progress);
        }
    }
    return nn::ResultSuccess();
}

} /* namespace common */