/***************************************************************************
*             __________               __   ___.
*   Open      \______   \ ____   ____ |  | _\_ |__   _______  ___
*   Source     |       _//  _ \_/ ___\|  |/ /| __ \ /  _ \  \/  /
*   Jukebox    |    |   (  <_> )  \___|    < | \_\ (  <_> > <  <
*   Firmware   |____|_  /\____/ \___  >__|_ \|___  /\____/__/\_ \
*                     \/            \/     \/    \/            \/
* $Id: $
*
* Plugin for reprogramming portions of the flash ROM chip with a new content.
* !!! BE CAREFUL THIS CAN DAMAGE YOUR PLAYER  !!!
*
* Copyright (C) 2007,2008 Karl Kurbjun
*
* 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 2
* of the License, or (at your option) any later version.
*
* This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY
* KIND, either express or implied.
*
****************************************************************************/

#include "plugin.h"
#include "s3c2440.h"
#include "lib/md5.h"
#include "lib/helper.h"

PLUGIN_HEADER

/* This define disables some of the error checking to allow untested bootloaders
 * to be written to flash.
 */
#if 0
#define ENABLE_DEV
#endif

#ifndef UINT8
#define UINT8 unsigned char
#endif

#ifndef UINT16
#define UINT16 unsigned short
#endif

#ifndef UINT32
#define UINT32 unsigned int
#endif

#define DELAY(timeout,comp,p) ({while(*p!=comp && --timeout);})

static bool do_flash_dump(void);

/* Take advantage of the fact that the flash repeats in memory up to 0x10000000
 * Ram is mapped from 0x00000000 to 0x02000000, so use 0x04000000.
 */
static volatile UINT8* FB = (UINT8*)0x04000000; /* Flash base address */

UINT8 sector_buffer[0x10000];

void mask_int(bool mask)
{
    static unsigned int oldmask;
    if(mask==true&&INTMSK!=-1)
    {
        oldmask=INTMSK;
        INTMSK= -1;
    }
    else if(mask==false)
    {
        INTMSK=oldmask;
    }
}

inline void FlashReset(void)
{
    FB[0x0] = 0xF0; /* reset */
}

void splash_pause(bool pause_key, const char *format, ...)
{
    static char printfbuf[256];
    static char presskey[]=" [Press a button to continue]";
    int len;
    va_list ap;
    va_start(ap, format);

    len = rb->vsnprintf(printfbuf, sizeof(printfbuf), format, ap);
    va_end(ap);
    
    if(pause_key)
    {
        if(rb->strlen(printfbuf)+rb->strlen(presskey)<sizeof(printfbuf))
        {
            rb->strcat(printfbuf, presskey);
        }
    }

    rb->lcd_clear_display();
    rb->splashf(HZ, printfbuf);
    
    if(pause_key)
    {
        /* Clear the button queue */
        while (rb->button_get(false) != BUTTON_NONE)
        {
            rb->yield();
        }

        while(true)
        {
            rb->button_get(true);
            break;
        }
    }
}

/* erase the sector which contains the given address */
bool EraseFlashSector(volatile UINT32 iAddr)
{
    mask_int(true);
    unsigned int timeout = 375000000; 
    /* the timeout loop should be no less than 15s */

    FB[0xAAA] = 0xAA; /* enter command mode */
    FB[0x555] = 0x55;
    FB[0xAAA] = 0x80; /* erase command */
    FB[0xAAA] = 0xAA; /* enter command mode */
    FB[0x555] = 0x55;
    FB[iAddr] = 0x30; /* erase the sector */

    DELAY(timeout,0xFF,&FB[iAddr]);

    mask_int(false);
    return (timeout == 0);
}

/* erase the sector which contains the given address */
bool EraseFlash(void)
{
    mask_int(true);
    unsigned timeout = 0xFFFFFFFF; 
    /* timeout loop should be no less than 200s; at 171s now (actually more) */

    FB[0xAAA] = 0xAA; /* enter command mode */
    FB[0x555] = 0x55;
    FB[0xAAA] = 0x80; /* erase command */
    FB[0xAAA] = 0xAA; /* enter command mode */
    FB[0x555] = 0x55;
    FB[0xAAA] = 0x10; /* erase the device */

    DELAY(timeout,0xFF,FB);

    mask_int(false);
    return (timeout == 0);
}

inline bool SetupFastWrite(void)
{
    FB[0xAAA] = 0xAA; /* enter command mode */
    FB[0x554] = 0x55;
    FB[0xAAA] = 0x20; /* fast write mode command */

    return false;
}

