#include <windows.h>
#include <map>
#include "OnDiskReader.h"
#include "debug.h"
#include "cenk.h"
#include "PatchFunctions.h"
#include "dllmain.h"
#include "FileFunctions.h"
#include "cever_deps.h"

#ifdef WITHDEBUG
#define ondiskdebug debug
#else
#define ondiskdebug while(0)
#endif

/*
 *  the ondisk.dll devicedriver forwards disk requests to the kernel,
 *  where the real driver is located.
 *
 *  ondisk.dll is found on these devices:
 *  01013fc0 htc:excalibur oxygen pharos juno
 *  01013fc4 samsung_sgh_i900
 *
 *
 *
==== unlocking upper part:
BML_Read     800503EC : 8a000004  -> ea000004

.text:800503E4 E5963010                 LDR     R3, [R6,#0x10]
.text:800503E8 E1530007                 CMP     R3, R7
.text:800503EC 8A000004                 BHI     loc_80050404
.text:800503F0 E59F0318                 LDR     R0, =aBifErrOutOfBou ; L"[BIF:ERR]  Out of bound(nVol:%d,Vsn:%d)"


BML_SglRead  80050890 : 9a0000e2  -> e1a00000

.text:80050888 E5953010                 LDR     R3, [R5,#0x10]
.text:8005088C E1530007                 CMP     R3, R7
.text:80050890 9A0000E2                 BLS     err_outofbound

.text:80050C20             err_outofbound                          ; CODE XREF: BML_sglread+168j
.text:80050C20 E59F0050                 LDR     R0, =aBifErrOutOfBou ; L"[BIF:ERR]  Out of bound(nVol:%d,Vsn:%d)"

BML_Write    80050DAC : 8a000004  -> ea000004

.text:80050DA4 E5963010                 LDR     R3, [R6,#0x10]
.text:80050DA8 E1530004                 CMP     R3, R4
.text:80050DAC 8A000004                 BHI     loc_80050DC4
.text:80050DB0 E59F0404                 LDR     R0, =aBifErrOutOfBou ; L"[BIF:ERR]  Out of bound(nVol:%d,Vsn:%d)"


BML_sglwrite 80051394 : 9a00013b  -> e1a00000

.text:8005138C E5953010                 LDR     R3, [R5,#0x10]
.text:80051390 E1530006                 CMP     R3, R6
.text:80051394 9A00013B                 BLS     loc_80051888

.text:80051888             loc_80051888                            ; CODE XREF: BML_sglwrite+1C0j
.text:80051888 E59F005C                 LDR     R0, =aBifOutOfBoundN ; L"[BIF:   ]  Out of bound(nVol:%d,Vsn:%d)"



"[BIF:ERR]  Out of bound(nVol:%d,Vsn:%d)"
"[BIF:   ]  Out of bound(nVol:%d,Vsn:%d)"

psetmem -4 0x800503ec 0xea000004
psetmem -4 0x80050890 0xe1a00000
psetmem -4 0x80050DAC 0xea000004
psetmem -4 0x80051394 0xe1a00000


-- remove write protect:
psetmem -4 0x8C15A0B4 0xe3a00000

--8004EE0C;8C0CDFB4;8C155200;8C0D3CF4
check_writeprotect
             E52DE004  STR     LR, [SP,#4]!
             E5913088  LDR     R3, [R1,#0x88]  ; note: this varies: #0x80, #0x84, #0x88 depending on os version
             E5912014  LDR     R2, [R1,#0x14]
             E593E00C  LDR     LR, [R3,#0xC]
             E2831010  ADD     R1, R3, #0x10
             E3520001  CMP     R2, #1
             1A00000F  BNE     loc_8004EE68

return_FALSE
             E3A00000  MOV     R0, #0
             E49DE004  LDR     LR, [SP],#4
             E12FFF1E  BX      LR

loc_8004EE34
             E5913004  LDR     R3, [R1,#4]
             E24EE001  SUB     LR, LR, #1
             E3130002  TST     R3, #2
             0A000007  BEQ     loc_8004EE64
             E5912008  LDR     R2, [R1,#8]
             E1520000  CMP     R2, R0
             8A000004  BHI     loc_8004EE64
             E591300C  LDR     R3, [R1,#0xC]
             E0833002  ADD     R3, R3, R2
             E2433001  SUB     R3, R3, #1
             E1530000  CMP     R3, R0
             2A000003  BCS     loc_8004EE74

loc_8004EE64
             E2811010  ADD     R1, R1, #0x10

loc_8004EE68
             E35E0000  CMP     LR, #0
             1AFFFFF0  BNE     loc_8004EE34
             EAFFFFEC  B       return_FALSE

loc_8004EE74
             EB001EBE  BL      writeprotectenabled
             E49DE004  LDR     LR, [SP],#4
             E12FFF1E  BX      LR

 */
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_READ        0x05 // 5 params
#define ONDISK_WRITE       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  ?write
#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 // ?
// 0x12
#define ONDISK_COPYBACK    0x13 // 1 param, returns 0x10
#define ONDISK_FLUSHOP     0x14 // 1 param
#define ONDISK_REPLACEBLK  0x15
// 0x16
#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
// 0x1C
// 0x1D
// 0x1E
// 0x1F
#define ONDISK_GETDEVINFO  0x20 // no params, returns 0x30 bytes

