/* (C) 2003-2007 Willem Jan Hengeveld <itsme@xs4all.nl>
 * Web: http://www.xs4all.nl/~itsme/
 *      http://wiki.xda-developers.com/
 *
 * $Id: pnewbootloader.cpp 1921 2008-07-31 17:08:31Z itsme $
 */

// todo: 
//    - usage of pflashrom, should not mention pnewbootloader
//    - add auto find offset feature.
//    - add option to try anyway, without patching.
#include <windows.h>
#include "ItsUtils.h"
#include "dllversion.h"

#include "debug.h"
#include "stringutils.h"
#include "args.h"
#include <stdio.h>
#include <string.h>
#include <rapi.h>

#include <map>
#include <algorithm>

typedef std::map<int,DWORD> IntDwordMap;

char *msgFlashStartErase= "FLASH: Start erase address %xh\r\n";

int nVerbose= 0;

// to flash the bootloader, we have to unprotect it:
//

// look where the strings 'FLASH: Start erase address' and 'FLASH: +FlashErase
// Start' are referenced, between these there should be 4 consequetive RAM
// addresses, the lowest is the bootloader address
// these 4 addresses contain: 
//       bootloaderstart, bootloadersize, kernelstart, kernelsize.
//
//  dwAddressList points to a section just after the subroutine that erases flash blocks.
//       DCD  msgFlashStartErase   ; points to string
//       DCD  pdwBootloaderStart   ; == dwBootloaderAddress from struct below.
//       DCD  pdwBootloaderSize
//       DCD  pdwKernelStart
//       DCD  pdwKernelSize

//  these are all consequetive addresses.
//
// dwBootloaderStart pointed to by pdwBootloaderStart to a list in RAM of the following:
//       DCD 0x80000000
// dwBootloaderSize
//       DCD 0x00040000
// dwKernelStart
//       DCD 0x80040000
// dwKernelSize
//       DCD 0x00150000
// dwRomStart
//       DCD 0x80000000
// dwRomSize
//       DCD 0x02000000
//
struct versiondata {
    DWORD dwVersion;
    WCHAR *versionstring;
    DWORD dwAddressList;        // points to pointer to 'FLASH: ...' text
    DWORD dwBootloaderAddress;  // 2nd address from the above list
} versions[]= {
    // xda1 / wallaby
    { 0x80041918, L"3.04.00 ENG", 0x8005FDCC, 0x8C0C60E8   },
    { 0x80055434, L"3.16.52 ENG", 0x80086610, 0x8C0BA204   },
    { 0x80055434, L"3.17.03 ENG", 0x80086648, 0x8C0BA204   },
    { 0x80055414, L"3.18.04 ENG", 0x80086628, 0x8C0BA204   },
    { 0x80055414, L"3.19.01 GER", 0x80086628, 0x8C0BA204   },
    { 0x80055be4, L"4.00.10 ENG", 0x80096F48, 0x8C0D62D8   },
    { 0x80055be4, L"4.00.01 ENG", 0x80096184, 0x8C0D62D8   },
    { 0x80055be4, L"4.00.05 ENG", 0x80096F48, 0x8C0D62D8   },
    { 0x80055be4, L"4.00.16 ENG", 0x80096F48, 0x8C0D62D8   },
    { 0x80054594, L"4.01.16 ENG", 0x80063210, 0x8C0D60F8   },

    // xda2 / himalaya
    { 0x800433B0, L"1.60.00WWE",  0x8008DAF8, 0x901C63D0   },
    { 0x800433B0, L"1.52.00WWE",  0x8008DAF8, 0x901C63D0   },
    { 0x800433B0, L"1.66.00WWE",  0x80091734, 0x901c6f68   },
    { 0x800433B0, L"1.66.01GER",  0x80091734, 0x901c6f68   },
    { 0x800433B0, L"1.72.00WWE",  0x80091744, 0x901C6F68   },
    { 0x800433B0, L"1.75.00WWE",  0x80091748, 0x901C6F68   },
    { 0x800433B0, L"1.75.04ITA",  0x8009174c, 0x901C6F68   },

    // i-mate jam / magician
    { 0x80043670, L"0.68.00 WWE", 0x8008C498, 0x905C63D0   },
    { 0x800435D0, L"1.00.00 WWE", 0x8008B794, 0x906663D4   },
    { 0x800435D0, L"1.03.00 WWE", 0x8008BA74, 0x906663D4   },
    { 0x800435D0, L"1.12.00 WWE", 0x8008BFF4, 0x906663D4   },
};
#define N_VERSIONS (sizeof(versions)/sizeof(*versions))