inline bool ExitFastWrite(void)
{
    FB[0xAAA] = 0x90; /* enter command mode */
    FB[0x554] = 0xF0;

    return false;
}

inline bool FProgramWord(volatile UINT16* pAddr, UINT16 data)
{
    unsigned timeout = 25000; /* the timeout loop should be no less than 1ms */

    if (~*pAddr & data) /* just a safety feature, not really necessary */
        return false; /* can't set any bit from 0 to 1 */

    FB[0xAAA] = 0xA0; /* Fast write mode */

    *pAddr = data;

    DELAY(timeout,data,pAddr);
    
    return (timeout == 0);
}

static bool write_block(volatile UINT16* pAddr, UINT16* data, UINT32 length)
{
    bool retval=false;
    unsigned int i;
    SetupFastWrite();
    for(i=0;i<length;i++)
    {
        if(FProgramWord(pAddr+i, data[i]))
        {
            retval=true;
        }
    }

    ExitFastWrite();
    return retval;
}

static bool write_file(volatile UINT16* pAddr, char name[], UINT32 length)
{
    int fd;
    bool retval=false;
    UINT32 i, intread, amt;
    
    fd = rb->open(name, O_RDONLY);

    if(fd < 0) /* no options to read, set defaults */
    {
        splash_pause(true, "ABORT!: %s could not be opened", name);
        return true;
    }

    SetupFastWrite();
    for(i=0;i<length;i+=2)
    {
        amt=rb->read(fd,&intread, sizeof(UINT32));
        if(amt<=0)
            break;
        if(FProgramWord(pAddr+i, (UINT16)(intread)))
        {
            retval=true;
        }
        if(FProgramWord((UINT16*)(pAddr+i+1), (UINT16)(intread>>16)))
        {
            retval=true;
        }
    }
    ExitFastWrite();
    rb->close(fd);
    return retval;
}

static bool md5sum_calc(char *fname, UINT32 *ret_digest)
{
    char buffer[512];
    ssize_t len;
    int fd;
    struct md5_s md5;
    
    InitMD5( &md5 );
    if(fname!=0x0)
    {
        fd = rb->open( fname, O_RDONLY );
        if( fd < 0 ) 
        {
            return true;
        }
        while( ( len = rb->read( fd, buffer, 512 ) ) > 0 )
        {
            AddMD5( &md5, buffer, len );
        }
        rb->close( fd );
    }
    else
    {
        AddMD5( &md5, (char *)FB, 0x100000 );
    }
    
    EndMD5( &md5 );

    for(len=0; len<4; len++)
    {
        ret_digest[len]=md5.p_digest[len];
    }
    return false;
}

static int md5sum_lookup(UINT32 *calc_digest, char *vstring, char *rstring)
{
    /* Most likely any flash that starts with 0xEA00001A will work
     * (that is FB[0]=0x1A, FB[1]=0x00, FB[2]=0x00, and FB[3]=0xEA). Unknown flash
     * images should still be tested to verify functionality though.
     *
     * This array is for the full 1MB images (file or flash)
     * The format for this array is as follows:
     *  full_digest[x][0] == Firmware version number, two major, two minor 
     *      (upper 16 bits reserved)
     *  full_digest[x][1] == Bootloader release for the patched images (0x00==OF).
     *      This matches the string in crt0.S.
     *  full_digest[x][2-5] == the calculated digest
     */

    UINT32 full_digest[8][6]=
    {
        /* These digests are for the unpatched OF images v0x00 */
        { (int)"V 02.00\0", (int)"O 02.00\0", 
            0xAE145195, 0xC35D8768, 0xCA46960D, 0x54D34D68 }, 
            
        { (int)"V 03.00\0", (int)"O 03.00\0", 
            0xCE1062D6, 0x4D94D20A, 0x9A212230, 0x5D8EB8A7 },
            
        { (int)"V 03.02\0", (int)"O 03.02\0", 
            0x3EB3F6F0, 0xFEA90DB0, 0xF41D5247, 0x185185E9 },
        
        /* These digests are for the patched OF images v>0x00 */  
        { (int)"V 02.00\0", (int)"R 03.00\0", 
            0x9BBFDE73, 0xF6DDD178, 0x4C3D75BA, 0x6C29AE25 }, /* V2.00 R 03.00*/
        
        { (int)"V 03.00\0", (int)"R 03.00\0", 
            0x3C4AF860, 0xCAFB8BFE, 0x34169358, 0xD56A7568 }, /* V3.00 R 03.00*/
        
        { (int)"V 03.02\0", (int)"R 03.00\0",     
            0x3DA3E71C, 0xE8DC5A65, 0x6ABCBAD9, 0xC14665A0 },
            
        /* These digests are for the bootloader images */
        { (int)"BOOT   \0", (int)"B 03.00\0", /* BOOTLOADER (only) R 03.00 */
            0x4199B276, 0x5F53B70C, 0x31F42A8B, 0x5F02D5C7 }, 
    };
    
    int i, match;

    for(i=0; i<8; i++)
    {
        match=0;
        int j;
        for(j=0; j<4; j++)
        {
            if(full_digest[i][j+2]==calc_digest[j])
            {
                match++;
            }
        }
        if(match==4)
        {
            rb->snprintf(vstring, 8, "%s", (char *)full_digest[i][0]);
            rb->snprintf(rstring, 8, "%s", (char *)full_digest[i][1]);
            if( *((char *)full_digest[i][1])=='O' )
            {
                return 0; /* Unpatched OF */
            }
            else if( *((char *)full_digest[i][1])=='R' )
            {
                return 1; /* Patched OF w/ known BL*/
            }
            else
            {
                return 2; /* Raw BL image (no OF) */
            }
        }
    }
    return -1; /* No matched found (untested/unknown)*/
}

