#include <windows.h>
#include <map>
#include <algorithm>
#include "FlashDrvReader.h"
//#include "devicedriverloader.h"
#include "debug.h"
#include "cenk.h"
#include "kernelmisc.h"
#include "cever_deps.h"
#include "deviceinfo.h"
#include "dllmain.h"
#include "PatchFunctions.h"
#include "FileFunctions.h"

#ifdef WITHDEBUG
#define flashdrvdebug  debug
#define flashdrvdebugt debugt
#else
#define flashdrvdebug  while(0)
#define flashdrvdebugt while(0) 
#endif

/*
flashdrv.dll is found on these devices:
   htc:artemis diamond hermes kaiser niki galaxy breeze diamondpro s740 trinity
   sony_xperia treo700 eten_m700
*/
#define IOCTL_FMD_NAND_READ_WLAN_DATA  0x71fe8
#define IOCTL_FMD_NAND_WRITE_WLAN_DATA  0x71fc0
#define IOCTL_FMD_WRITE_VERSION_INFO  0x71fd8

DWORD arm_mov(int dst, int src)
{
    return 0xe1a00000+(dst<<12)+src;
}

const DWORD ARM_NOP= 0xE1A00000;

bool makenop(DWORD vofs, DwordPtrVector &offsets, DwordVector &oldvalues)
{
    DWORD pp= VirtToPhys(vofs);
    if (pp==INVALID_PHYSICAL_ADDRESS)
        return false;
    DWORD *vpv= (DWORD*)PhysToVirt(pp);
    if (vpv==0)
        return false;

    offsets.push_back(vpv);
    oldvalues.push_back(*offsets.back());
    *offsets.back()= ARM_NOP;

    return true;
}
bool makejmp(DWORD vofs, DwordPtrVector &offsets, DwordVector &oldvalues)
{
    DWORD pp= VirtToPhys(vofs);
    if (pp==INVALID_PHYSICAL_ADDRESS)
        return false;
    DWORD *vpv= (DWORD*)PhysToVirt(pp);
    if (vpv==0)
        return false;

    offsets.push_back(vpv);
    oldvalues.push_back(*offsets.back());
    *offsets.back()= (*offsets.back()&0x0fffffff)|0xE0000000;

    return true;
}

DWORD searchcode(DWORD vbase, DWORD vsize, const DWORD*writeflashcode, DWORD codesize)
{
    flashdrvdebug("searchcode(%08x, %08x, %08x, %08x)\n", vbase,vsize, writeflashcode,codesize);
    return (DWORD)std::search((const DWORD*)vbase, (const DWORD*)(vbase+vsize), writeflashcode, writeflashcode+codesize);
}

