#include  <windows.h>
#include "pkfuncs.h"
#include "debug.h"
#include "stringutils.h"
#include "args.h"
#include <map>

// SLOCK 1
// SUNLOCK 1
// MFGDATA 1
// HTCLOGO 1
// GSMDATA 1
// GSMCODE 0x18
// SPLASH 2
// HTCSRV 1
//
// volumespec:
// nPgsPerBlk;   0040    
// nSctsPerPg;   04      
// nBadPos;      02      
// nEccPos;      08      
// nTrTime;      0000c350
// nTwTime;      00055730
// nTeTime;      001e8480
// nTfTime;      00000032
// nTotalBlks;   000003e9
// nMEFlag;      00000001
// val_field30;  00000002
// aUID[16];     ec 00 ec ee e9 6a 32 32 33 33 32 31 07 07 09 06 
//
//            partnr    attr    start     size   ( in blocks )
// pi0.2:  00000002 00000002 00000005 00000020  gsm etc
// pi0.3:  00000003 00000002 00000025 00000160  os
// pi0.6:  00000006 00000022 00000001 00000004  spl
// pi0.8:  00000008 00000001 00000185 00000264  userfs
//
// 1 block = 256 sectors
// 1 sector = 512 bytes
//
// unistore docu:
//   part2 = stage 3 bootloader
//   part3 = os image
//   part8 = fat filesystem
//
//  partition attr:
//    1 BML_PI_ATTR_RW
//    2 BML_PI_ATTR_RO
//   20 BML_PI_ATTR_FROZEN
//   
// c:\local\wince500\PUBLIC\COMMON\OAK\INC\diskio.h

/*  -- from the samsung sgh i900  ondisk.dll
BML_Init         0x1013FC4, {0,    0x01, 0,0, 0,0,0,0, 0,0}, 0x28, 0, 0                         
BML_Open         0x1013FC4, {arg0, 0x02, 0,0, 0,0,0,0, 0,0}, 0x28, 0, 0                      
BML_Close        0x1013FC4, {arg0, 0x03, 0,0, 0,0,0,0, 0,0}, 0x28, 0,0                       
// 4
BML_Read         0x1013FC4, {arg0, 0x05, 0, arg1, arg2, arg3, 0, arg4, arg5, 0}, 0x28, 0,0   
BML_Write        0x1013FC4, {arg0, 0x06, 0, arg1, arg2, arg3, 0, arg4, arg5, 0}, 0x28, 0,0   
BML_EraseBlk     0x1013FC4, {arg0, 0x07, arg1, 0, 0,0,0,0, arg2, 0}, 0x28, 0, 0              
BML_GetVolId     0x1013FC4, {arg0, 0x08, 0,0, 0,0,0,0, 0,0}, 0x28, arg1, 0x10
// 9
BML_ReadRes      0x1013FC4, {arg0, 0x0A, 0, arg1, arg2, arg3, 0, arg4, arg5, 0}, 0x28, 0,0 
BML_WriteRes     0x1013FC4, {arg0, 0x0B, 0, arg1, arg2, arg3, 0, arg4, arg5, 0}, 0x28, 0,0 
BML_MEraseBlk    0x1013FC4, {arg0, 0x0C, 0, 0, 0,0,0,0, arg2,arg1}, 0x28,0,0               
// D
// E
BML_LoadPIEntry  0x1013FC4, {arg0, 0x0F, ?,0, 0,0,0,0, 0,0},0x28, arg2={arg1}, 0x10          
BML_GetVolInfo   0x1013FC4, {arg0, 0x10, 0,0, 0,0,0,0, 0,0}, 0x28, arg1, 0x34              
BML_Format       0x1013FC4, {arg0, 0x11,  ??
// 12
BML_CopyBack     0x1013FC4, {arg0, 0x13, 0, 0, 0, 0, 0, 0, arg2, 0}, 0x28, arg1, 0x10     
BML_FlushOp      0x1013FC4, {arg0, 0x14, 0, 0, 0,0,0,0, arg1, 0}, 0x28, 0, 0              
// 15
// 16 ?
BML_LoadPIExt    0x1013FC4, {arg0, 0x17, arg1,0, 0,0,0,0, 0,0},0x28, arg2,0xC              
// 18 ?
BML_Copy         0x1013FC4, {arg0, 0x19, 0,0, 0,0,0,0, arg2,0},0x28, arg1,0x38               
BML_SGLRead      0x1013FC4, {arg0, 0x1A, 0,arg1, arg2,0,arg3,arg4, arg5,0}, 0x28, 0,0       
BML_SGLWrite     0x1013FC4, {arg0, 0x1B, 0,arg1, arg2,0,arg3,arg4, arg5,0}, 0x28, 0,0      
// 1C ?
// 1D ?
// 1E ?
// 1F ?
BML_GetDevInfo   0x1013FC4, {arg0, 0x20, 0, ...}, 0x28, arg1, 0x30                        



di0:  
0040        nPgsPerBlk; 
04          nSctsPerPg; 
02          nBadPos;    
00000008    nEccPos;    
0000c350    nTrTime;    
00055730    nTwTime;    
001e8480    nTeTime;    
00000032    nTfTime;    
000007d2    nTotalBlks; 
00000001    nMEFlag;    
00000002    val_field30;
ec 00 a2 0d ad 78 32 31 32 32 30 30 06 06 08 08   aUID

pi<volnr>.<partnr>: id  unknown start size
pi0.3:  00000003 00000022 00000012 000003ea   DSK1:
pi0.5:  00000005 00000022 00000005 00000003
pi0.6:  00000006 00000022 00000001 00000004
pi0.7:  00000007 00000001 000007d1 00000001
pi0.8:  00000008 00000001 000003fc 0000030d   DSK2:
pi0.9:  00000009 00000022 00000008 0000000a
pi0.A:  0000000a 00000001 00000709 000000c8


*/