static bool md5sum_write(char *fname, UINT32 *calc_digest)
{
    int fd;
    fd = rb->open(fname, O_WRONLY| O_CREAT);

    if(fd < 0)
    {
        splash_pause(true, "ABORT: Unable to open %s!", fname);
        return true;
    }

    rb->fdprintf(fd,"{ 0x%8x, 0x%8x, 0x%8x, 0x%8x }\n\n",
        calc_digest[0], calc_digest[1], calc_digest[2], calc_digest[3]);
    rb->close(fd);
    return false;
}

static int md5sum_testflash(void)
{
    UINT32 calc_digest[4]={0xFFFFFFFF,0xFFFFFFFF,0xFFFFFFFF,0xFFFFFFFF};
    char vstring[8], rstring[8];
    int check;
    
    splash_pause(false, "INFO: Checking to for tested image version");

    /* This should never error, but check just to be sure */
    if(md5sum_calc(0x0, calc_digest))
    {
        splash_pause(true, "ABORT: The MD5sum was not able to be calculated!");
        return -2;
    }
       
    check=md5sum_lookup(calc_digest, vstring, rstring);
    
    /* Check should never equal 2 (possible collision) */
    if(check<0 || check==2)
    {
        splash_pause(true, "ABORT: This is an untested flash image. Please run "
            "\"Backup Flash and MD5sum\" and post the backup.bin and "
            "backup.md5 files to FS#7505 or create a new message on the "
            "rockbox forums.");
    }
    else if(check==1)
    {
        splash_pause(false, "INFO: OK, this is a patched flash image %s : %s", 
            vstring, rstring);
    }
    else
    {
        splash_pause(false, "INFO: OK, This is an unpatched flash "
            "image: %s : %s", vstring, rstring);
    }

    return check;
}

static bool md5sum_testfile(char *fname)
{
    UINT32 calc_digest[4]={0xFFFFFFFF,0xFFFFFFFF,0xFFFFFFFF,0xFFFFFFFF};
    char vstring[8], rstring[8];
    
    splash_pause(false, "INFO: Checking to for tested image version");
    
    if(md5sum_calc(fname, calc_digest))
    {
        splash_pause(true, "ABORT: Error opening %s", fname);
        return true;
    }
    else if(md5sum_lookup(calc_digest, vstring, rstring)<0)
    {
        splash_pause(true, "ABORT: Unknown/Untested %s", fname);
        return true;
    }

    splash_pause(false, "INFO: OK, %s is a tested image: %s : %s", fname, 
        vstring, rstring);

    return false;
}

static bool md5sum_bootloader(void)
{
    UINT32 calc_digest[4]={0xFFFFFFFF,0xFFFFFFFF,0xFFFFFFFF,0xFFFFFFFF};
    
    splash_pause(false, "INFO: Calculating MD5 hash and storing bootloader.md5");
    
    md5sum_calc("/bootloader.bin", calc_digest);
    
    if(md5sum_write("/bootloader.md5", calc_digest))
    {
        return true;
    }
    return false;
}

static bool md5sum_flash(void)
{
    UINT32 calc_digest[4]={0xFFFFFFFF,0xFFFFFFFF,0xFFFFFFFF,0xFFFFFFFF};
    
    splash_pause(false, "INFO: Calculating MD5 hash and storing flash.md5");
    
    md5sum_calc(0x0, calc_digest);
    
    if(md5sum_write("/flash.md5", calc_digest))
    {
        return true;
    }
    return false;
}

