#include <windows.h>
#include <psyscall.h>
#include <storemgr.h>
#include "cenk.h"
#include "kernelmisc.h"
#include "ceioctl.h"
#include "debug.h"
#include "stringutils.h"

#define BDK_SIGNATURE "BIPO"

std::string cinfostring(const CINFO *ci)
{
    if (ci)
        return stringformat("cinfo{%c%c%c%c, d=%02x t=%02x n=%d, svr=%08lx}", 
                ci->acName[0], ci->acName[1], ci->acName[2], ci->acName[3], 
                ci->disp, ci->type, ci->cMethods, ci->pServer);
    else
        return "(null)";
}
std::string handleinfostring(HANDLE h)
{
	HDATA *hi= cvHandle2HDataPtr(h);
    return stringformat("hdata{%hs hv=%08lx l=%08lx c=%08lx pci=%08lx o=%08lx i=%08lx}", 
                cinfostring(hi->pci).c_str(), hi->hValue, hi->lock, hi->ref.count, hi->pci, hi->pvObj, hi->dwInfo);
}

HANDLE WINAPI OpenStore(LPCTSTR szDeviceName)
{
    debug("OpenStore('%ls')\n", szDeviceName);
    return CreateFile( ToWString(stringformat("\\StoreMgr\\%ls", szDeviceName)).c_str(), 0, 0, NULL, 0, 0, NULL);
}
bool OpenTFFSDisk(const WCHAR *szDeviceName, const WCHAR *szPartitionName, HANDLE& hStore, HANDLE& hPartition)
{
    debug("OpenTFFSDisk('%ls', '%ls')\n", szDeviceName, szPartitionName);
    if (szDeviceName) {
        hStore = OpenStore(szDeviceName);
        if (hStore == INVALID_HANDLE_VALUE || hStore == NULL) {
            error("OpenTFFSDisk: OpenStore('%ls')", szDeviceName);
            return false;
        }
        debug("hStore: %08lx  %hs\n", hStore, handleinfostring(hStore).c_str());
        hPartition = PSLOpenPartition(hStore, szPartitionName);
        if (hPartition == INVALID_HANDLE_VALUE || hPartition == NULL) {
            error("PSLOpenPartition('%ls')", szPartitionName);
            CloseHandle(hStore);
            hStore= NULL;
            return false;
        }
        debug("hPartition: %08lx  %hs\n", hPartition, handleinfostring(hPartition).c_str());
    }
    else {
        hStore= NULL;
        hPartition= CreateFile( szPartitionName,
            GENERIC_READ, FILE_SHARE_READ|FILE_SHARE_WRITE,
            NULL, OPEN_EXISTING, 0, 0 );
        debug("hPartition: %08lx  %hs\n", hPartition, handleinfostring(hPartition).c_str());
    }

    return true;
}
bool TFFS_GetUniqueId(HANDLE hDisk)
{
    flUniqueIdOutput out; memset(&out, 0, sizeof(out));

    DWORD nReturned;
    if (!DeviceIoControl(hDisk, FL_IOCTL_UNIQUE_ID, NULL, 0, &out, sizeof(out), &nReturned, NULL)) {
        error("DeviceIoControl(FL_IOCTL_UNIQUE_ID)");
        return false;
    }
    if (out.status != flOK) {
        error("DeviceIoControl(FL_IOCTL_UNIQUE_ID) - status=%08lx", out.status);
        SetLastError(ERROR_INTERNAL_ERROR);
        return false;
    }
    debug("getuniqueid: %hs\n", hexdump(out.id, 16).c_str());
    return true;
}
bool TFFS_GetNrOfPartitions(HANDLE hDisk, DWORD& dwPartitionCount)
{
    flCountPartitionsOutput out; memset(&out, 0, sizeof(out));

    DWORD nReturned;
    if (!DeviceIoControl(hDisk, FL_IOCTL_NUMBER_OF_PARTITIONS, NULL, 0, &out, sizeof(out), &nReturned, NULL)) {
        error("DeviceIoControl(FL_IOCTL_NUMBER_OF_PARTITIONS)");
        return false;
    }
    if (out.status != flOK) {
        error("DeviceIoControl(FL_IOCTL_NUMBER_OF_PARTITIONS) - status=%08lx", out.status);
        SetLastError(ERROR_INTERNAL_ERROR);
        return false;
    }
    debug("nrofpartitions: %d\n", out.noOfPartitions);
    dwPartitionCount= out.noOfPartitions;
    return true;
}
bool TFFS_BDK_GetInfo(HANDLE hDisk, DWORD dwBinaryPartitionNr, DWORD &dwLength, DWORD &dwCount)
{
    flBDKOperationInput in;  memset(&in, 0, sizeof(in));

    in.partitionNumber = (BYTE)dwBinaryPartitionNr;
    in.type = BDK_GET_INFO;
    memcpy(in.bdkStruct.oldSign, BDK_SIGNATURE, strlen(BDK_SIGNATURE));
    in.bdkStruct.signOffset = 8;
    in.bdkStruct.startingBlock = 0;

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

    DWORD nReturned;
    if (!DeviceIoControl(hDisk, FL_IOCTL_BDK_OPERATION, &in, sizeof(in), &out, sizeof(out), &nReturned, NULL)) {
        error("DeviceIoControl(FL_IOCTL_BDK_OPERATION, BDK_GET_INFO)");
        return false;
    }
    if (out.status != flOK) {
        error("DeviceIoControl(FL_IOCTL_BDK_OPERATION, BDK_GET_INFO) - status=%08lx", out.status);
        SetLastError(ERROR_INTERNAL_ERROR);
        return false;
    }
    debug("bdk_getinfo: size=%08lx used=%08lx count=%08lx\n", 
            in.bdkStruct.startingBlock, in.bdkStruct.length, in.bdkStruct.flags);

    dwLength= in.bdkStruct.startingBlock;
    dwCount= in.bdkStruct.flags;

    return true;
}
bool TFFS_OTP_GetSize(HANDLE hDisk, DWORD& dwSize)
{
    flOtpInput in;  memset(&in, 0, sizeof(in));
    flOutputStatusRecord out; memset(&out, 0, sizeof(out));

    in.type= OTP_SIZE;

    DWORD nReturned;
    if (!DeviceIoControl(hDisk, FL_IOCTL_OTP, &in, sizeof(in), &out, sizeof(out), &nReturned, NULL)) {
        error("DeviceIoControl(FL_IOCTL_OTP)");
        return false;
    }
    if (out.status != flOK) {
        error("DeviceIoControl(FL_IOCTL_OTP) - status=%08lx", out.status);
        SetLastError(ERROR_INTERNAL_ERROR);
        return false;
    }
    debug("otp_getsize:  len=%08lx used=%08lx flag=%08lx\n",
            in.length, in.usedSize, in.lockedFlag);

    dwSize= in.usedSize;
    return true;
}
bool TFFS_OTP_Read(HANDLE hDisk, DWORD dwSize)
{
    flOtpInput in;  memset(&in, 0, sizeof(in));
    flOutputStatusRecord out; memset(&out, 0, sizeof(out));

    BYTE localbuf[6144]; memset(localbuf, 0, sizeof(localbuf));

    in.type   = OTP_READ;
    in.usedSize= 0;
    in.length = dwSize<sizeof(localbuf) ? dwSize : sizeof(localbuf);
    in.buffer= localbuf;

    DWORD nReturned;
    if (!DeviceIoControl(hDisk, FL_IOCTL_OTP, &in, sizeof(in), &out, sizeof(out), &nReturned, NULL)) {
        error("DeviceIoControl(FL_IOCTL_OTP)");
        return false;
    }
    if (out.status != flOK) {
        error("DeviceIoControl(FL_IOCTL_OTP) - status=%08lx", out.status);
        SetLastError(ERROR_INTERNAL_ERROR);
        return false;
    }
    for (int i=0 ; i<sizeof(localbuf) ; i+=256)
        debug("%hs\n", hexdump(i, localbuf+i, 256).c_str());

    return true;
}
bool TFFS_GetCustomerId(HANDLE hDisk)
{
    flCustomerIdOutput out; memset(&out, 0, sizeof(out));

    DWORD nReturned;
    if (!DeviceIoControl(hDisk, FL_IOCTL_CUSTOMER_ID, NULL, 0, &out, sizeof(out), &nReturned, NULL)) {
        error("DeviceIoControl(FL_IOCTL_CUSTOMER_ID)");
        return false;
    }
    if (out.status != flOK) {
        error("DeviceIoControl(FL_IOCTL_CUSTOMER_ID) - status=%08lx", out.status);
        SetLastError(ERROR_INTERNAL_ERROR);
        return false;
    }
    debug("getcustomerid: %hs\n", hexdump(out.id, 4).c_str());
    return true;
}
std::string GetCInfoName(const CINFO *ci)
{
    std::string name;
	for (int i=0 ; i<4 ; i++)
		name += ci->acName[i];
	return name;
}

