#include <gccore.h>
#include <malloc.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <wiiuse/wpad.h>

#include "AnyRegionChanger.h"
#include "controller.h"
#include "detect_settings.h"
#include "dopios.h"
#include "Error.h"
#include "fat.h"
#include "gecko.h"
#include "IOSMenu.h"
#include "IOSCheck.h"
#include "IOSPatcher.h"
#include "sha1.h"
#include "Settings.h"
#include "title_install.h"
#include "tools.h"
#include "ticket_dat.h"
#include "tmd_bin.h"
#include "video.h"

#define BOOTMII_VERSION     254
#define CSV_VERSION         2
#define WIIMOD_CSV          "sd:/wad/wiimod/wiimod.csv"
#define ES_ERROR_1028 -1028 // No ticket installed
#define ES_ERROR_1029 -1029 // Installed Ticket/TMD is invalid
#define ES_ERROR_1035 -1035 // Title with a higher version is already installed 
#define ES_ERROR_2011 -2011 // Signature check failed (Needs Fakesign)
s32 __u8Cmp(const void *a, const void *b);

static u8 certs[0xA00] ATTRIBUTE_ALIGN(32);
static const char certs_fs[] ATTRIBUTE_ALIGN(32) = "/sys/cert.sys";
static const char flash_fs[] ATTRIBUTE_ALIGN(32) = "/dev/flash";
static const char usb2_fs[] ATTRIBUTE_ALIGN(32) = "/dev/usb2";
static const char boot2_fs[] ATTRIBUTE_ALIGN(32) = "/dev/boot2";

int get_certs(void) {
    int fd, ret;
    fd = IOS_Open(certs_fs, 1);
    ret = IOS_Read(fd, certs, sizeof(certs));
    if (ret < sizeof(certs)) {
        ret = -1;
    } else {
        IOS_Close(fd);
    }
    return ret;
}

int check_fakesig(void) {
    int ret = 0;
    ret = ES_AddTitleStart((signed_blob*) tmd_bin, tmd_bin_size,
        (signed_blob *) certs, sizeof(certs), NULL, 0);
    if (ret >= 0) ES_AddTitleCancel();
    // we don't care if IOS bitches that our version is to low because we
    // don't even want to install this title
    if (ret == ES_ERROR_1035) ret = 0;
    if (ret == ES_ERROR_1028) ret = 0;
    return ret;
}

void zeroSignature(signed_blob *sig) {
    u8 *psig = (u8*) sig;
    memset(psig + 4, 0, SIGNATURE_SIZE(sig) - 4);
}

int bruteTMD(tmd *ptmd) {
    u16 fill;
    for (fill = 0; fill < 65535; fill++) {
        ptmd->fill3 = fill;
        sha1 hash;
        SHA1((u8 *) ptmd, TMD_SIZE(ptmd), hash);
        if (hash[0] == 0) return 0;
    }
    return -1;
}

void forgeTMD(signed_blob *signedTmd) {
    zeroSignature(signedTmd);
    bruteTMD((tmd*) SIGNATURE_PAYLOAD(signedTmd));
}

int bruteTicket(tik *ticket) {
    u16 fill;
    for (fill = 0; fill < 65535; fill++) {
        ticket->padding = fill;
        sha1 hash;
        SHA1((u8*) ticket, sizeof(ticket), hash);
        if (hash[0] == 0) return 0;
    }
    return -1;
}

void forgeTicket(signed_blob *signedTicket) {
    zeroSignature(signedTicket);
    bruteTicket((tik*) SIGNATURE_PAYLOAD(signedTicket));
}

int check_identify(void) {
    int ret = 0;
    ret = ES_Identify((signed_blob *) certs, sizeof(certs),
        (signed_blob *) tmd_bin, tmd_bin_size, (signed_blob *) ticket_dat,
        ticket_dat_size, NULL);

    // we don't care about invalid signatures here.
    // patched versions return -1017 before checking them.
    if (ret == ES_ERROR_2011) ret = 0;
    if (ret == ES_ERROR_1029) ret = 0;
    return ret;
}

int check_flash(void) {
    int ret;
    ret = IOS_Open(flash_fs, 0);
    if (ret >= 0) IOS_Close(ret);
    return ret;
}

int check_usb2(void) {
    int ret = IOS_Open(usb2_fs, 0);
    if (ret >= 0) IOS_Close(ret);
    return ret;
}

int check_boot2(void) {
    int ret;
    ret = IOS_Open(boot2_fs, 0);
    if (ret >= 0) IOS_Close(ret);
    return ret;
}

int check_nandp(void) {
    int ret;
    ret = IOS_Open("/title/00000001/00000002/content/title.tmd", 1);
    if (ret >= 0) IOS_Close(ret);
    return ret;
}

