/* (C) 2003 XDA Developers  itsme@xs4all.nl
 *
 * $Header$
 */
#define NOMINMAX
#include  <windows.h>
#include  <Pwinuser.h>
#include "ceioctl.h"
#include "debug.h"
#include <string>

#include "stringutils.h"
#include "kernelmisc.h"

#define XDA2PASSWORD "aYaLaMiH"
#define XDA3PASSWORD "eGnAeUlB"
// Prophet password:  a!032191

//inline int min(int a, int b) { return a<b ? a : b; }

std::string FLStatus_string(FLStatus status);

// ------ from blockdev.h
#define FL_NOT_DOC    0 // - Not DiskOnChip
#define FL_DOC        1 // - DiskOnChip 2000
#define FL_MDOC       2 // - MDOC8-based devices
#define FL_MDOCP_16   4 // - MDOCP16-based devices (interleave-1)
#define FL_MDOCP      5 // - MDOCP32-based devices (interleave-2)
#define FL_DOC_OREN   6 // - OREN-based devices
#define FL_MDOC512_G3 7 // - MDOC512(G3)-based devices
#define FL_MDOC256_P3 8 // - MDOC256(P3)

// --------- to string translations
std::string Protection_string(byte protection)
{
    StringList list;
    if (protection&PROTECTABLE          ) list.push_back("PROTECTABLE");
    if (protection&READ_PROTECTED       ) list.push_back("READ_PROTECTED");
    if (protection&WRITE_PROTECTED      ) list.push_back("WRITE_PROTECTED");
    if (protection&LOCK_ENABLED         ) list.push_back("LOCK_ENABLED");
    if (protection&LOCK_ASSERTED        ) list.push_back("LOCK_ASSERTED");
    if (protection&KEY_INSERTED         ) list.push_back("KEY_INSERTED");
    if (protection&CHANGEABLE_PROTECTION) list.push_back("CHANGEABLE_PROTECTION");

    return JoinStringList(list, ",");
}

std::string FLStatus_string(FLStatus status)
{
switch(status) {
case flOK: return "flOK";
case flBadFunction: return "flBadFunction";
case flFileNotFound: return "flFileNotFound";
case flPathNotFound: return "flPathNotFound";
case flTooManyOpenFiles: return "flTooManyOpenFiles";
case flNoWriteAccess: return "flNoWriteAccess";
case flBadFileHandle: return "flBadFileHandle";
case flDriveNotAvailable: return "flDriveNotAvailable";
case flNonFATformat: return "flNonFATformat";
case flFormatNotSupported: return "flFormatNotSupported";
case flNoMoreFiles: return "flNoMoreFiles";
case flWriteProtect: return "flWriteProtect";
case flBadDriveHandle: return "flBadDriveHandle";
case flDriveNotReady: return "flDriveNotReady";
case flUnknownCmd: return "flUnknownCmd";
case flBadFormat: return "flBadFormat";
case flBadLength: return "flBadLength";
case flDataError: return "flDataError";
case flUnknownMedia: return "flUnknownMedia";
case flSectorNotFound: return "flSectorNotFound";
case flOutOfPaper: return "flOutOfPaper";
case flWriteFault: return "flWriteFault";
case flReadFault: return "flReadFault";
case flGeneralFailure: return "flGeneralFailure";
case flDiskChange: return "flDiskChange";
case flVppFailure: return "flVppFailure";
case flBadParameter: return "flBadParameter";
case flNoSpaceInVolume: return "flNoSpaceInVolume";
case flInvalidFATchain: return "flInvalidFATchain";
case flRootDirectoryFull: return "flRootDirectoryFull";
case flNotMounted: return "flNotMounted";
case flPathIsRootDirectory: return "flPathIsRootDirectory";
case flNotADirectory: return "flNotADirectory";
case flDirectoryNotEmpty: return "flDirectoryNotEmpty";
case flFileIsADirectory: return "flFileIsADirectory";
case flAdapterNotFound: return "flAdapterNotFound";
case flFormattingError: return "flFormattingError";
case flNotEnoughMemory: return "flNotEnoughMemory";
case flVolumeTooSmall: return "flVolumeTooSmall";
case flBufferingError: return "flBufferingError";
case flFileAlreadyExists: return "flFileAlreadyExists";
case flIncomplete: return "flIncomplete";
case flTimedOut: return "flTimedOut";
case flTooManyComponents: return "flTooManyComponents";
case flTooManyDrives: return "flTooManyDrives";
case flTooManyBinaryPartitions: return "flTooManyBinaryPartitions";
case flPartitionNotFound: return "flPartitionNotFound";
case flFeatureNotSupported: return "flFeatureNotSupported";
case flWrongVersion: return "flWrongVersion";
case flTooManyBadBlocks: return "flTooManyBadBlocks";
case flNotProtected: return "flNotProtected";
case flUnchangeableProtection: return "flUnchangeableProtection";
case flBadDownload: return "flBadDownload";
case flBadBBT: return "flBadBBT";
case flInterleaveError: return "flInterleaveError";
case flWrongKey: return "flWrongKey";
case flHWProtection: return "flHWProtection";
case flLeftForCompetability: return "flLeftForCompetability";
case flMultiDocContradiction: return "flMultiDocContradiction";
case flCanNotFold: return "flCanNotFold";
case flBadIPLBlock: return "flBadIPLBlock";
case flIOCommandBlocked: return "flIOCommandBlocked";
default: return stringformat("FLStatus_%d", status);
}
}