std::string GetHandleListString(const std::string& magic, HANDLE hSeed)
{
    StringList lines;  lines.push_back("handle list");
	HDATA *hi= cvHandle2HDataPtr(hSeed);
	HDATA *ha;
	ha=hi;
	do {
        if (GetCInfoName(ha->pci)==magic) {
		    lines.push_back(stringformat("%08lx", ha->hValue));
        }

		ha= (HDATA*)ha->linkage.fwd;
	} while (ha!=hi && ha!=NULL);

    if (ha==NULL)
        debug("WARNING: ha==NULL\n");

    return JoinStringList(lines, "\n");
}
int WINAPI WinMain( HINSTANCE hInstance,
                   HINSTANCE hPrevInstance,
                   LPTSTR    lpCmdLine,
                   int       nCmdShow)
{
    DebugSetLogfile("opendisk.log");

    KernelMode _km;

    debug("cmdline='%ls'\n", lpCmdLine);
    StringList args;
    if (!SplitString(ToString(lpCmdLine), args)) {
        debug("error splitting into args\n");
        return 1;
    }

    debug("cmdline - found %d args\n", args.size());

    if (args.size()==2) {
        debug("opening('%hs', '%hs')\n", args[0].c_str(), args[1].c_str());

        HANDLE hStore, hPartition;
        if (!OpenTFFSDisk(ToWString(args[0]).c_str(), ToWString(args[1]).c_str(), hStore, hPartition)) {
            error("OpenTFFSDisk");
            return 1;
        }

        MessageBox(0,ToWString(GetHandleListString("STRG", hPartition)).c_str(),L"title",0);

        CloseHandle(hPartition);
        CloseHandle(hStore);

        debug("done\n");
        return 0;
    }
    else {
        HANDLE hPartition= (HANDLE)strtoul(args[0].c_str(), 0, 0);

        DWORD dwNrParts;
        TFFS_GetNrOfPartitions(hPartition, dwNrParts);

        DWORD dwNrBParts=0;
        DWORD dwLength=0;
        TFFS_BDK_GetInfo(hPartition, 0, dwLength, dwNrBParts);
        for (DWORD p=1 ; p<dwNrBParts ; p++)
            TFFS_BDK_GetInfo(hPartition, p, dwLength, dwNrBParts);

        DWORD dwOtpSize;
        TFFS_OTP_GetSize(hPartition, dwOtpSize);
        if (dwOtpSize)
            TFFS_OTP_Read(hPartition, dwOtpSize);

        TFFS_GetCustomerId(hPartition);

        TFFS_GetUniqueId(hPartition);
    }

    return 0;
}