#define CMDDWSIZE 10

// todo: there is also the 'samsungflash' interface, using ioctl 0x1013f98
class ondiskdevice {
public:
    ondiskdevice() : _devid(-1), _ok(false), _cmdsize(0x24), _ondiskioctl(0x1013fc0) { }
    ~ondiskdevice() { close(); }

    void open(int devid)
    {
        if (GetFileInfo(std::string("\\windows\\ondisk.dll"))!=AT_ISFILE)
            return;

        ondiskdebug("ondisk open(%d)\n", devid);
        if (!init_flash()) {
            // try again with different ioctl.
            _cmdsize= 0x28;
            _ondiskioctl= 0x1013fc4;
            if (!init_flash()) {
                ondiskdebug("ERROR initializing ondisk flash\n");
                return;
            }
        }
        if (!open_volume(devid)) {
            ondiskdebug("error opening ondisk volume\n");
            return;
        }
        if (!getdiskinfo(devid, &_di)) {
            ondiskdebug("error getting diskinfo\n");
            return;
        }
        ondiskdebug("ondisk opened: %08lx\n", _ondiskioctl);

        _devid= devid;
        _ok= true;
    }
    void close()
    {
        _devid=-1;
        _ok= false;
    }
    bool GetDiskSize(ULONGLONG& llDiskSize)
    {
        if (!_ok) return false;

        llDiskSize= ULONGLONG(_di.nTotalBlks)*_di.nPgsPerBlk*_di.nSctsPerPg;
        return true;
    }
    bool GetParams(DWORD &dwBlockSize, DWORD &dwSectorSize)
    {
        if (!_ok) return false;
        dwSectorSize = 512;
        dwBlockSize= dwSectorSize * _di.nPgsPerBlk*_di.nSctsPerPg;
        return true;
    }
    bool Read(ULONGLONG llOffset, DWORD dwLength, BYTE *buffer, DWORD &dwRead)
    {
        if (!_ok) return false;
        if (!read_sector(_devid, DWORD(llOffset/512), dwLength/512, buffer))
            return false;
        dwRead= dwLength;
        return true;
    }
    bool Write(ULONGLONG llOffset, DWORD dwLength, const BYTE *buffer, DWORD &dwWritten)
    {
        if (!_ok) return false;
        DWORD dwBlockSize, dwSectorSize;
        GetParams(dwBlockSize, dwSectorSize);

        DWORD dwStartBlock= DWORD(llOffset/dwBlockSize);
        // only erase start when '0' is included
        if (llOffset%dwBlockSize)
            dwStartBlock++;

        DWORD dwEndBlock= DWORD((llOffset+dwLength)/dwBlockSize)-1;
        if ((llOffset+dwLength)%dwBlockSize)
            dwEndBlock++;

        //debug("erasing [%08x-%08x] bs=%08lx blocks %d .. %d\n", DWORD(llOffset), DWORD(llOffset+dwLength), dwBlockSize, dwStartBlock, dwEndBlock);

        for (DWORD blocknr= dwStartBlock ; blocknr<=dwEndBlock ; blocknr++) {
            if (!erase_block(_devid, blocknr))
                return false;
        }
        if (!write_sector(_devid, DWORD(llOffset/dwSectorSize), dwLength/dwSectorSize, buffer))
            return false;
        dwWritten= dwLength;
        return true;
    }

private:
    int _devid;
    bool _ok;
    VolumeSpec _di;
    DWORD _ondiskioctl;
    DWORD _cmdsize;