static bool do_update_flash(void)
{
    int size, fd;

    FlashReset();

    if(!rb->file_exists("/backup.bin"))
    {
        splash_pause(false, "INFO: Backup image does not exist. Forcing backup.");
        if(do_flash_dump())
        {
            return true;
        }
    }

    if(md5sum_testfile("/bootloader.bin"))
    {
        md5sum_bootloader();
#if !defined(ENABLE_DEV)
        return true;
#endif
    }

    fd = rb->open("/bootloader.bin", O_RDONLY);

    if(fd < 0)
    {
        splash_pause(true, "ABORT: Unable to open /bootloader.bin");
        return true;
    }

    size=rb->filesize(fd);
    if(size>0x10000)
    {
        splash_pause(true, "ABORT: Bootloader image too large: 0x%x > 0x10000", 
            size);
        rb->close(fd);
        return true;
    }
    rb->close(fd);

    splash_pause(false, "INFO: Erasing sector 10 (0xA0000)");
    if(EraseFlashSector(0xA0000))
    {
        splash_pause(true, "ABORT: Sector 10 (0xA0000) erasure failed!");
        return true;
    }

    splash_pause(false, "INFO: Writing bootloader");
    /* The bootloader is placed at 0xA0000 since this location contains only
     * graphics which are not critical to the player's functionality.
     */
    if(write_file((UINT16*)(FB+0xA0000), "/bootloader.bin", 0x10000))
    {
        splash_pause(true, "ABORT: Bootloader programming failed!");
        return true;
    }

    splash_pause(false, "INFO: Bootloader installed");

    return 0;
}

static bool do_flash_restore(void)
{
    int fd, size;
        
    if(!rb->charger_inserted() || (rb->battery_level()<75))
    {
        splash_pause(true, "ABORT: Please insert the charger before running "
            "and charge the battery to at least 75 percent. Battery is "
            "currently %d percent.", rb->battery_level());
        return true;
    }
    
    splash_pause(false, "INFO: Restoring flash from /backup.bin");

    if(!rb->file_exists("/backup.bin"))
    {
        splash_pause(true, "ABORT: Backup image does not exist!");
        return true;
    }
    
    if(md5sum_testfile("/backup.bin"))
    {
        return true;
    }

    fd = rb->open("/backup.bin", O_RDONLY);

    if(fd < 0)
    {
        splash_pause(true, "ABORT: Unable to open backup image");
        return true;
    }

    size=rb->filesize(fd);
    if(size!= 0x100000)
    {
        splash_pause(true, "ABORT: Backup image too small: 0x%x bytes, "
            "expected 0x100000 bytes", size);
        rb->close(fd);
        return true;
    }

    splash_pause(false, "INFO: Erasing Flash");
    if(EraseFlash())
    {
        splash_pause(true, "ABORT: Flash erasure failed!");
        rb->close(fd);
        return true;
    }

    splash_pause(false, "INFO: Programming flash");    
    if(write_file((UINT16*)FB, "/backup.bin", 0x100000))
    {
        splash_pause(true, "ABORT: Programming failed!");
    }

    if(md5sum_testflash()>=2)
    {
        splash_pause(true, "POTENTIAL FATAL ERROR: Backup Restore Failed! DO "
            "NOT TURN OFF THE PLAYER! LEAVE CONNECTED TO CHARGER! Please go to "
            "#rockbox on irc.freenode.org or create a message on the rockbox "
            "forums.");
    }
    else
    {
        splash_pause(false, "INFO: OK, backup restored.");
    }

    rb->close(fd);
    return false;
}