HANDLE OpenDisk(int n)
{
    TCHAR diskname[16];
    _sntprintf(diskname, 16, L"DSK%u:", n);

    HANDLE hDsk= CreateFile( diskname,
            GENERIC_READ|GENERIC_WRITE, FILE_SHARE_READ|FILE_SHARE_WRITE,
            NULL, OPEN_EXISTING, 0, 0 );
    if (hDsk==NULL || hDsk==INVALID_HANDLE_VALUE) {
        return NULL;
    }
    return hDsk;
}

bool PartitionInsertKey(HANDLE hDsk, bool bInsert, const std::string& passwd)
{
    //flProtectionInput in; memset(&in, 0, sizeof(in));
    BYTE in[10];

    in[0]= 0;
    memcpy(in+1, passwd.c_str(), 8);
    in[9]= bInsert?PROTECTION_INSERT_KEY:PROTECTION_REMOVE_KEY;

    flProtectionOutput out;  memset(&out, 0, sizeof(out));

    DWORD nReturned=0;
    if (!DeviceIoControl(hDsk, FL_IOCTL_BDTL_HW_PROTECTION, &in, sizeof(in), &out, sizeof(out), &nReturned, NULL)) {
        error("FL_IOCTL_BDTL_HW_PROTECTION-%s-key", bInsert?"insert":"remove");
        return false;
    }

    return true;
}
bool BDKPartitionInsertKey(HANDLE hDsk, bool bInsert, int partnr, const std::string& passwd)
{
    //flBDKProtectionInput in; memset(&in, 0, sizeof(in));
    BYTE in[11];

    in[0]= partnr;
    in[1]= 0;
    memcpy(in+2, passwd.c_str(), 8);
    in[10]= bInsert?PROTECTION_INSERT_KEY:PROTECTION_REMOVE_KEY;

    flProtectionOutput out;  memset(&out, 0, sizeof(out));

    DWORD nReturned=0;
    if (!DeviceIoControl(hDsk, FL_IOCTL_BINARY_HW_PROTECTION, &in, sizeof(in), &out, sizeof(out), &nReturned, NULL)) {
        error("FL_IOCTL_BINARY_HW_PROTECTION-%s-key", bInsert?"insert":"remove");
        return false;
    }

    return true;
}

bool PartitionChangeType(HANDLE hDsk, DWORD dwType, const std::string& passwd)
{
    BYTE in[10];

    in[0]= (BYTE)dwType;
    memcpy(in+1, passwd.c_str(), 8);
    in[9]= PROTECTION_CHANGE_TYPE;

    flProtectionOutput out;  memset(&out, 0, sizeof(out));

    DWORD nReturned=0;
    if (!DeviceIoControl(hDsk, FL_IOCTL_BDTL_HW_PROTECTION, &in, sizeof(in), &out, sizeof(out), &nReturned, NULL)) {
        error("FL_IOCTL_BDTL_HW_PROTECTION-changetype(%d)", dwType);
        return false;
    }

    return true;
}


bool unlockdisk(HANDLE hDsk, const std::string& passwd)
{
    PartitionInsertKey(hDsk, true, passwd);
    BDKPartitionInsertKey(hDsk, true, 0, passwd);
    BDKPartitionInsertKey(hDsk, true, 1, passwd);
    PartitionChangeType(hDsk, PROTECTABLE, passwd);

    CloseHandle(hDsk);

    return true;
}

bool lockdisk(HANDLE hDsk, const std::string& passwd)
{
    PartitionChangeType(hDsk, PROTECTABLE|WRITE_PROTECTED, passwd);
    PartitionInsertKey(hDsk, false, passwd);
    BDKPartitionInsertKey(hDsk, false, 0, passwd);
    BDKPartitionInsertKey(hDsk, false, 1, passwd);

    CloseHandle(hDsk);

    return true;
}