void patch_flashdrv(DwordPtrVector &offsets, DwordVector &oldvalues)
{
//   .text:03DE8978 [E28E1509]   ADD R1, LR, #0x2400000 ; offset to 'APPSBL' 
    DWORD dwAddInstruction= 0xE28E1509;
    DWORD dwRegisterMask  = 0x000ff000;
    DWORD dwOperandMask  =  0x00000fff;
    // apply 'unlimited' patch
    //
    // ADD: 
    // cccc00I0100SnnnnddddOOOOOOOOOOOO
    // S : setflags bit
    // nnnn = source reg
    // dddd = dest reg
    // OOOOOOOOOOOO = shifter operand

    //
    // MOV:
    // cccc00I1101S0000ddddOOOOOOOOOOOO

    // I=1 : immediate : value = RRRRiiiiiiii : val=iiiiiiii>>(2*RRRR)
    // I=0 : O.654 = shifttype   
    //    0  LSL#   iiiii 000 mmmm
    //       MOV    00000 000 mmmm
    //    1  LSL    ssss0 001 mmmm
    //    2  LSR#   iiiii 010 mmmm
    //    3  LSR    ssss0 011 mmmm
    //    4  ASR#   iiiii 100 mmmm
    //    5  Asr    ssss0 101 mmmm
    //    6  ror#   iiiii 110 mmmm
    //       RRX    00000 110 mmmm
    //    7  Ror    ssss0 111 mmmm
    //    ?         ....1 ..1 ....  not dataprocessing insn.


    MODULE *m= FindModuleForName(L"flashdrv.dll");
    if (m==NULL) {
        flashdrvdebug("flashdrv.dll not found\n");
        return ;
    }
    DWORD vbase=0;
    DWORD vsize=0;
    if (!GetModuleVirtualCodeRange(m, &vbase, &vsize)) {
        flashdrvdebug("could not find flash driver code range\n");
        return ;
    }

    searchforinstruction(vbase, vsize, 2, ~dwRegisterMask, dwAddInstruction, offsets);

    if (offsets.size()!=2) {
        if (offsets.size())
            debug("NOTE: %d flashdrv patch offsets were found, requiring exactly 2 to work\n", offsets.size());
        offsets.clear();
        oldvalues.clear();
        return;
    }

    for (unsigned i=0 ; i<offsets.size() ; i++) {
        oldvalues.push_back(*offsets[i]);
        DWORD dwNewValue= (oldvalues.back()&~dwOperandMask);  // change to ADD Rx,Ry,#0
        *offsets[i]= dwNewValue;
        flashdrvdebug("found %08lx at %08lx, changed to %08lx\n", oldvalues.back(), offsets[i], dwNewValue);
    }
}
void patch_nanddll(DwordPtrVector &offsets, DwordVector &oldvalues)
{
/*
-- in write_flash: only the config area is allowed

03DE6DC4  E59D2014                 LDR     R2, [SP,#0x48+vararg0_address]
03DE6DC8  E3A0346D E383372A        MOV     R3, 0x6DA80000
03DE6DD0  E1520003                 CMP     R2, R3
03DE6DD4  3A0000DC                 BCC     return_0     ; < 6da80000
03DE6DD8  E3A0346D E38338AE        MOV     R3, 0x6DAE0000
03DE6DE0  E1520003                 CMP     R2, R3
03DE6DE4  8A0000D8                 BHI     return_0     ; >= 6dae0000

-- in FMD_WriteVersionInfo the write size is limited to 0x100 bytes
03DE4BAC  E5961004                 LDR     R1, [R6,#struc_flashreq.size]
03DE4BB0  E3A03000                 MOV     R3, #0
03DE4BB4  E1A07000                 MOV     R7, R0
03DE4BB8  E3510C01                 CMP     R1, #0x100
03DE4BBC  E58D3004                 STR     R3, [SP,#0x1C+var_18]
03DE4BC0  9A000005                 BLS     loc_3DE4BDC

*/


    MODULE *m= FindModuleForName(L"nand.dll");
    if (m==NULL) {
        flashdrvdebug("nand.dll not found\n");
        return ;
    }
    DWORD vbase=0;
    DWORD vsize=0;
    if (!GetModuleVirtualCodeRange(m, &vbase, &vsize)) {
        flashdrvdebug("could not find nand driver code range\n");
        return ;
    }


    static DWORD writeflashcode[9]= {0xE59D2014, 0xE3A0346D, 0xE383372A, 0xE1520003, 0x3A0000DC, 0xE3A0346D, 0xE38338AE, 0xE1520003, 0x8A0000D8};
    static DWORD writeversioncode[6]={0xE5961004, 0xE3A03000, 0xE1A07000, 0xE3510C01, 0xE58D3004, 0x9A000005};

    DWORD vofs1= searchcode(vbase, vsize, writeflashcode, sizeof(writeflashcode)/sizeof(DWORD));
    DWORD vofs2= searchcode(vbase, vsize, writeversioncode, sizeof(writeversioncode)/sizeof(DWORD));

    if (vofs1==0 || vofs2==0) {
        debug("NOTE: nanddll code not found: %08x %08x\n", vofs1, vofs2);
        return;
    }
    makenop(vofs1+4*4, offsets, oldvalues);
    makenop(vofs1+8*4, offsets, oldvalues);
    makejmp(vofs2+5*4, offsets, oldvalues);

    if (offsets.size()!=3) {
        unpatch(offsets, oldvalues);
        debug("patch unsuccesfull\n");
    }
    else {
        debug("nanddll patch: %08x and %08x\n", vofs1, vofs2);
    }
}

void do_flashdrv_patch(bool bProtectionEnabled)
{
    static DwordPtrVector offsets;
    static DwordVector oldvalues;

    if (bProtectionEnabled) {
        unpatch(offsets, oldvalues);
        ReleaseItsutils();

        flashdrvdebug("flashdrv: re-protected baseband flash\n");
    }
    else {
        // disable nand-mpu ( phys 0xa0b00000 )
        *(DWORD*)0xb1200000=0;

        patch_flashdrv(offsets, oldvalues);
        KeepItsutils();

        flashdrvdebug("flashdrv: un-protected baseband flash\n");
    }
    flashdrvdebug("do_flashdrv_patch[patch=%d:%d]\n", offsets.size(), oldvalues.size());
}
void do_nanddll_patch(bool bProtectionEnabled)
{
    static DwordPtrVector offsets;
    static DwordVector oldvalues;

    if (bProtectionEnabled) {
        unpatch(offsets, oldvalues);
        ReleaseItsutils();

        flashdrvdebug("flashdrv: re-protected nand write\n");
    }
    else {
        patch_nanddll(offsets, oldvalues);
        KeepItsutils();

        flashdrvdebug("flashdrv: un-protected nand write\n");
    }
    flashdrvdebug("do_nanddll_patch[patch=%d:%d]\n", offsets.size(), oldvalues.size());
}