extern "C"
BOOL KernelIoControl(DWORD dwIoControlCode, LPVOID lpInBuf, DWORD nInBufSize, LPVOID lpOutBuf, DWORD nOutBufSize, LPDWORD lpBytesReturned);

struct VolumeSpec {     // size=0x34
    WORD   nPgsPerBlk;     // 00
    BYTE   nSctsPerPg;     // 02
    BYTE   nBadPos;        // 03
    BYTE   nEccPos;        // 04
    BYTE   field_5;        // 05
    BYTE   field_6;        // 06
    BYTE   field_7;        // 07
    DWORD  nTrTime;        // 08
    DWORD  nTwTime;        // 0C
    DWORD  nTeTime;        // 10
    DWORD  nTfTime;        // 14
    DWORD  nTotalBlks;     // 18
    DWORD  nMEFlag;        // 1C
    DWORD  val_field30;    // 20
    BYTE   aUID[16];       // 24
};
struct PartSpec {       // size= 0x10
    DWORD id;
    DWORD unknown;
    DWORD start;
    DWORD size;
};
struct DevInfo {
    DWORD d[12];
};
#define ONDISK_INIT        0x01 // no params
#define ONDISK_OPENVOL     0x02 // no params
#define ONDISK_CLOSEVOL    0x03 // no params
// 4
#define ONDISK_WRITE       0x05 // 5 params
#define ONDISK_READ        0x06 // 5 params
#define ONDISK_ERASEBLK    0x07 // 2 params
#define ONDISK_GETVOLID    0x08 // returns 0x10 bytes
// 9
#define ONDISK_READRES     0x0a // 5 params
#define ONDISK_WRITERES    0x0b // 5 params
#define ONDISK_MERASEBLK   0x0C // 2 params
// d
// e
#define ONDISK_LOADPIENTRY 0x0f // no params, returns 0x10 bytes
#define ONDISK_GETVOLINFO  0x10 // no params, returns 0x34 bytes
#define ONDISK_FORMAT      0x11 // ?
// 12
#define ONDISK_COPYBACK    0x13 // 1 param, returns 0x10
#define ONDISK_FLUSHOP     0x14 // 1 param
// 15

#define ONDISK_LOADPIEXT   0x17 // 1 param, returns 0x0C
// 0x18
#define ONDISK_COPY        0x19 // 1 param, returns 0x38
#define ONDISK_SGLREAD     0x1A // 5 params
#define ONDISK_SGLWRITE    0x1B // 5 params
#define ONDISK_GETDEVINFO  0x20 // no params, returns 0x30 bytes

