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

int dumpBDKZone_MR(char *storeData, int partitionNumber, int binaryPartitionNumber, int startBlock, unsigned char *buffer, int length);
int dumpBDKZone_IT(TCHAR *szDeviceName, TCHAR *szPartitionName, DWORD dwBinaryPartitionNumber, DWORD dwOffset, unsigned char *buffer, int dwSize);

void initbuf(unsigned char *buf, int n)
{
    for (int i=0 ; i<n ; i++) 
        buf[i]=(unsigned char)i&0xff;
}
int WINAPI WinMain( HINSTANCE hInstance,
                   HINSTANCE hPrevInstance,
                   LPTSTR    lpCmdLine,
                   int       nCmdShow)
{
    DebugSetLogfile("tsttffswork.log");

    debug("starting tsttffswork\n");
	unsigned char *buffer = (unsigned char*)malloc(0x10000);
    if (buffer==NULL) {
        error("malloc");
        return 1;
    }
    initbuf(buffer, 0x10000);

    debug("testing arisme version\n");
    int rc= dumpBDKZone_MR("TrueFFS", 0, 1, 0, buffer, 0x10000);

    debug("result= %d\n", rc);

    debug("MR: %hs\n", hexdump(buffer, 0x100).c_str());

    initbuf(buffer, 0x10000);

    debug("testing itsme version\n");
    rc= dumpBDKZone_IT(L"TrueFFS", L"Part00", 1, 0, buffer, 0x10000);

    debug("result= %d\n", rc);

    debug("IT: %hs\n", hexdump(buffer, 0x100).c_str());

    return 0;
}


#define BDK_SIGNATURE "BIPO"
#define TFFS_BLOCK_SIZE 0x10000
#define TFFS_SECTOR_SIZE 512

int dumpBDKZone_MR(char *storeData, int partitionNumber, int binaryPartitionNumber, int startBlock, unsigned char *buffer, int length)
{
	flBDKOperationInput BDKInput = {0};
	flOutputStatusRecord BDKOutput;
	HANDLE storageHandle;
	TCHAR storageFile[100];
	TCHAR partitionFile[10];
	HANDLE partHandle;
	DWORD dummy;
	unsigned char TFFSBuffer[TFFS_SECTOR_SIZE];
	int remainingLength = length;
	int currentOffset = 0;

	if (length % 512)
		return 10;

	wsprintf(storageFile, TEXT("\\StoreMgr\\%hs"), storeData);
	storageHandle = CreateFile(storageFile, 0, 0, NULL, 0, 0, NULL);
	if (storageHandle == INVALID_HANDLE_VALUE) {
		return 1;
	}
	wsprintf(partitionFile, TEXT("Part%.2d"), partitionNumber);
	partHandle = PSLOpenPartition(storageHandle, partitionFile);
	if (partHandle == INVALID_HANDLE_VALUE) {
		CloseHandle(storageHandle);
		return 2;
	}

	debug("readbin('%ls', '%ls', %d, %08lx, %08lx, %08lx)\n", storageFile, partitionFile, binaryPartitionNumber, startBlock, buffer, length);

	BDKInput.partitionNumber = binaryPartitionNumber;
	BDKInput.type = BDK_INIT_READ;
	memcpy(BDKInput.bdkStruct.oldSign, BDK_SIGNATURE, strlen(BDK_SIGNATURE));
	BDKInput.bdkStruct.signOffset = 8;
	BDKInput.bdkStruct.startingBlock = startBlock;
	BDKInput.bdkStruct.length = length;

	debug("bdkinit: %hs\n", hexdump((BYTE*)&BDKInput, sizeof(BDKInput)/sizeof(DWORD), sizeof(DWORD)).c_str());

	if (!DeviceIoControl(partHandle, FL_IOCTL_BDK_OPERATION, &BDKInput, sizeof(BDKInput), &BDKOutput, sizeof(BDKOutput), &dummy, NULL)) {
		CloseHandle(partHandle);
		CloseHandle(storageHandle);
		return 3;
	}
	if (BDKOutput.status != flOK) {
		CloseHandle(partHandle);
		CloseHandle(storageHandle);
		return 4;
	}

	while (remainingLength) {
		int toRead = (remainingLength > TFFS_SECTOR_SIZE ? TFFS_SECTOR_SIZE : remainingLength);
		BDKInput.type = BDK_READ;
		BDKInput.bdkStruct.length = TFFS_SECTOR_SIZE;
		BDKInput.bdkStruct.bdkBuffer = TFFSBuffer;

        initbuf(TFFSBuffer, TFFS_SECTOR_SIZE);

        debug("-----------in: %08lx out: %08lx\n", &BDKInput, &BDKOutput);
		debug("bdkread: %hs\n", hexdump((BYTE*)&BDKInput, sizeof(BDKInput)/sizeof(DWORD), sizeof(DWORD)).c_str());
		debug("    pre: %hs\n", hexdump(TFFSBuffer, 16).c_str());
		if (!DeviceIoControl(partHandle, FL_IOCTL_BDK_OPERATION, &BDKInput, sizeof(BDKInput), &BDKOutput, sizeof(BDKOutput), &dummy, NULL)) {
			CloseHandle(partHandle);
			CloseHandle(storageHandle);
			return 5;
		}
		debug("   post: %hs\n", hexdump(TFFSBuffer, 16).c_str());

		if (BDKOutput.status != flOK) {
			CloseHandle(partHandle);
			CloseHandle(storageHandle);
			return 6;
		}

		memcpy(buffer + currentOffset, TFFSBuffer, toRead);
		currentOffset += toRead;
		remainingLength -= toRead;
	}

	return 0;
}