HANDLE findBDEV()
{
    HDATA *hi= cvHandle2HDataPtr((HANDLE)GetCurrentThreadId());
    HDATA *ha;
    ha=hi;

    int total=0;
    for ( ; total==0 || (ha!=hi && ha!=NULL) ; ha= (HDATA*)ha->linkage.fwd) {
        if ((DWORD(ha->hValue)&3)!=2)
            continue;
        if (memcmp("BDEV", ha->pci->acName, 4)==0)
        {
            return ha->hValue;
        }
        total++;
    }
    return NULL;
}
class flashdrvdevice {
    enum { PATCH_NONE, PATCH_FLASHDRV, PATCH_NANDDLL };
public:
    flashdrvdevice() 
        : _devid(-1), _ok(false), _dwHandle(0), _writeioctl(0), _flashdrvpatch(PATCH_NONE), _baseofs(0)
    {
        flashdrvdebugt("flashdrv constructor(%08lx)\n", this);
    }
    ~flashdrvdevice()
    {
        flashdrvdebugt("flashdrv destructor(%08lx)\n", this);
        close();
    }
    void open(int devid)
    {
        bool flashdrv= GetFileInfo(std::string("\\windows\\FlashDrv.dll"))==AT_ISFILE;
        bool nanddll= GetFileInfo(std::string("\\windows\\nand.dll"))==AT_ISFILE;

        if (!flashdrv && !nanddll) {
            flashdrvdebugt("flashdrv did not find either flashdrv.dll or nand.dll\n");
            return;
        }
        if (flashdrv) {
            _writeioctl= IOCTL_FMD_NAND_WRITE_WLAN_DATA;
            _flashdrvpatch= PATCH_FLASHDRV;
            _baseofs= 0;
        }
        else if (nanddll) {
            _writeioctl= IOCTL_FMD_WRITE_VERSION_INFO;
            _flashdrvpatch= PATCH_NANDDLL;
            _baseofs= 0x50000000;
        }

        // first try to find already open 'BDEV' device
        _dwHandle= findBDEV();
        if (_dwHandle==0 || _dwHandle==INVALID_HANDLE_VALUE) {
            _dwHandle= CreateFile(L"DSK1:",
                GENERIC_READ|GENERIC_WRITE, FILE_SHARE_READ|FILE_SHARE_WRITE,
                NULL, OPEN_EXISTING, 0, 0 );
        }
        if (_dwHandle==0 || _dwHandle==INVALID_HANDLE_VALUE) {
            flashdrvdebugt("error %08x opening DSK1:\n", GetLastError());
            return;
        }

/*
        flashdrvdebugt("flashdrvdevice opening\n");
        _dvr.load(_T("Drivers\\BuiltIn\\FLASHDRV"));

        flashdrvdebug("procs loaded: init=%08lx, open=%08lx, ioctl=%08lx\n", _dvr.DrvInit, _dvr.DrvOpen, _dvr.DrvIOControl);
        _dwHandle= _dvr.open(0, 0);
*/
        if (_flashdrvpatch==PATCH_NANDDLL)
            do_nanddll_patch(false);

        flashdrvdebugt("open: handle=%08lx stat=%08lx\n", _dwHandle, GetLastError());

        _devid= devid;
        _ok= true;
    }
    void close()
    {
        if (!_ok) return;
/*
        flashdrvdebugt("calling close %08lx\n", _dvr.DrvClose);
        BOOL clres= _dvr.DrvClose(_dwHandle);
        flashdrvdebugt("close: res=%08lx stat=%08lx\n", clres, GetLastError());
*/
        if (_flashdrvpatch==PATCH_NANDDLL)
            do_nanddll_patch(true);

        _devid=-1;
        _ok= false;
    }