bool isIOSstub(u8 ios_number) {
    u32 tmd_size = 0;
    tmd *ios_tmd ATTRIBUTE_ALIGN(32);

    ES_GetStoredTMDSize(0x0000000100000000ULL | ios_number, &tmd_size);
    if (!tmd_size) {
        //getting size failed. invalid or fake tmd for sure!
        gprintf("failed to get tmd for ios %d\n",ios_number);
        return true;
    }
    signed_blob *ios_tmd_buf = (signed_blob *) AllocateMemory(tmd_size);
    memset(ios_tmd_buf, 0, tmd_size);

    ES_GetStoredTMD(0x0000000100000000ULL | ios_number, ios_tmd_buf, tmd_size);
    ios_tmd = (tmd *) SIGNATURE_PAYLOAD(ios_tmd_buf);
    free(ios_tmd_buf);
    gprintf("IOS %d is rev %d(0x%x) with tmd size off %u and %u contents\n", ios_number, ios_tmd->title_version, ios_tmd->title_version,
            tmd_size, ios_tmd->num_contents);
    /*Stubs have a few things in common¨:
     - title version : it is mostly 65280 , or even better : in hex the last 2 digits are 0.
     example : IOS 60 rev 6400 = 0x1900 = 00 = stub
     - exception for IOS21 which is active, the tmd size is 592 bytes
     - the stub ios' have 1 app of their own (type 0x1) and 2 shared apps (type 0x8001).

     eventho the 00 check seems to work fine , we'll only use other knowledge as well cause some
     people/applications install an ios with a stub rev >_> ...*/
    u8 Version = ios_tmd->title_version;
#ifdef DEBUG
    gprintf("Version = 0x%x\n", Version);
#endif
    //version now contains the last 2 bytes. as said above, if this is 00, its a stub
    if (Version == 0) {
        if ((ios_tmd->num_contents == 3) && (ios_tmd->contents[0].type == 1
                && ios_tmd->contents[1].type == 0x8001
                && ios_tmd->contents[2].type == 0x8001)) {
            gprintf("IOS %d is a stub\n", ios_number);
            return true;
        } else {
            gprintf("IOS %d is active\n", ios_number);
            return false;
        }
    }gprintf("IOS %d is active\n", ios_number);
    return false;
}

u32 GetTitleVersion(u64 tid) {
    static u32 tmd_size ATTRIBUTE_ALIGN(32);

    s32 r = ES_GetStoredTMDSize(tid, &tmd_size);
#ifdef DEBUG
    printf("ES_GetStoredTMDSize(%llX, %08X):%d\n", tid, (u32) (&tmd_size), r);
#endif
    if (r < 0) return 0;

    signed_blob *TMD = (signed_blob *) AllocateMemory(tmd_size);
    memset(TMD, 0, tmd_size);
    
    r = ES_GetStoredTMD(tid, TMD, tmd_size);
#ifdef DEBUG
    printf("ES_GetStoredTMD(%llX, %08X, %d):%d\n", tid, (u32) (TMD), tmd_size, r);
#endif
    if (r < 0) {
        free(TMD);
        return 0;
    }

    tmd *rTMD = (tmd *) (TMD + (0x140 / sizeof(tmd *)));
    u32 version = rTMD->title_version;
    free(TMD);
    return version;
}

int checkIOSMenu() {
    static int menuItem = 0;
    u32 button = 0;
    bool displayCurIOS = false, skipChecking = false, dumpHashes = false;

    while (!(button & WPAD_BUTTON_B)) {
        initNextBuffer();
        PrintBanner();
        printf("\n\n\n\t\t\t\t%3s Check IOSs\n", (menuItem == 0) ? "-->" : " ");
        printf("\t\t\t\t%3s Skip checking IOSs: %d\n", (menuItem == 1)
            ? "-->" : " ", (skipChecking) ? 1 : 0);
        printf("\t\t\t\t%3s Display current IOS: %d\n", (menuItem == 2)
            ? "-->" : " ", (displayCurIOS) ? 1 : 0);
        printf("\t\t\t\t%3s Output IOS csv: %d\n", (menuItem == 3)
            ? "-->" : " ", theSettings.outputios);
        printf("\t\t\t\t%3s Dump Hashes: %d\n", (menuItem == 4) ? "-->" : " ",
            (dumpHashes) ? 1 : 0);

        Console_SetPosition(ConsoleRows - 5, 0);
        horizontalLineBreak();
        printf("\n\t(%s)/(%s) (%s)/(%s)\t\tChange Selection\n", UP_ARROW,
            DOWN_ARROW, LEFT_ARROW, RIGHT_ARROW);
        printf("\t(A)\t\t\t\t\tSelect\t\t\t\t\t(B)\tBack\n");
        printf("\t(HOME)/GC:(START)\t\tExit");
        flushBuffer();
        button = WaitButtons();
        if (ExitMainThread) return 0;
        if (button & WPAD_BUTTON_HOME) ReturnToLoader();
        if (button & WPAD_BUTTON_A) {
            if (menuItem == 0) {
                PrintBanner();
                checkmyios(displayCurIOS, skipChecking, dumpHashes);
                WPAD_Init();
                PromptAnyKeyContinue();
            }
        }
        if ((button & WPAD_BUTTON_LEFT) || (button & WPAD_BUTTON_RIGHT)) {
            if (menuItem == 1) {
                skipChecking = !skipChecking;
                if (skipChecking) {
                    displayCurIOS = false;
                }
            } else if (menuItem == 2) {
                if (!skipChecking) {
                    displayCurIOS = !displayCurIOS;
                }
            } else if (menuItem == 3) {
                theSettings.outputios = !theSettings.outputios;
                if (!theSettings.outputios) {
                    dumpHashes = false;
                }
            } else if (menuItem == 4) {
                if(theSettings.outputios) {
                    dumpHashes = !dumpHashes;
                }
            }
        }
        if (button & WPAD_BUTTON_DOWN) menuItem++;
        if (button & WPAD_BUTTON_UP) menuItem--;
        
        if (menuItem > 4) menuItem = 0;
        if (menuItem < 0) menuItem = 4;
    }
    return 0;
}