//..............................

HANDLE WINAPI OpenStore(LPCTSTR szDeviceName)
{
    TCHAR szFileName[MAX_PATH];
    memset( szFileName, 0, sizeof(szFileName));
    wcscpy( szFileName, L"\\StoreMgr\\");
    __try {
        wcsncat( szFileName, szDeviceName, MAX_PATH-15);
    } __except(EXCEPTION_EXECUTE_HANDLER) {
        SetLastError(ERROR_BAD_ARGUMENTS);
    }
    return CreateFile( szFileName, 0, 0, NULL, 0, 0, NULL);
}

bool OpenTFFSDisk(const WCHAR *szDeviceName, const WCHAR *szPartitionName, HANDLE& hStore, HANDLE& hPartition)
{
    hStore = OpenStore(szDeviceName);
    if (hStore == INVALID_HANDLE_VALUE || hStore == NULL) {
        error("OpenTFFSDisk: OpenStore('%ls')", szDeviceName);
        return false;
    }
    hPartition = PSLOpenPartition(hStore, szPartitionName);
    if (hPartition == INVALID_HANDLE_VALUE || hPartition == NULL) {
        error("PSLOpenPartition('%ls')", szPartitionName);
        CloseHandle(hStore);
        hStore= NULL;
        return false;
    }

    return true;
}

bool TFFS_BDK_InitRead(HANDLE hDisk, DWORD dwBinaryPartitionNr, DWORD dwOffset, DWORD dwSize)
{
    flBDKOperationInput BDKInput;  memset(&BDKInput, 0, sizeof(BDKInput));

    BDKInput.partitionNumber = (BYTE)dwBinaryPartitionNr;
    BDKInput.type = BDK_INIT_READ;
    memcpy(BDKInput.bdkStruct.oldSign, BDK_SIGNATURE, strlen(BDK_SIGNATURE));
    BDKInput.bdkStruct.signOffset = 8;
    BDKInput.bdkStruct.startingBlock = dwOffset/TFFS_BLOCK_SIZE;
    BDKInput.bdkStruct.length = dwSize;

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

    debug("bdkinit: %hs\n", hexdump((BYTE*)&BDKInput, sizeof(BDKInput)/sizeof(DWORD), sizeof(DWORD)).c_str());
    DWORD nReturned;
    if (!DeviceIoControl(hDisk, FL_IOCTL_BDK_OPERATION, &BDKInput, sizeof(BDKInput), &BDKOutput, sizeof(BDKOutput), &nReturned, NULL)) {
        error("DeviceIoControl(FL_IOCTL_BDK_OPERATION, BDK_INIT_READ)");
        return false;
    }
    if (BDKOutput.status != flOK) {
        error("DeviceIoControl(FL_IOCTL_BDK_OPERATION, BDK_INIT_READ) - status=%08lx", BDKOutput.status);
        SetLastError(ERROR_INTERNAL_ERROR);
        return false;
    }
    return true;
}