bool ITFlashROM(DWORD dwType, DWORD dwOffset, BYTE *buffer, DWORD dwBufLength);
bool ITWriteProcessMemory(HANDLE hProc, DWORD dwOffset, BYTE *buffer, DWORD dwBufLength, DWORD *pdwNumberOfBytesWritten);
bool ITReadProcessMemory(HANDLE hProc, DWORD dwOffset, BYTE *buffer, DWORD dwBytesWanted, DWORD *pdwNumberOfBytesRead);

bool DisableWriteProtect(DWORD dwFlashInfoOffset, DwordVector& orig, const IntDwordMap& patch);
bool EnableWriteProtect(DWORD dwFlashInfoOffset, const DwordVector& orig);
bool FindFlashInfoOffset(DWORD& dwFlashInfoOffset);

bool ReadFromFile(char *filename, DWORD dwOffset, DWORD dwSize, ByteVector& buffer);
DWORD CalcCheckSum(BYTE *buffer, DWORD dwSize);

DWORD GetFileSize(const std::string& filename)
{
    HANDLE hSrc = CreateFile(filename.c_str(), GENERIC_READ, FILE_SHARE_READ,
                NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
    if (INVALID_HANDLE_VALUE == hSrc)
    {
        error("Unable to open file %hs", filename.c_str());
        return 0;
    }

    DWORD dwSize= GetFileSize(hSrc, NULL);

    CloseHandle(hSrc);

    return dwSize;
}

void usage()
{
    printf("(C) 2003-2008 Willem jan Hengeveld  itsme@xs4all.nl\n");
    printf("Usage: pnewbootloader [options] filename\n");
    printf("  -f : force - ignoring warnings about rom contents\n");
    printf("  -a : use flash access api\n");
    printf("  -v : verbose\n");
    printf("  -o : rom offset\n");
    printf("  -s : file seek offset\n");
    printf("  -l : length\n");
    printf("  -n : don't remove rom protection\n");
    printf("  -p : N:DWORD - specify patch value\n");
}

// this tries to identify a bootloader by the following:
//   * at offset 0x40 : 'ECEC'
//   * in the romhdr cpuid=0x1c0
//   * the romhdr says numfiles=0, nummods=1
bool isBootloader(const ByteVector& buffer)
{
    ByteVector ecec;
    ecec.push_back(0x45);
    ecec.push_back(0x43);
    ecec.push_back(0x45);
    ecec.push_back(0x43);

    ByteVector cpuid;
    cpuid.push_back(0xc0);
    cpuid.push_back(0x01);
    cpuid.push_back(0x00);
    cpuid.push_back(0x00);

    typedef std::vector<ByteVector::const_iterator> BVIteratorList;
    BVIteratorList ececlist;
    BVIteratorList cpuidlist;

    ByteVector::const_iterator i_ecec= buffer.begin();
    while ((i_ecec= std::search(i_ecec, buffer.end(), ecec.begin(), ecec.end()))!=buffer.end()) {
        ececlist.push_back(i_ecec);
        ++i_ecec;
    }
    ByteVector::const_iterator i_cpuid= buffer.begin();
    while ((i_cpuid= std::search(i_cpuid, buffer.end(), cpuid.begin(), cpuid.end()))!=buffer.end()) {
        cpuidlist.push_back(i_cpuid);
        ++i_cpuid;
    }

    for (BVIteratorList::iterator iie= ececlist.begin() ; iie!=ececlist.end() ; iie++)
    {
        ByteVector::const_iterator ie= *iie;
        if ((iteratorptr(ie)-0x40) != iteratorptr(buffer.begin())) {
            // todo:  instead of ignoring roms not starting at offset 0
            //  i could add code to just flash the rom from whereever it was found
            continue;
        }
        DWORD romhdr= *(DWORD*)(iteratorptr(ie)+4);
        for (BVIteratorList::iterator iic= cpuidlist.begin() ; iic!=cpuidlist.end() ; iic++)
        {
            ByteVector::const_iterator ic= *iic;
            DWORD physstart= *(DWORD*)(iteratorptr(ic)-0x3c);
            if (romhdr-physstart == (iteratorptr(ic)-0x44) - (iteratorptr(ie)-0x40)) {
                debug("found romhdr at fileofs=%08lx, physstart=%08lx\n",
                        (iteratorptr(ic)-0x44)-iteratorptr(buffer.begin()), physstart);

                DWORD numfiles= *(DWORD*)(iteratorptr(ic)-0x14);
                DWORD nummods= *(DWORD*)(iteratorptr(ic)-0x34);
                
                if (numfiles==0 && nummods==1) {
                    return true;
                }
            }
        }
    }
    return false;
}

int main( int argc, char *argv[])
{
    DebugStdOut();

    char *filename=NULL;
    bool bForce= false;
    bool bUseFlashAccessApi= false;
    bool bRemoveRomProtection= true;
    bool bBootloaderTool= strstr(argv[0], "pnewbootloader")!=0;

    DWORD dwRomOffset= 0x80000000;
    DWORD dwLength= bBootloaderTool?0x40000:0;
    DWORD dwFileOffset= 0;

    IntDwordMap patch;

    int argsfound=0;
    for (int i=1 ; i<argc ; i++)
    {
        if (argv[i][0]=='-') switch(argv[i][1])
        {
            case 'f': bForce= true; break;
            case 'a': bUseFlashAccessApi= true; break;
            case 'n': bRemoveRomProtection= false; break;
            case 'v': nVerbose++; break;
            case 'o': HANDLEULOPTION(dwRomOffset, DWORD); break;
            case 's': HANDLEULOPTION(dwFileOffset, DWORD); break;
            case 'l': HANDLEULOPTION(dwLength, DWORD); break;
            case 'p': if (strlen(argv[i])<5 || argv[i][3]!=':') {
                          usage();
                      }
                      patch[argv[i][2]-'0']= strtoul(argv[i]+4,0,16); break;
            default:
                usage();
                return 1;
        }
        else switch (argsfound++)
        {
            case 0: filename= argv[i]; break;
        }
    }

    if (patch.size()==0)
    {
        patch[0]= 0x60000000;
        patch[2]= 0x60040000;
        patch[5]= 0x04000000;    // extend rom size to include extended rom.
    }
    if (argsfound!=1)
    {
        usage();
        return 1;
    }

    if (dwLength==0)
        dwLength= GetFileSize(filename)-dwFileOffset;

    ByteVector buffer; 
    if (!ReadFromFile(filename, dwFileOffset, dwLength, buffer))
    {
        debug("ERROR reading bootloader\n");
        return 1;
    }
    // verify that input file really resembles a bootloader
    if (bBootloaderTool && !bForce && !isBootloader(buffer)) 
    {
        debug("this is possibly not a bootloader!!\n");
        debug("%hs\n", hexdump(dwRomOffset, vectorptr(buffer), 80, 4, 4).c_str());
        return 1;
    }

    if (bBootloaderTool && MessageBox(0,
"Do not run this tool unless :\n"
"    - you want to downgrade your 5.17 bootloader to 5.15\n"
"    - you have somehow managed to overwrite your bootloader and still\n"
"      have a working ( ce bootable ) device, and still have the operating\n"
"      system running\n"
"\n"
"If the above does not apply to you, press NO\n",
"WARNING", MB_YESNO)!=IDYES)
        return 1;

    CheckITSDll();

    int res= 0;

    DWORD dwFlashInfoOffset=0;
    DwordVector protectparams;
    if (bRemoveRomProtection) {
        if (!FindFlashInfoOffset(dwFlashInfoOffset))
        {
            debug("Unable to find flash info offset, cannot disable bootloader writeprotect\n");
            return 1;
        }
        debug("protection found at %08lx\n", dwFlashInfoOffset);

        if (!DisableWriteProtect(dwFlashInfoOffset, protectparams, patch))
        {
            debug("unable to disable bootloader writeprotect\n");
            return 1;
        }
    }

    if (!ITFlashROM(bUseFlashAccessApi?FLASH_TYPE_MAGICIAN:FLASH_TYPE_XIPFLASH, dwRomOffset, vectorptr(buffer), buffer.size()))
    {
        debug("ERROR in ITFlashROM\n");
        res= 1;
    }

    if (bRemoveRomProtection) {
        if (!EnableWriteProtect(dwFlashInfoOffset, protectparams))
        {
            debug("unable to re-enable bootloader writeprotect\n");
            return 1;
        }
    }
    CeRapiUninit();

    return res;
}

bool FindFlashInfoOffset(DWORD& dwFlashInfoOffset)
{
    for (int i=0 ; i<N_VERSIONS ; i++)
    {
        WCHAR vstr[16];
        DWORD nRead;
        if (!ITReadProcessMemory(NULL, versions[i].dwVersion, (BYTE*)vstr, sizeof(vstr), &nRead))
            return false;
        if (wcscmp(vstr, versions[i].versionstring)) {

            if (nVerbose>1) {
                printf("%08lx:  no version '%ls' but\n   %hs\n\n", 
                        versions[i].dwVersion, 
                        versions[i].versionstring, 
                        ascdump(ByteVector((BYTE*)vstr, (BYTE*)(vstr+8))).c_str()); 
            }

            continue;
        }

        DWORD addrs[5];
        if (!ITReadProcessMemory(NULL, versions[i].dwAddressList, (BYTE*)addrs, sizeof(addrs), &nRead))
            return false;

        if (nVerbose) {
            printf("%08lx: addrlist: %hs\n", 
                    versions[i].dwAddressList, 
                    hexdump(ByteVector((BYTE*)addrs, (BYTE*)(addrs+5)), 4).c_str());
        }
        char txt[256];
        if (!ITReadProcessMemory(NULL, addrs[0], (BYTE*)txt, sizeof(txt), &nRead))
            continue;
        if (strcmp(txt, msgFlashStartErase)) {

            if (nVerbose) {
                printf("%08lx:  no txt '%hs' but\n   %hs\n\n", 
                        addrs[0],
                        msgFlashStartErase, 
                        ascdump(ByteVector((BYTE*)txt, (BYTE*)(txt+16))).c_str()); 
            }

            continue;
        }
// TODO: !!!!! figure out if this -4 is only for xda2.
        dwFlashInfoOffset= versions[i].dwBootloaderAddress;
        return true;

    }
    return false;
}

// this changes the protected area to 60000000 ( and thus opening up 80000000 )
bool DisableWriteProtect(DWORD dwFlashInfoOffset, DwordVector& orig, const IntDwordMap& patch)
{
    orig.resize(7);

    DWORD nRead;
    if (!ITReadProcessMemory(NULL, dwFlashInfoOffset, (BYTE*)vectorptr(orig), sizeof(DWORD)*orig.size(), &nRead)) {
        debug("ERROR: reading from %08lx\n", dwFlashInfoOffset);
        return false;
    }
    if (nVerbose) {
        printf("%08lx: orig addrlist: %hs\n", dwFlashInfoOffset, hexdump((BYTE*)vectorptr(orig), orig.size(), 4).c_str());
    }

    DwordVector data= orig;

    for (IntDwordMap::const_iterator i= patch.begin() ; i!=patch.end() ; ++i)
        data[(*i).first]= (*i).second;

    if (nVerbose) {
        printf("%08lx:  new addrlist: %hs\n", dwFlashInfoOffset, hexdump((BYTE*)vectorptr(data), data.size(), 4).c_str());
    }
    DWORD nWritten;
    return ITWriteProcessMemory(NULL, dwFlashInfoOffset, (BYTE*)vectorptr(data), sizeof(DWORD)*data.size(), &nWritten);
}
// reset it to normal
bool EnableWriteProtect(DWORD dwFlashInfoOffset, const DwordVector& orig)
{
    DWORD nWritten;
    return ITWriteProcessMemory(NULL, dwFlashInfoOffset, (BYTE*)vectorptr(orig), sizeof(DWORD)*orig.size(), &nWritten);
}
bool ITFlashROM(DWORD dwType, DWORD dwOffset, BYTE *buffer, DWORD dwBufLength)
{
    DWORD insize= sizeof(FlashROMParams)+dwBufLength;
    FlashROMParams *req= (FlashROMParams*)LocalAlloc(LPTR, insize);

    req->dwType= dwType;
    req->dwOffset= dwOffset;
    req->dwSize= dwBufLength;
    memcpy(req->buffer, buffer, dwBufLength);
    req->dwCheckSum= CalcCheckSum(req->buffer, req->dwSize);

    DWORD outsize=0;
    FlashROMResult *outbuf=NULL;

    HRESULT res= ItsutilsInvoke(L"ITFlashROM",
            insize, (BYTE*)req,
            &outsize, (BYTE**)&outbuf);
    if (res || outbuf==NULL)
    {
        error(res, "ITFlashROM(%d, %08lx, %08lx)", dwType, dwOffset, dwBufLength);
        return false;
    }
    LocalFree(outbuf);

    return true;
}


bool ReadFromFile(char *filename, DWORD dwOffset, DWORD dwSize, ByteVector& buffer)
{
    FILE *f= fopen(filename, "rb");
    if (f==NULL)
    {
        error("fopen(%hs)", filename);
        return false;
    }
    if (0!=fseek(f, dwOffset, SEEK_SET))
    {
        error("fseek(%hs, %08lx)", filename, dwOffset);
        fclose(f);
        return false;
    }
    buffer.resize(dwSize);
    if (1!=fread(vectorptr(buffer), dwSize, 1, f))
    {
        error("fread(%hs, %08lx)", filename, dwSize);
        fclose(f);
        return false;
    }

    fclose(f);
    return true;
}

DWORD CalcCheckSum(BYTE *buffer, DWORD dwSize)
{
    DWORD sum= 0;
    while (dwSize--)
        sum+=*buffer++;

    return sum;
}

bool ITReadProcessMemory(HANDLE hProc, DWORD dwOffset, BYTE *buffer, DWORD dwBytesWanted, DWORD *pdwNumberOfBytesRead)
{
    ReadProcessMemoryParams inbuf;
    DWORD outsize=0;
    ReadProcessMemoryResult *outbuf=NULL;

    inbuf.hProcess= hProc;
    inbuf.dwOffset= dwOffset;
    inbuf.nSize= dwBytesWanted;
    inbuf.nDataAccess= 0;
    outbuf= NULL; outsize= 0;
    HRESULT res= ItsutilsInvoke(L"ITReadProcessMemory",
            sizeof(ReadProcessMemoryParams), (BYTE*)&inbuf,
            &outsize, (BYTE**)&outbuf);
    if (res || outbuf==NULL)
    {
        error(res, "ITReadProcessMemory");
        return false;
    }
    memcpy(buffer, &outbuf->buffer, outbuf->dwNumberOfBytesRead);
    *pdwNumberOfBytesRead= outbuf->dwNumberOfBytesRead;
    LocalFree(outbuf);

    return true;
}
bool ITWriteProcessMemory(HANDLE hProc, DWORD dwOffset, BYTE *buffer, DWORD dwBufLength, DWORD *pdwNumberOfBytesWritten)
{
    WriteProcessMemoryParams *inbuf;
    DWORD outsize=0;
    WriteProcessMemoryResult *outbuf=NULL;

    DWORD insize= sizeof(WriteProcessMemoryParams)+dwBufLength;
    inbuf= (WriteProcessMemoryParams *)LocalAlloc(LPTR, insize);

    inbuf->hProcess= hProc;
    inbuf->dwOffset= dwOffset;
    inbuf->nSize= dwBufLength;
    inbuf->nDataAccess= 0;
    memcpy(inbuf->buffer, buffer, dwBufLength);

    outbuf= NULL; outsize= 0;
    HRESULT res= ItsutilsInvoke(L"ITWriteProcessMemory",
            insize, (BYTE*)inbuf,
            &outsize, (BYTE**)&outbuf);
    if (res || outbuf==NULL)
    {
        error(res, "ITWriteProcessMemory(%08lx, %08lx, %08lx)", hProc, dwOffset, dwBufLength);
        return false;
    }
    *pdwNumberOfBytesWritten= outbuf->dwNumberOfBytesWritten;
    LocalFree(outbuf);

    return true;
}