bool SetRegistryDword(HKEY hive, TCHAR *szPath, TCHAR *szKey, DWORD dwValue)
{
    HKEY hk;
    LONG res= RegOpenKeyEx(hive, szPath, 0, 0, &hk);
    if (res!=ERROR_SUCCESS) {
        error(res, "ERROR opening regpath %08lx %ls\n", hive, szPath);
        return false;
    }
    res= RegSetValueEx(hk, szKey, 0, REG_DWORD, (BYTE*)&dwValue, sizeof(DWORD));
    if (res!=ERROR_SUCCESS) {
        error(res, "ERROR setting regvalue %08lx %ls : %ls\n", hive, szPath, szKey);
        return false;
    }
    RegCloseKey(hk);

    return true;
}
bool RemoveRegistryValue(HKEY hive, TCHAR *szPath, TCHAR *szKey)
{
    HKEY hk;
    LONG res= RegOpenKeyEx(hive, szPath, 0, 0, &hk);
    if (res!=ERROR_SUCCESS) {
        //error(res, "ERROR opening regpath %ls\n", szPath);
        return false;
    }
    res= RegDeleteValue(hk, szKey);
    if (res!=ERROR_SUCCESS) {
        //error(res, "ERROR setting regvalue %ls\n", szKey);
        return false;
    }
    RegCloseKey(hk);

    return true;
}
bool hidedisk()
{
    return SetRegistryDword(HKEY_LOCAL_MACHINE, L"System\\StorageManager\\Profiles\\TRUEFFS_DOC", L"MountFlags", 1);
}
bool unhidedisk()
{
    return RemoveRegistryValue(HKEY_LOCAL_MACHINE, L"System\\StorageManager\\Profiles\\TRUEFFS_DOC", L"MountFlags");
}
int WINAPI WinMain( HINSTANCE hInstance,
                   HINSTANCE hPrevInstance,
                   LPTSTR    lpCmdLine,
                   int       nCmdShow)
{
    KernelMode _km;

    DebugMessagebox();

    bool bDoLockDevice= false;
    bool bDoUnlockDevice= false;
    bool bDoHideDevice= false;
    bool bDoUnhideDevice= false;
    std::string passwd= XDA2PASSWORD;
    int disknr=0;
    DWORD dwHandle=0;


    StringList args;
    if (!SplitString(ToString(lpCmdLine), args, false)) {
        error("Error in commandline");
        return false;
    }
    for (StringList::iterator i= args.begin() ; i!=args.end() ; ++i)
    {
        std::string arg= *i;
		chomp(arg);
        if (arg=="-lock") {
            bDoLockDevice= true;
        }
        else if (arg=="-unlock") {
            bDoUnlockDevice= true;
        }
        else if (arg=="-hide") {
            bDoHideDevice= true;
        }
        else if (arg=="-unhide") {
            bDoUnhideDevice= true;
        }
        else if (arg=="-xda3") {
            passwd= XDA3PASSWORD;
        }
        else if (arg=="-passwd") {
            if (++i == args.end())
                break;
            passwd= *i;
        }
        else if (arg=="-handle") {
            if (++i == args.end())
                break;
            dwHandle= strtoul((*i).c_str(), 0, 0);
        }
    }
    if (!bDoLockDevice && !bDoUnlockDevice && !bDoHideDevice && !bDoUnhideDevice)
        debug("Usage: unlockextrom [-lock | -unlock] [-hide | -unhide] [-handle HANDLE] [-passwd PASSWD] [-xda3]");

    // default disk is DSK1:
    if (disknr==0 && dwHandle==0)
        disknr= 1;

    if (disknr) {
        dwHandle= (DWORD)OpenDisk(disknr);
        if (dwHandle==NULL) {
            error("Open");
            return 1;
        }
    }

    if (bDoLockDevice) {
        if (!lockdisk((HANDLE)dwHandle, passwd))
            return 1;
        debug("extrom locked");
    }
    if (bDoUnlockDevice) {
        if (!unlockdisk((HANDLE)dwHandle, passwd))
            return 1;
        debug("extrom unlocked");
    }

    if (bDoHideDevice) {
        if (!hidedisk())
            return 1;
        debug("extrom hidden, reboot to finish");
    }
    if (bDoUnhideDevice) {
        if (!unhidedisk())
            return 1;
        debug("extrom revealed, reboot to finish");
    }

    return 0;
}