bool TFFS_BDK_ReadNextSector(HANDLE hDisk, DWORD dwBinaryPartitionNr, DWORD dwOffset, BYTE *buffer)
{
    flBDKOperationInput BDKInput;  memset(&BDKInput, 0, sizeof(BDKInput));
    unsigned char localbuf[TFFS_SECTOR_SIZE];


    BDKInput.partitionNumber = (BYTE)dwBinaryPartitionNr;
    BDKInput.type = BDK_READ;
    memcpy(BDKInput.bdkStruct.oldSign, BDK_SIGNATURE, strlen(BDK_SIGNATURE));
    BDKInput.bdkStruct.signOffset = 8;
    BDKInput.bdkStruct.startingBlock = dwOffset/TFFS_BLOCK_SIZE;
    BDKInput.bdkStruct.length = TFFS_SECTOR_SIZE;
    BDKInput.bdkStruct.bdkBuffer = localbuf;

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

    initbuf(localbuf, TFFS_SECTOR_SIZE);

    debug("-----------in: %08lx out: %08lx\n", &BDKInput, &BDKOutput);
    debug("bdkread: %hs\n", hexdump((BYTE*)&BDKInput, sizeof(BDKInput)/sizeof(DWORD), sizeof(DWORD)).c_str());
    debug("    pre: %hs\n", hexdump(localbuf, 16).c_str());
    DWORD nReturned;
    if (!DeviceIoControl(hDisk, FL_IOCTL_BDK_OPERATION, &BDKInput, sizeof(BDKInput), &BDKOutput, sizeof(BDKOutput), &nReturned, NULL)) {
        error("DeviceIoControl(FL_IOCTL_BDK_OPERATION, BDK_READ)");
        return false;
    }
    if (BDKOutput.status != flOK) {
        error("DeviceIoControl(FL_IOCTL_BDK_OPERATION, BDK_READ) - status=%08lx", BDKOutput.status);
        SetLastError(ERROR_INTERNAL_ERROR);
        return false;
    }
    debug("   post: %hs\n", hexdump(localbuf, 16).c_str());

    memcpy(buffer, localbuf, TFFS_SECTOR_SIZE);

    return true;
}

int dumpBDKZone_IT(TCHAR *szDeviceName, TCHAR *szPartitionName, DWORD dwBinaryPartitionNr, DWORD dwOffset, unsigned char *buffer, int dwSize)
{
    if (dwSize%TFFS_SECTOR_SIZE) {
        return ERROR_INVALID_PARAMETER;
    }
    HANDLE hStore, hPartition;

    if (!OpenTFFSDisk(szDeviceName, szPartitionName, hStore, hPartition))
        return GetLastError();

// using hStore here will return ERROR_DEVICE_NOT_AVAILABLE
// using hPartition here will return ERROR_INVALID_HANDLE

    debug("readbin('%ls', '%ls', %d, %08lx, %08lx, %08lx)\n", szDeviceName, szPartitionName, dwBinaryPartitionNr, dwOffset, buffer, dwSize);

    if (!TFFS_BDK_InitRead(hPartition, dwBinaryPartitionNr, dwOffset, dwSize))
        return GetLastError();

    DWORD dwRemainingLength= dwSize;
    DWORD dwBufOffset= 0;
    while (dwRemainingLength)
    {
        DWORD dwWanted= (dwRemainingLength > TFFS_SECTOR_SIZE ? TFFS_SECTOR_SIZE : dwRemainingLength);

        if (!TFFS_BDK_ReadNextSector(hPartition, dwBinaryPartitionNr, dwOffset, buffer+dwBufOffset))
            break;
        dwBufOffset += dwWanted;
        dwRemainingLength -= dwWanted;
    }

    debug("read %08lx bytes\n", dwBufOffset);

    return 0;
}