    bool GetDiskSize(ULONGLONG& llDiskSize)
    {
        if (!_ok) return false;

        llDiskSize= 0x10000000;
        return true;
    }
    bool GetParams(DWORD &dwBlockSize, DWORD &dwSectorSize)
    {
        if (!_ok) return false;
        dwSectorSize = 0x800;
        dwBlockSize= 0x20000;
        return true;
    }
    bool Read(ULONGLONG llOffset, DWORD dwLength, BYTE *buffer, DWORD &dwRead)
    {
        if (!_ok) return false;

        // the buffer passed in the request must be allocated with vmbase. otherwise it will be incorrectly
        // translated to be in filesys.exe's process space.
        DWORD req[3];
        req[0]= DWORD(llOffset)+_baseofs;
        req[1]= dwLength;
        req[2]= (DWORD)globalptr(buffer);
        DWORD nRet;
        flashdrvdebug("flashdrv_device::Read(%I64x, %08lx, %08lx):rq=%08lx [base=%x]\n", llOffset, dwLength, buffer, req, _baseofs);
/*
        BOOL iores= _dvr.DrvIOControl(_dwHandle, IOCTL_FMD_NAND_READ_WLAN_DATA, NULL, 0, (PBYTE)req, 12, &nRet);
*/
        BOOL iores= DeviceIoControl(_dwHandle, IOCTL_FMD_NAND_READ_WLAN_DATA, NULL, 0, (PBYTE)req, 12, &nRet, NULL);
        dwRead= dwLength;
        flashdrvdebug("->%08lx, ret=%08lx\n", iores, nRet);
        return true;
    }
    bool Write(ULONGLONG llOffset, DWORD dwLength, const BYTE *buffer, DWORD &dwWritten)
    {
        if (!_ok) return false;

        DWORD req[3];
        req[0]= DWORD(llOffset)+_baseofs;
        req[1]= dwLength;
        req[2]= (DWORD)globalptr(buffer);
        DWORD nRet;
        flashdrvdebug("flashdrv_device::Write(%I64x, %08lx, %08lx):rq=%08lx [base=%x, ioctl=%x]\n", llOffset, dwLength, buffer, req, _baseofs, _writeioctl);
/*
        BOOL iores= _dvr.DrvIOControl(_dwHandle, IOCTL_FMD_NAND_WRITE_WLAN_DATA, (PBYTE)req, 12, NULL, 0, &nRet);
*/
        BOOL iores= DeviceIoControl(_dwHandle, _writeioctl, (PBYTE)req, 12, NULL, 0, &nRet, NULL);
        flashdrvdebug("->%08lx, ret=%08lx\n", iores, nRet);


        dwWritten= dwLength;
        return true;
    }
    bool Protect(bool bProtectionEnabled)
    {
        flashdrvdebug("+FlashDrv_Protection(%d)\n", bProtectionEnabled);

        if (_flashdrvpatch==PATCH_FLASHDRV) {
            do_flashdrv_patch(bProtectionEnabled);
        }
        else {
            if (bProtectionEnabled) {
                _baseofs= 0x50000000;
                flashdrvdebug("nand: re-protected baseband flash\n");
            }
            else {
                // disable nand-mpu ( phys 0xa0b00000 )
                *(DWORD*)0xb1400000=0;

                _baseofs= 0x4dc00000;
                flashdrvdebug("nand: un-protected baseband flash\n");
            }
        }
        return true;
    }

private:
    int _devid;
    bool _ok;

    HANDLE _dwHandle;
    DWORD _writeioctl;
    int _flashdrvpatch;

    DWORD _baseofs;

//    devicedriver _dvr;
};

typedef std::map<int, flashdrvdevice> flashdrvmap_t;
typedef std::pair<flashdrvmap_t::iterator,bool> flashdrvmap_insert;

static flashdrvmap_t g_flashdrvmap;

flashdrvdevice& getflashdrv(int devid)
{
    flashdrvmap_t::iterator i= g_flashdrvmap.find(devid);
    if (i==g_flashdrvmap.end()) {
        flashdrvmap_insert ins= g_flashdrvmap.insert(
                flashdrvmap_t::value_type(devid, flashdrvdevice())
                );
        i= ins.first;
        flashdrvdebug("new flashdrvdevice: %d %d\n", ins.second, devid);
        if (ins.second) {
            //KeepItsutils();
            (*i).second.open(devid);
        }
    }
    else {
        //debug("existing flashdrv device %d\n", devid);
    }
    return (*i).second;
}

bool FlashDrv_GetDiskSize(int devid, ULONGLONG &llDiskSize)
{
    flashdrvdevice& dev= getflashdrv(devid);

    return dev.GetDiskSize(llDiskSize);
}
bool FlashDrv_GetDiskParams(int devid, DWORD &dwBlockSize, DWORD &dwSectorSize)
{
    flashdrvdevice& dev= getflashdrv(devid);

    return dev.GetParams(dwBlockSize, dwSectorSize);
}
bool FlashDrv_Read(int devid, ULONGLONG llOffset, DWORD dwLength, BYTE *buffer, DWORD &dwRead)
{
    //debug("FlashDrv_Read(%d, %I64x, %08lx)\n", devid, llOffset, dwLength);
    flashdrvdevice& dev= getflashdrv(devid);

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

    return dev.Write(llOffset, dwLength, buffer, dwWritten);
}
bool FlashDrv_Protection(int devid, bool bProtectionEnabled)
{
    flashdrvdevice& dev= getflashdrv(devid);

    return dev.Protect(bProtectionEnabled);
}