static bool do_write_flash(void)
{
    int i, md5ret;
    
    if(!rb->charger_inserted() || (rb->battery_level()<75))
    {
        splash_pause(true, "ABORT: Please insert the charger before running "
            "and charge the battery to at least 75 percent. Battery is "
            "currently %d percent.", rb->battery_level());
        return true;
    }

    md5ret=md5sum_testflash();
    if(md5ret<0)
    {
#if !defined(ENABLE_DEV)
        return true; /* Not a tested flash image, get out of here */
#endif
    }

    if(do_update_flash()) /* Install/update the Rockbox bootloader */
    {
        return true;
    }

    if(md5ret==1)   /* Sector 0 was patched so we don't need to repatch it */
    {
        splash_pause(false, "INFO: Sector 0 already patched from previous "
            "install");
    }
    else
    {
        splash_pause(false, "INFO: Reading sector 0 (0x0)");
        for(i=0; i<0x10000; i++)
        {
            sector_buffer[i]=FB[i];
        }
        
        splash_pause(false, "INFO: Patching reset vector and erasing sector 0 "
            "(0x0)");
        
        /* This points the reset vector to 0xA0000 instead of 0x70 */
        sector_buffer[0]=0x0A;
        sector_buffer[1]=0xF8;
        sector_buffer[2]=0xA0;
        sector_buffer[3]=0xE3;

        if(EraseFlashSector(0x00000))
        {
            splash_pause(true, "ABORT: Sector 0 (0x0) erasure failed!");
            return true;
        }

        splash_pause(false, "INFO: Writing patched sector 0 (0x0)");
        if(write_block((UINT16*)FB, (UINT16*)sector_buffer, 0x10000/2))
        {
            splash_pause(true, "ABORT: Writing patched sector 0 failed!");
            return true;
        }
    }
    
    if(md5sum_testflash()<0)
    {
        splash_pause(true, "ERROR: Bootloader improperly installed");
        md5sum_flash();
#if !defined(ENABLE_DEV)
        splash_pause(true, "ERROR: Forcing backup restoration!");
        do_flash_restore();
#endif
    }
    else
    {
        splash_pause(true, " INFO: OK, Reset vector patched and bootloader "
            "sucessfully installed");
    }

    return 0;
}

static bool do_flash_dump(void)
{
    int fd;
    UINT32 calc_digest[4]={0xFFFFFFFF,0xFFFFFFFF,0xFFFFFFFF,0xFFFFFFFF};

    if(rb->file_exists("/backup.bin"))
    {
        splash_pause(true, "INFO: Backup image already exists! Skipping "
            "backup.");
        return true;
    }

    splash_pause(false, "INFO: Saving backup to /backup.bin");

    fd = rb->open("/backup.bin", O_WRONLY| O_CREAT);

    if(fd < 0)
    {
        splash_pause(true, "ABORT: Unable to open file for backup!");
        return true;
    }

    if(rb->write(fd, (char *)FB, 0x100000)!=rb->filesize(fd))
    {
        splash_pause(true, "ABORT: Unable to write complete data to "
            "/backup.bin!");
        rb->close(fd);
        return true;
    }

    rb->close(fd);

    splash_pause(false, "INFO: Saving backup md5 sum to /backup.md5");
    md5sum_calc(0x0, calc_digest);
    
    return md5sum_write("/backup.md5", calc_digest);
}

/*** Main Menu Creation */
MENUITEM_FUNCTION(md5sum_item, MENU_FUNC_USEPARAM, "Check flash MD5 (version)",
    md5sum_testflash, (void*)true, NULL, Icon_NOICON);

MENUITEM_FUNCTION(backup_item, MENU_FUNC_USEPARAM, "Backup flash and MD5 hash",
    do_flash_dump, (void*)true, NULL, Icon_NOICON);

MENUITEM_FUNCTION(restore_item, MENU_FUNC_USEPARAM, "Restore flash backup",
    do_flash_restore, (void*)true, NULL, Icon_NOICON);

MENUITEM_FUNCTION(patchflash_item, MENU_FUNC_USEPARAM, 
    "Install/Update bootloader", do_write_flash, (void*)true, NULL, 
    Icon_NOICON);

MENUITEM_RETURNVALUE(quit_item, "Exit Gigabeat Flashwriter", 
    MENU_SELECTED_EXIT, NULL, Icon_NOICON);

MAKE_MENU(main_menu, "Gigabeat Flashwriter", NULL, Icon_NOICON, &md5sum_item,
    &backup_item, &restore_item, &patchflash_item, &quit_item);
/*** End of Main Menu Creation */

/***************** Plugin Entry Point *****************/

enum plugin_status plugin_start(const void* parameter)
{
    int ret, selection=0;
    (void) parameter;

    backlight_force_on();

    splash_pause(true, "** WARNING ** This software is distributed on an \"AS "
        "IS\" basis, WITHOUT WARRANTY OF ANY KIND, either express or implied.");
        
    splash_pause(true, "There are preventative measures in place, but this "
        "plugin may erase your flash and leave you with a broken player.");
        
    splash_pause(true, "Please pay close attention to the messages shown when "
        "using this tool.");

    ret=rb->do_menu(&main_menu, &selection, NULL, false);
    
    backlight_use_settings();

    return PLUGIN_OK;
}