s32 __u8Cmp(const void *a, const void *b) {
    return *(u8 *) a - *(u8 *) b;
}

u8 *get_ioslist(u32 *cnt) {
    u32 i, icnt;
    u8 *ioses = NULL;
    
    ioses = (u8 *) malloc(sizeof(u8) * wiiSettings.iosCount);
    icnt = 0;
    for (i = 3; i < 256; i++) {
        if (ios_found[i] > 0) {
            ioses[icnt++] = i;
        }
    }

    *cnt = wiiSettings.iosCount;
    return ioses;
}

int writebackLog(bool dumpHashes) {
    int i, ret, tmp;
    u8 *list;
    u32 len;
    list = get_ioslist(&len);
    
    ret = SDCard_Init();
    if (ret >= 0) {
        FILE *logFile = fopen(WIIMOD_CSV, "wb");
        if (logFile != NULL) {
            fprintf(logFile, "%d,%d\n", CSV_VERSION, (dumpHashes) ? 1 : 0);
            fprintf(
                logFile,
                "IOS number,version, ,Flash access, FakeSign, ES_Identify, Usb2.0 IOS tree, Boot2 access, NAND Permssions, sysmenuver");
            if(dumpHashes) {
                fprintf( logFile, ",hash 1,hash 2,hash 3,hash 4,hash 5\n");
            } else {
                fprintf( logFile, "\n");
            }
            fprintf(logFile, "%u\n", len);
            for (i = len - 1; i >= 0; i -= 1) {
                tmp = list[i];
                fprintf(logFile, "%d,%d, ,%d,%d,%d,%d,%d,%d,%d", tmp,
                    ios_found[tmp], ios_flash[tmp], ios_fsign[tmp],
                    ios_ident[tmp], ios_usb2m[tmp], ios_boot2[tmp],
                    ios_nandp[tmp], ios_sysmen[tmp]);
                if (dumpHashes) {
                    fprintf(logFile, ",%x,%x,%x,%x,%x\n", ios_hash[tmp][0],
                        ios_hash[tmp][1], ios_hash[tmp][2], ios_hash[tmp][3],
                        ios_hash[tmp][4]);
                } else {
                    fprintf(logFile, "\n");
                }
                printf(".");
            }
            if (wiiSettings.bcVersion > 0) {
                fprintf(logFile, "BC v%u\n", wiiSettings.bcVersion);
            }
            if (wiiSettings.miosVersion > 0) {
                fprintf(logFile, "MIOS v%u\n", wiiSettings.miosVersion);
            }
            fprintf(logFile, "Region %s\n", regions_names[wiiSettings.game]);
            fprintf(logFile, "Hollywood v0x%x Console ID: %d\n",
                wiiSettings.hollywoodVersion, wiiSettings.deviceId);
            fprintf(logFile, "Boot2 v%u\n", wiiSettings.boot2version);
            fprintf(logFile, "System menu %d\n", wiiSettings.sysMenuVer);
            fprintf(logFile, "System menu v%.1lf%c on IOS%d\n",
                wiiSettings.sysMenuNinVersion, wiiSettings.sysMenuRegion,
                wiiSettings.sysMenuIOS);
            fprintf(logFile, "Found %u titles of which %u are IOSs.\n",
                wiiSettings.titleCnt, wiiSettings.iosCount);
            fprintf(logFile, "Wii Mod v%.1lf\n", WII_MOD_VERSION);
            fclose(logFile);
        }
        SDCard_deInit();
        printf("Wrote %s\n", WIIMOD_CSV);
        ret = 0;
    }
    free(list);
    return ret;
}

