/*
*   This file is part of Luma3DS
*   Copyright (C) 2016 Aurora Wright, TuxSH
*
*   This program is free software: you can redistribute it and/or modify
*   it under the terms of the GNU General Public License as published by
*   the Free Software Foundation, either version 3 of the License, or
*   (at your option) any later version.
*
*   This program is distributed in the hope that it will be useful,
*   but WITHOUT ANY WARRANTY; without even the implied warranty of
*   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
*   GNU General Public License for more details.
*
*   You should have received a copy of the GNU General Public License
*   along with this program.  If not, see <http://www.gnu.org/licenses/>.
*
*   Additional Terms 7.b of GPLv3 applies to this file: Requiring preservation of specified
*   reasonable legal notices or author attributions in that material or in the Appropriate Legal
*   Notices displayed by works containing it.
*/

#include "fs.h"
#include "memory.h"
#include "strings.h"
#include "fmt.h"
#include "crypto.h"
#include "cache.h"
#include "screen.h"
#include "draw.h"
#include "utils.h"
#include "config.h"
#include "fatfs/ff.h"
#include "buttons.h"
#include "firm.h"
#include "crypto.h"
#include "../build/bundled.h"

static FATFS sdFs,
             nandFs;

static bool switchToMainDir(bool isSd)
{
    const char *mainDir = isSd ? "/luma" : "/rw/luma";

    switch(f_chdir(mainDir))
    {
        case FR_OK:
            return true;
        case FR_NO_PATH:
            f_mkdir(mainDir);
            return switchToMainDir(isSd);
        default:
            return false;
    }
}

bool mountFs(bool isSd, bool switchToCtrNand)
{
    return isSd ? f_mount(&sdFs, "0:", 1) == FR_OK && switchToMainDir(true) :
                  f_mount(&nandFs, "1:", 1) == FR_OK && (!switchToCtrNand || (f_chdrive("1:") == FR_OK && switchToMainDir(false)));
}

u32 fileRead(void *dest, const char *path, u32 maxSize)
{
    FIL file;
    u32 ret = 0;

    if(f_open(&file, path, FA_READ) != FR_OK) return ret;

    u32 size = f_size(&file);
    if(dest == NULL) ret = size;
    else if(size <= maxSize)
        f_read(&file, dest, size, (unsigned int *)&ret);
    f_close(&file);

    return ret;
}

u32 getFileSize(const char *path)
{
    return fileRead(NULL, path, 0);
}

bool fileWrite(const void *buffer, const char *path, u32 size)
{
    FIL file;

    switch(f_open(&file, path, FA_WRITE | FA_OPEN_ALWAYS))
    {
        case FR_OK:
        {
            unsigned int written;
            f_write(&file, buffer, size, &written);
            f_truncate(&file);
            f_close(&file);

            return (u32)written == size;
        }
        case FR_NO_PATH:
            for(u32 i = 1; path[i] != 0; i++)
                if(path[i] == '/')
                {
                    char folder[i + 1];
                    memcpy(folder, path, i);
                    folder[i] = 0;
                    f_mkdir(folder);
                }

            return fileWrite(buffer, path, size);
        default:
            return false;
    }
}

void fileDelete(const char *path)
{
    f_unlink(path);
}

static __attribute__((noinline)) bool overlaps(u32 as, u32 ae, u32 bs, u32 be)
{
    if (as <= bs && bs <= ae)
        return true;
    else if (bs <= as && as <= be)
        return true;
    return false;
}