    bool getdiskinfo(int devid, struct VolumeSpec *di)
    {
        DWORD cmd[CMDDWSIZE]; 
        DWORD nReturned=0;
        cmd[0]=devid; 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(_ondiskioctl, cmd, _cmdsize, (BYTE*)di, sizeof(struct VolumeSpec), &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(_ondiskioctl, cmd, _cmdsize, NULL, 0, &nReturned)) {
            //error("kioctl(%08lx, FLASH:0x%x, init1)", _ondiskioctl, _cmdsize);
            return false;
        }
        return true;
    }
    bool open_volume(int devid)   // ioctl 2
    {
        DWORD cmd[CMDDWSIZE];
        DWORD nReturned=0;
        cmd[0]=devid; 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(_ondiskioctl, cmd, _cmdsize, NULL, 0, &nReturned)) {
            error("kioctl(FLASH, open)");
            return false;
        }
        return true;
    }
    bool read_sector(DWORD devid, DWORD secnr, DWORD nsec, BYTE *buf)
    {
        ondiskdebug("readsector(%d, %08lx, %d)\n", devid, secnr, nsec);
        DWORD cmd[CMDDWSIZE];
        DWORD nReturned=0;
        cmd[0]=devid;
        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>");
        //ondiskdebug("cmd: %s\n", hexdump((BYTE*)cmd, 9, 4).c_str());
        if (!KernelIoControl(_ondiskioctl, cmd, _cmdsize, NULL, 0, &nReturned)) {
            error("kioctl(FLASH, read)");
            return false;
        }
        //ondiskdebug("nret=%08lx   cmd=%s\n", nReturned, hexdump((BYTE*)cmd, 9, 4).c_str());
        //OutputDebugString(L"<2>");
        return true;
    }
    bool erase_block(DWORD devid, DWORD blknr)
    {
        ondiskdebug("eraseblock(%d, %08lx)\n", devid, blknr);
        DWORD cmd[CMDDWSIZE];
        DWORD nReturned=0;
        cmd[0]=devid;
        cmd[1]=ONDISK_ERASEBLK;
        cmd[2]=blknr;
        cmd[3]=0;
        cmd[4]=0;
        cmd[5]=0;
        cmd[6]=0;
        cmd[7]=0;
        cmd[8]=0;
        cmd[9]=0;
        //OutputDebugString(L"<1>");
        //ondiskdebug("cmd: %s\n", hexdump((BYTE*)cmd, 9, 4).c_str());
        if (!KernelIoControl(_ondiskioctl, cmd, _cmdsize, NULL, 0, &nReturned)) {
            error("kioctl(FLASH, erase)");
            return false;
        }
        //ondiskdebug("nret=%08lx   cmd=%s\n", nReturned, hexdump((BYTE*)cmd, 9, 4).c_str());
        //OutputDebugString(L"<2>");
        return true;
    }

    bool write_sector(DWORD devid, DWORD secnr, DWORD nsec, const BYTE *buf)
    {
        ondiskdebug("writesector(%d, %08lx, %d)\n", devid, secnr, nsec);
        DWORD cmd[CMDDWSIZE];
        DWORD nReturned=0;
        cmd[0]=devid;
        cmd[1]=ONDISK_WRITERES;
        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>");
        //ondiskdebug("cmd: %s\n", hexdump((BYTE*)cmd, 9, 4).c_str());
        if (!KernelIoControl(_ondiskioctl, cmd, _cmdsize, NULL, 0, &nReturned)) {
            error("kioctl(FLASH, write)");
            return false;
        }
        //ondiskdebug("nret=%08lx   cmd=%s\n", nReturned, hexdump((BYTE*)cmd, 9, 4).c_str());
        //OutputDebugString(L"<2>");
        return true;
    }

};
typedef std::map<int, ondiskdevice> ondiskmap_t;
typedef std::pair<ondiskmap_t::iterator,bool> ondiskmap_insert;

static ondiskmap_t g_ondiskmap;

ondiskdevice& getondisk(int devid)
{
    ondiskmap_t::iterator i= g_ondiskmap.find(devid);
    if (i==g_ondiskmap.end()) {
        ondiskmap_insert ins= g_ondiskmap.insert(
                ondiskmap_t::value_type(devid, ondiskdevice())
                );
        i= ins.first;
        if (ins.second)
            (*i).second.open(0);
    }
    return (*i).second;
}