int checkmyios(bool displayIOS, bool skipcheckingIOSs, bool dumpHashes) {
    int i, ret, tmp;
    u8 *list;
    u32 len, version;
    list = get_ioslist(&len);
    if (len == 0) {
        return 0;
    }

    version = IOS_GetVersion();

    memset(ios_ident, -1, 256 * sizeof(int));
    memset(ios_fsign, -1, 256 * sizeof(int));
    memset(ios_flash, -1, 256 * sizeof(int));
    memset(ios_boot2, -1, 256 * sizeof(int));
    memset(ios_usb2m, -1, 256 * sizeof(int));
    memset(ios_nandp, -1, 256 * sizeof(int));
    memset(ios_sysmen, -1, 256 * sizeof(int));

    get_certs();
    printf("\nChecking IOSs.");
    SpinnerStart();
    for (i = 0; !skipcheckingIOSs && i < len; i += 1) {
        tmp = list[i];
        if (tmp == BOOTMII_VERSION) {
            continue;
        }
        if (ios_found[tmp]) {
            if (ios_isStub[tmp] > 1) {
                printf("\b\nStub %d\n", tmp);
                fflush(stdout);
                ret = -1;
            } else {
                if (displayIOS) {
                    printf("\b\nChecking IOS %d\n", tmp);
                    fflush(stdout);
                    sleep(1);
                }
                ret = ReloadIosNoInit(tmp);
            }
            if (ret >= 0) {
                ios_flash[tmp] = check_flash();
                ios_boot2[tmp] = check_boot2();
                ios_usb2m[tmp] = check_usb2();
                ios_nandp[tmp] = check_nandp();
                ios_fsign[tmp] = check_fakesig();
                ios_sysmen[tmp] = check_sysmenuver();
                ios_ident[tmp] = check_identify();
            }
        }
        if (i % 2 == 0) {
            printf("\b..");
            fflush(stdout);
        }
    }
    free(list);
    ret = ReloadIosNoInit(version);
    SpinnerStop();
    printf("\b.");
    if (theSettings.outputios) {
        writebackLog(dumpHashes);
    }
    return 1;
}

int loadiosfromcsv() {
    int i, tmp2 = 0, ret = -1;
    char in[80];
    u32 len;

    memset(ios_ident, -1, 256 * sizeof(int));
    memset(ios_fsign, -1, 256 * sizeof(int));
    memset(ios_flash, -1, 256 * sizeof(int));
    memset(ios_boot2, -1, 256 * sizeof(int));
    memset(ios_usb2m, -1, 256 * sizeof(int));
    memset(ios_nandp, -1, 256 * sizeof(int));
    memset(ios_sysmen, -1, 256 * sizeof(int));

    ret = SDCard_Init();
    if (ret >= 0) {
        FILE * logFile = fopen(WIIMOD_CSV, "rb");
        if (logFile != NULL) {
            //printf("Loading from csv...\n");
            int ver = 0, hashes = 0;
            fscanf(logFile, "%d,%d\n", &ver, &hashes);
            if (ver != CSV_VERSION) {
                //printf("load failed\n");
                return -1;
            }
            len = (hashes)?19:14;
            for (i = 0; i < len; i += 1) {
                fscanf(logFile, "%s", in);
            }
            fscanf(logFile, "%u", &len);
            for (i = len; i > 0; i -= 1) {
                fscanf(logFile, "%d,", &tmp2);
                if (tmp2 > 2 && tmp2 < 256) {
                    fscanf(logFile, "%d, ,%d,%d,%d,%d,%d,%d,%d",
                        &ios_found[tmp2], &ios_flash[tmp2], &ios_fsign[tmp2],
                        &ios_ident[tmp2], &ios_usb2m[tmp2], &ios_boot2[tmp2],
                        &ios_nandp[tmp2], &ios_sysmen[tmp2]);
                    if (hashes) {
                        fscanf(logFile, ",%x,%x,%x,%x,%x\n", &ver, &ver, &ver,
                            &ver, &ver);
                    } else {
                        fscanf(logFile, "\n");
                    }
                } else {
                    fscanf(logFile, "%d, ,%d,%d,%d,%d,%d,%d,%d", &ver, &ver,
                        &ver, &ver, &ver, &ver, &ver, &ver);
                    if (hashes) {
                        fscanf(logFile, ",%x,%x,%x,%x,%x\n", &ver, &ver, &ver,
                            &ver, &ver);
                    } else {
                        fscanf(logFile, "\n");
                    }
                }
                //printf(".");
            }
            fclose(logFile);
        }
        SDCard_deInit();
    }
    return 0;
}