#define CMDDWSIZE 10

int g_cmdsize= 0x24;
int g_ondiskioctl= 0x1013fc0;

bool getdiskinfo(DWORD vol, struct VolumeSpec *di)
{
    DWORD cmd[CMDDWSIZE]; 
    DWORD nReturned=0;
    cmd[0]=vol; cmd[1]=ONDISK_GETVOLINFO; cmd[2]=0; cmd[3]=0; cmd[4]=0; cmd[5]=0; cmd[6]=0; cmd[7]=0; cmd[8]=0; cmd[9]=0;
    if (!KernelIoControl(g_ondiskioctl, cmd, g_cmdsize, (BYTE*)di, sizeof(struct VolumeSpec), &nReturned)) {
        error("kioctl(FLASH, diskinfo)");
        return false;
    }
    return true;
}
bool getdevinfo(DWORD vol, struct DevInfo *di)
{
    DWORD cmd[CMDDWSIZE]; 
    DWORD nReturned=0;
    cmd[0]=vol; cmd[1]=ONDISK_GETVOLINFO; cmd[2]=0; cmd[3]=0; cmd[4]=0; cmd[5]=0; cmd[6]=0; cmd[7]=0; cmd[8]=0; cmd[9]=0;

    if (!KernelIoControl(g_ondiskioctl, cmd, g_cmdsize, (BYTE*)di, sizeof(struct DevInfo), &nReturned)) {
        error("kioctl(FLASH, diskinfo)");
        return false;
    }
    return true;
}

bool getpartinfo(DWORD vol, DWORD part, struct PartSpec*pi)
{
    DWORD cmd[CMDDWSIZE];
    DWORD nReturned=0;
    pi->id=part;
    cmd[0]=vol; cmd[1]=ONDISK_LOADPIENTRY; cmd[2]=0; cmd[3]=0; cmd[4]=0; cmd[5]=0; cmd[6]=0; cmd[7]=0; cmd[8]=0; cmd[9]=0;
    if (!KernelIoControl(g_ondiskioctl, cmd, g_cmdsize, (BYTE*)pi, sizeof(struct PartSpec), &nReturned)) {
        error("kioctl(FLASH, diskinfo)");
        return false;
    }
    return true;
}

bool init_flash()   // ioctl 1
{
    DWORD cmd[CMDDWSIZE];
    DWORD nReturned=0;
    cmd[0]=0; cmd[1]=ONDISK_INIT; cmd[2]=0; cmd[3]=0; cmd[4]=0; cmd[5]=0; cmd[6]=0; cmd[7]=0; cmd[8]=0; cmd[9]=0;
    if (!KernelIoControl(g_ondiskioctl, cmd, g_cmdsize, NULL, 0, &nReturned)) {
        error("kioctl(FLASH, init1)");
        return false;
    }
    return true;
}
bool open_flash(DWORD vol)   // ioctl 2
{
    DWORD cmd[CMDDWSIZE];
    DWORD nReturned=0;
    cmd[0]=vol; cmd[1]=ONDISK_OPENVOL; cmd[2]=0; cmd[3]=0; cmd[4]=0; cmd[5]=0; cmd[6]=0; cmd[7]=0; cmd[8]=0; cmd[9]=0;
    if (!KernelIoControl(g_ondiskioctl, cmd, g_cmdsize, NULL, 0, &nReturned)) {
        error("kioctl(FLASH, open)");
        return false;
    }
    return true;
}

bool read_sector(DWORD vol, DWORD secnr, DWORD nsec, BYTE *buf)
{
    DWORD cmd[CMDDWSIZE];
    DWORD nReturned=0;
    cmd[0]=vol;
    cmd[1]=ONDISK_READRES;
    cmd[2]=0;
    cmd[3]=secnr;
    cmd[4]=nsec;
    cmd[5]=(DWORD)buf;
    cmd[6]=0;
    cmd[7]=0;   // ptr to spare array of nand memory
    cmd[8]=2;   // LLD_FLAG_ECC_ON
    cmd[9]=0;
    //OutputDebugString(L"<1>");
    //debug("cmd: %s\n", hexdump((BYTE*)cmd, 9, 4).c_str());
    if (!KernelIoControl(g_ondiskioctl, cmd, g_cmdsize, NULL, 0, &nReturned)) {
        error("kioctl(FLASH, read)");
        return false;
    }
    //debug("nret=%08lx   cmd=%s\n", nReturned, hexdump((BYTE*)cmd, 9, 4).c_str());
    //OutputDebugString(L"<2>");
    return true;
}