bool OnDisk_GetDiskSize(int devid, ULONGLONG &llDiskSize)
{
    ondiskdevice& dev= getondisk(devid);

    return dev.GetDiskSize(llDiskSize);
}
bool OnDisk_GetDiskParams(int devid, DWORD &dwBlockSize, DWORD &dwSectorSize)
{
    ondiskdevice& dev= getondisk(devid);

    return dev.GetParams(dwBlockSize, dwSectorSize);
}
bool OnDisk_Read(int devid, ULONGLONG llOffset, DWORD dwLength, BYTE *buffer, DWORD &dwRead)
{
    ondiskdevice& dev= getondisk(devid);

    return dev.Read(llOffset, dwLength, buffer, dwRead);
}
bool OnDisk_Write(int devid, ULONGLONG llOffset, DWORD dwLength, const BYTE *buffer, DWORD &dwWritten)
{
    ondiskdevice& dev= getondisk(devid);

    return dev.Write(llOffset, dwLength, buffer, dwWritten);
}
void patch_ondiskdriver(DwordPtrVector &offsets, DwordVector &oldvalues)
{
    // 0 : patch conditional branch
    // 1 : patch 'MOV R0,#0'
    // 2 : patch 'BX LR'
    enum patch_types {
        PATCH_UNCONDITIONAL,
        PATCH_MOV_R0_0,
        PATCH_BX_LR,
        PATCH_NOP,
    };
    std::vector<patch_types> types;
// search kernel for E59[56]3010 E153000[467] [89]A000...
    PROCESS *p= FindProcessForName(L"nk.exe");
    if (p==NULL) {
        ondiskdebug("nk.exe not found\n");
        return ;
    }
    DWORD vbase=0;
    DWORD vsize=0;
    GetProcessVirtualRange(p, &vbase, &vsize);


    for (DWORD *p=(DWORD*)vbase ; p<(DWORD*)(vbase+vsize-0x20) ; p++)
    {
        // find 'unlock-upper section' patch
        if ((p[0]&0xfff0ffff)==0xe5903010)    // LDR     R3, [R?,#0x10]
        if ((p[1]&0xfffffff0)==0xe1530000)    // CMP     R3, R?
        if ((p[2]&0xeffff000)==0x8a000000     // BHI/BLS  PC+?
                || (p[2]&0xeffff000)==0x8afff000) {  // BLS  PC-?
            offsets.push_back(p+2);
            oldvalues.push_back(p[2]);

            // 8a...... : BHI
            // 9a...... : BLS
            if ((p[2]&0xFF000000)==0x8A000000)
                types.push_back(PATCH_UNCONDITIONAL);
            else
                types.push_back(PATCH_NOP);
            ondiskdebug("found: patch0: %08lx: %08lx %08lx %08lx\n", p, p[0], p[1], p[2]);
        }

        // find 'writeprotect'-patch
        if (p[0]==0xE52DE004)              // STR     LR, [SP,#4]!
        if ((p[1]&0xffffff00)==0xE5913000) // LDR     R3, [R1,#0x88]  ; note: this varies: #0x80, #0x84, #0x88 depending on os version
        if (p[2]==0xE5912014)              // LDR     R2, [R1,#0x14]
        if (p[3]==0xE593E00C)              // LDR     LR, [R3,#0xC]
        if (p[4]==0xE2831010) {            // ADD     R1, R3, #0x10
             // E3520001  CMP     R2, #1
             // 1A00000F  BNE     loc_8004EE68   <-- change this to NOP
            offsets.push_back(p+6);
            oldvalues.push_back(p[6]);
            types.push_back(PATCH_NOP);

            ondiskdebug("found: patch1: %08lx: %08lx %08lx %08lx\n", p, p[0], p[1], p[2]);
        }
    }
    if (offsets.size()<4) {
        if (offsets.size()>0)
            debug("NOTE: only %d ondisk patch offsets were found, requiring minimum 4, to work\n", offsets.size());
        offsets.clear();
        oldvalues.clear();
        return;
    }
    for (unsigned i=0 ; i<offsets.size() ; i++)
    {
        DWORD *p= offsets[i];

        switch(types[i])
        {
            case PATCH_UNCONDITIONAL: p[0]|=0xE0000000; break;  // change to unconditional (jump)
            case PATCH_NOP:           p[0]= 0xE1A00000; break;  // change to nop
            case PATCH_MOV_R0_0:      p[0]= 0xE3A00000; break;  // change to mov r0,#0
            case PATCH_BX_LR:         p[0]= 0xE12FFF1E; break;  // change to bx lr
        }

        ondiskdebug("patching %08lx : %d : %08lx -> %08lx\n", p, types[i], oldvalues[i], *p);
    }
}
bool OnDisk_Protection(bool bEnabled)
{
    static DwordPtrVector offsets;
    static DwordVector oldvalues;

    ondiskdebug("+OnDisk_Protection(%d)[patch=%d:%d]\n", bEnabled, offsets.size(), oldvalues.size());

    if (bEnabled) {
        unpatch(offsets, oldvalues);
        ReleaseItsutils();
    }
    else {
        patch_ondiskdriver(offsets, oldvalues);
        KeepItsutils();
    }
    ondiskdebug("-OnDisk_Protection[patch=%d:%d]\n", offsets.size(), oldvalues.size());
    return true;
}