static bool checkFirmPayload(void)
{
    if(memcmp(firm->magic, "FIRM", 4) != 0)
        return false;

    if(firm->arm9Entry == NULL)  //allow for the arm11 entrypoint to be zero in which case nothing is done on the arm11 side
        return false;

    u32 size = 0x200;
    for(u32 i = 0; i < 4; i++)
        size += firm->section[i].size;

    bool arm9EpFound = false, arm11EpFound = false;
    for(u32 i = 0; i < 4; i++)
    {
        __attribute__((aligned(4))) u8 hash[0x20];

        FirmSection *section = &firm->section[i];

        // allow empty sections
        if (section->size == 0)
            continue;

        if(section->offset < 0x200)
            return false;

        if(section->address + section->size < section->address) //overflow check
            return false;

        if(((u32)section->address & 3) || (section->offset & 0x1FF) || (section->size & 0x1FF)) //alignment check
            return false;

        if(overlaps((u32)section->address, (u32)section->address + section->size, 0x27FFE000, 0x28000000))
            return false;
        else if(overlaps((u32)section->address, (u32)section->address + section->size, 0x27FFE000 - 0x1000, 0x27FFE000))
            return false;
        else if(overlaps((u32)section->address, (u32)section->address + section->size, (u32)firm, (u32)firm + size))
            return false;

        sha(hash, (u8 *)firm + section->offset, section->size, SHA_256_MODE);
        if(memcmp(hash, section->hash, 0x20) != 0)
            return false;

        if(firm->arm9Entry >= section->address && firm->arm9Entry < (section->address + section->size))
            arm9EpFound = true;

        if(firm->arm11Entry >= section->address && firm->arm11Entry < (section->address + section->size))
            arm11EpFound = true;
    }

    return arm9EpFound && (firm->arm11Entry == NULL || arm11EpFound);
}

void loadPayload(u32 pressed, const char *payloadPath)
{
    u32 *loaderAddress = (u32 *)0x27FFE000;
    u32 payloadSize = 0,
        maxPayloadSize = (u32)((u8 *)loaderAddress - (u8 *)firm);

    char absPath[24 + _MAX_LFN] = {0};
    char path[10 + _MAX_LFN] = {0};

    if(payloadPath == NULL)
    {
        const char *pattern;

        if(pressed & BUTTON_LEFT) pattern = PATTERN("left");
        else if(pressed & BUTTON_RIGHT) pattern = PATTERN("right");
        else if(pressed & BUTTON_UP) pattern = PATTERN("up");
        else if(pressed & BUTTON_DOWN) pattern = PATTERN("down");
        else if(pressed & BUTTON_START) pattern = PATTERN("start");
        else if(pressed & BUTTON_B) pattern = PATTERN("b");
        else if(pressed & BUTTON_X) pattern = PATTERN("x");
        else if(pressed & BUTTON_Y) pattern = PATTERN("y");
        else if(pressed & BUTTON_R1) pattern = PATTERN("r");
        else if(pressed & BUTTON_A) pattern = PATTERN("a");
        else pattern = PATTERN("select");

        DIR dir;
        FILINFO info;
        FRESULT result;

        result = f_findfirst(&dir, &info, "payloads", pattern);

        if(result != FR_OK) return;

        f_closedir(&dir);

        if(!info.fname[0]) return;

        sprintf(path, "payloads/%s", info.fname);

    }
    else sprintf(path, "%s", payloadPath);

    payloadSize = fileRead(firm, path, maxPayloadSize);

    if(!payloadSize || !checkFirmPayload()) return;

    writeConfig(true);

    if(memcmp(launchedPath, u"nand", 8))
        sprintf(absPath, "nand:/rw/luma/%s", path);
    else
        sprintf(absPath, "sdmc:/luma/%s", path);

    char *argv[1] = {absPath};
    memcpy(loaderAddress, loader_bin, loader_bin_size);

    initScreens();

    flushDCacheRange(loaderAddress, loader_bin_size);
    flushICacheRange(loaderAddress, loader_bin_size);

    ((void (*)(int, char **, u32))loaderAddress)(1, argv, 0x0000BEEF);
}