bool save(const TCHAR *file, DWORD vol, DWORD start, DWORD nsect)
{
    HANDLE h= CreateFile(file, GENERIC_WRITE, FILE_SHARE_READ,
                NULL, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);
    if (h==INVALID_HANDLE_VALUE || h==NULL) {
        error("creating %ls", file);
        return false;
    }
    //OutputDebugString(L"<0a>");
    ByteVector buf; buf.resize(0x20000);
    debug("saving %d: %d:%d to %ls\n", vol, start, nsect, file);
    //OutputDebugString(L"<0b>");
    for (DWORD secnr= start ; secnr < start+nsect; secnr+=buf.size()/0x200) {
        memset(vectorptr(buf), 0xaa, buf.size());
        //OutputDebugString(L"<0c>");
        if (!read_sector(vol, secnr, buf.size()/0x200, vectorptr(buf))) {
            //memset(buf, 0x55, sizeof(buf));
        }
        DWORD wrote=0;
        if (!WriteFile(h, vectorptr(buf), buf.size(), &wrote, NULL)) {
            error("write %08lx", secnr);
            CloseHandle(h);
            return false;
        }
    }
    CloseHandle(h);
    return true;
}
void usage() {
    debug("Usage: bkondisk [rootdir]\n");
    debug("        will save all partitions on all volumes in files on <rootdir>\n");
    debug("Usage: bkondisk [-vVOLUME] [-pPARTITION] [rootdir]\n");
    debug("        will save a specific partition on <rootdir>\n");
    debug("Usage: bkondisk [-vVOLUME] [-bSTARTBLOCK] [-nNBLOCKS]  [filename]\n");
    debug("        will save the specified blocks to <filename>\n");
    debug("Usage: bkondisk -i\n");
    debug("        will only list disk info in the logfile\n");
}
int WINAPI WinMain( HINSTANCE hInstance,
                   HINSTANCE hPrevInstance,
                   LPTSTR    lpCmdLine,
                   int       nCmdShow)
{
    DebugSetLogfile("bkondisk.log");

    BOOL bMode = SetKMode(TRUE);
    DWORD dwPerm = SetProcPermissions(0xFFFFFFFF);

    StringList args;
    if (!SplitString(ToString(lpCmdLine), args, false))
    {
        error("Error in commandline");
        return false;
    }
    DWORD volume=0;
    DWORD partition=-1;
    std::string root;
    DWORD start_block=-1;
    DWORD block_count=-1;
    bool info_only= false;
    for (StringList::iterator i= args.begin() ; i!=args.end() ; ++i)
    {
        std::string& arg= *i;
        if (arg[0]=='-') switch(arg[1])
        {
            case '0':
                g_cmdsize= 0x24;
                g_ondiskioctl= 0x1013fc0;
                break;
            case '1':
                g_cmdsize= 0x28;
                g_ondiskioctl= 0x1013fc4;
                break;
            case 'v': HANDLESTLULOPTION(volume, DWORD); break;
            case 'p': HANDLESTLULOPTION(partition, DWORD); break;
            case 'b': HANDLESTLULOPTION(start_block, DWORD); break;
            case 'n': HANDLESTLULOPTION(block_count, DWORD); break;
            case 'i': info_only=true; break;
            default: usage();

        }
        else {
            if (root.size())
                root += ' ';
            root += arg;
        }
    }

    if (!init_flash()) {
        debug("error initializing flash\n");
        return false;
    }

    // todo: find real storagecard by using ?? api to enum storage cards
    if (root.empty())
        root= "\\Storage Card";
    debug("listing partions for each volume in format:\n   pi<volnr>.<partnr>: id  unknown start size\n");
    for (DWORD vol=(volume==-1)?0:volume ; (volume==-1)?(vol<2):vol==volume ; vol++)
    {
        if (!open_flash(vol)) {
            debug("error initializing flash vol %d\n", vol);
            continue;
        }

        struct VolumeSpec di; memset(&di, 0, sizeof(di));
        if (!getdiskinfo(vol, &di)) {
            debug("error getting di(%d)\n", vol);
        }
        else {
            debug("di%d: %s\n", vol, hexdump((BYTE*)&di, 13, 4).c_str());
        }

        if (di.nTotalBlks==0)
            continue;

        if (start_block!=-1) {
            if (block_count==-1)
                block_count=di.nTotalBlks-start_block;
            save(ToWString(root).c_str(), vol, start_block*di.nSctsPerPg*di.nPgsPerBlk, block_count*di.nSctsPerPg*di.nPgsPerBlk);
            return 0;
        }

        typedef std::map<DWORD,DWORD> PartitionMap;
        PartitionMap savedblocks;

        DWORD minblock=0xFFFFFFFF;
        DWORD maxblock=0;
        for (DWORD part=(partition==-1)?0:partition ; (partition==-1)?(part<32):part==partition ; part++) {
            struct PartSpec pi; memset(&pi, 0, sizeof(pi));
            if (!getpartinfo(vol, part, &pi)) {
                debug("error getting pi(%d.%d)\n", vol, part);
            }
            else {
                if (pi.unknown || pi.start || pi.size) {
                    debug("pi%d.%d: %s\n", vol, part, hexdump((BYTE*)&pi, 4, 4).c_str());

                    if (!info_only)
                        save(ToWString(stringformat("%s\\bk_%d%d_%04x.img", root.c_str(), vol, part, pi.start)).c_str(), vol, pi.start*di.nSctsPerPg*di.nPgsPerBlk, pi.size*di.nSctsPerPg*di.nPgsPerBlk);

                    if (savedblocks.find(pi.start)!=savedblocks.end()) {
                        debug("!!! partition reuses block\n");
                    }
                    savedblocks[pi.start]= pi.size;

                    if (pi.start<minblock) 
                        minblock=pi.start;
                    if (pi.start+pi.size>maxblock)
                        maxblock=pi.start+pi.size;
                }
            }
        }
        if (partition==-1) {
            if (minblock) {
                if (info_only) {
                    debug("no partition for blocks 0x%x l=0x%x\n", 0, minblock);
                }
                else {
                    save(ToWString(stringformat("%s\\bk_%d__%04x.img", root.c_str(), vol, 0)).c_str(), vol, 0, minblock*di.nSctsPerPg*di.nPgsPerBlk);
                }
                savedblocks[0]= minblock;
            }
            if (maxblock<di.nTotalBlks) {
                if (info_only) {
                    debug("no partition for blocks 0x%x l=0x%x\n", maxblock, di.nTotalBlks-maxblock);
                }
                else {
                    save(ToWString(stringformat("%s\\bk_%d___%04x.img", root.c_str(), vol, maxblock)).c_str(), vol, maxblock*di.nSctsPerPg*di.nPgsPerBlk, (di.nTotalBlks-maxblock)*di.nSctsPerPg*di.nPgsPerBlk);
                }
                savedblocks[maxblock]= di.nTotalBlks-maxblock;
            }
        }

        DWORD blockptr=0;
        for (PartitionMap::iterator mi= savedblocks.begin() ; mi!=savedblocks.end() ; mi++) {
            if ((*mi).first>blockptr) {
                if (info_only) {
                    debug("no partition for blocks 0x%x l=0x%x\n", blockptr, (*mi).first-blockptr);
                }
                else {
                    save(ToWString(stringformat("%s\\bk_%d__%04x.img", root.c_str(), vol, blockptr)).c_str(), vol, blockptr*di.nSctsPerPg*di.nPgsPerBlk, ((*mi).first-blockptr)*di.nSctsPerPg*di.nPgsPerBlk);
                }
            }
            else if ((*mi).first<blockptr) {
                debug("partition overlapping by 0x%x blocks ( 0x%x - 0x%x )\n", blockptr-(*mi).first, (*mi).first, blockptr);
            }
            blockptr= (*mi).first+(*mi).second;
        }
    }

    SetProcPermissions(dwPerm);
    SetKMode(bMode);
    return 0;
}