void payloadMenu(void)
{
    DIR dir;
    char path[62] = "payloads";

    if(f_opendir(&dir, path) != FR_OK) return;

    FILINFO info;
    u32 payloadNum = 0;
    char payloadList[20][49];

    while(f_readdir(&dir, &info) == FR_OK && info.fname[0] != 0 && payloadNum < 20)
    {
        if(info.fname[0] == '.') continue;

        u32 nameLength = strlen(info.fname);

        if(nameLength < 6 || nameLength > 52) continue;

        nameLength -= 5;

        if(memcmp(info.fname + nameLength, ".firm", 5) != 0) continue;

        memcpy(payloadList[payloadNum], info.fname, nameLength);
        payloadList[payloadNum][nameLength] = 0;
        payloadNum++;
    }

    f_closedir(&dir);

    if(!payloadNum) return;

    u32 pressed = 0,
        selectedPayload = 0;

    if(payloadNum != 1)
    {
        initScreens();

        drawString(true, 10, 10, COLOR_TITLE, "Luma3DS chainloader");
        drawString(true, 10, 10 + SPACING_Y, COLOR_TITLE, "Press A to select, START to quit");

        for(u32 i = 0, posY = 10 + 3 * SPACING_Y, color = COLOR_RED; i < payloadNum; i++, posY += SPACING_Y)
        {
            drawString(true, 10, posY, color, payloadList[i]);
            if(color == COLOR_RED) color = COLOR_WHITE;
        }

        while(pressed != BUTTON_A && pressed != BUTTON_START)
        {
            do
            {
                pressed = waitInput(true);
            }
            while(!(pressed & MENU_BUTTONS));

            u32 oldSelectedPayload = selectedPayload;

            switch(pressed)
            {
                case BUTTON_UP:
                    selectedPayload = !selectedPayload ? payloadNum - 1 : selectedPayload - 1;
                    break;
                case BUTTON_DOWN:
                    selectedPayload = selectedPayload == payloadNum - 1 ? 0 : selectedPayload + 1;
                    break;
                case BUTTON_LEFT:
                    selectedPayload = 0;
                    break;
                case BUTTON_RIGHT:
                    selectedPayload = payloadNum - 1;
                    break;
                default:
                    continue;
            }

            if(oldSelectedPayload == selectedPayload) continue;

            drawString(true, 10, 10 + (3 + oldSelectedPayload) * SPACING_Y, COLOR_WHITE, payloadList[oldSelectedPayload]);
            drawString(true, 10, 10 + (3 + selectedPayload) * SPACING_Y, COLOR_RED, payloadList[selectedPayload]);
        }
    }

    if(pressed != BUTTON_START)
    {
        sprintf(path, "payloads/%s.firm", payloadList[selectedPayload]);
        loadPayload(0, path);
        error("The payload is too large or corrupted.");
    }

    while(HID_PAD & MENU_BUTTONS);
    wait(2000ULL);
}

u32 firmRead(void *dest, u32 firmType)
{
    const char *firmFolders[][2] = {{"00000002", "20000002"},
                                    {"00000102", "20000102"},
                                    {"00000202", "20000202"},
                                    {"00000003", "20000003"},
                                    {"00000001", "20000001"}};

    char folderPath[35],
         path[48];

    sprintf(folderPath, "1:/title/00040138/%s/content", firmFolders[firmType][ISN3DS ? 1 : 0]);

    DIR dir;
    u32 firmVersion = 0xFFFFFFFF;

    if(f_opendir(&dir, folderPath) != FR_OK) goto exit;

    FILINFO info;

    //Parse the target directory
    while(f_readdir(&dir, &info) == FR_OK && info.fname[0] != 0)
    {
        //Not a cxi
        if(info.fname[9] != 'a' || strlen(info.fname) != 12) continue;

        u32 tempVersion = hexAtoi(info.altname, 8);

        //Found an older cxi
        if(tempVersion < firmVersion) firmVersion = tempVersion;
    }

    f_closedir(&dir);

    if(firmVersion == 0xFFFFFFFF) goto exit;

    //Complete the string with the .app name
    sprintf(path, "%s/%08x.app", folderPath, firmVersion);

    if(fileRead(dest, path, 0x400000 + sizeof(Cxi) + 0x200) <= sizeof(Cxi) + 0x200) firmVersion = 0xFFFFFFFF;

exit:
    return firmVersion;
}

void findDumpFile(const char *folderPath, char *fileName)
{
    DIR dir;
    FRESULT result;

    for(u32 n = 0; n <= 99999999; n++)
    {
        FILINFO info;

        sprintf(fileName, "crash_dump_%08u.dmp", n);
        result = f_findfirst(&dir, &info, folderPath, fileName);

        if(result != FR_OK || !info.fname[0]) break;
    }

    if(result == FR_OK) f_closedir(&dir);
}