#include <windows.h>
#include <stdio.h>

#include <imagehlp.h>

// defined in mscoree.h, need to link with mscoree.lib
extern "C" HRESULT STDMETHODCALLTYPE GetCORVersion(LPWSTR pbuffer, DWORD cchBuffer, DWORD* dwlength);

// the result xml file to store dumped assembly info
extern FILE *resultXmlFile;

// the simple exe file name, not full path
extern char simpleExeFileName[]; 

// the default dump directory
extern char dumpDir[];

/* number of pad bytes to make 'len' bytes align to 'align' */
inline static unsigned roundUp(unsigned len, unsigned align) {
    return((len + align-1) & ~(align-1));
}

inline static unsigned padLen(unsigned len, unsigned align) {
    return(roundUp(len, align) - len);
}

// it's v1.0 or v1.1, 1.0.3705 or v1.1.4322
inline bool isVersion_1()
{	
	WCHAR version[MAX_PATH+1];
	DWORD len = 0;
	GetCORVersion(version, MAX_PATH, &len);

	return ((char)version[1]) == '1' && 
		   ((char)version[2]) == '.';
}

// save the memory into PE Format:
// header, section1, section2, ...
bool SaveFile(PBYTE pbImageBase, WCHAR *modulename)
{	
	PIMAGE_NT_HEADERS pNTHeader = NULL;
	char msg[4098];

	// only dump executable images
	__try {
        PIMAGE_DOS_HEADER pDosHeader = (PIMAGE_DOS_HEADER)pbImageBase;
        if (pDosHeader->e_magic != IMAGE_DOS_SIGNATURE) {
            return false;
        }

        pNTHeader = (PIMAGE_NT_HEADERS)((PBYTE)pDosHeader +
                                                          pDosHeader->e_lfanew);
        if (pNTHeader->Signature != IMAGE_NT_SIGNATURE) {
            return false;
        }        
    } __except(EXCEPTION_EXECUTE_HANDLER) {
       return false;
    }

	WCHAR filename[MAX_PATH+1];

	// retrieve simple name
	WCHAR *p = wcsrchr(modulename, '\\');
	if (p) {
		wcscpy(filename, p+1);
	} else {
		wcscpy(filename, modulename);
	}

	// figure out output file name, saved into sub dir "dumped" in the folder 
	char outputfile[MAX_PATH+1];			
	if (filename[0] != 0) {
		sprintf(outputfile, "%s\\%ls", dumpDir, filename);		
	} else
		sprintf(outputfile, "%s\\Image_%x.dll", dumpDir, pbImageBase);
	
	int numSections = pNTHeader->FileHeader.NumberOfSections;

	PIMAGE_SECTION_HEADER section = (PIMAGE_SECTION_HEADER)(pNTHeader + 1);		

	// the section header ends here, from here to the start of the first section, we can fill zeros
	PBYTE pLastSectionEnd = (PBYTE)(section + numSections); 
 
	ULONG file_alignment = pNTHeader->OptionalHeader.FileAlignment;	

	// write into a file
	FILE* stream = fopen(outputfile, "w+b" );
	if (stream == NULL ) {
	   return false;
	}				

	printf("pbImageBase[0]: 0x%x, first section starts at offset: 0x%x\n", pbImageBase[0], section->PointerToRawData);

	// write out header
	int headerLen = pLastSectionEnd-pbImageBase;
	int numwritten = fwrite( pbImageBase, 1, headerLen, stream );
	if (numwritten != headerLen) {
		printf("fwrite error, file: %s\n", outputfile);
	}

	// fill zeros upto the start of the first section
	char zero = 0;
	for (int i=headerLen; i<section->PointerToRawData; i++) { // roundUp(section->PointerToRawData, file_alignment)
		fwrite( &zero, 1, 1, stream );
	}

	MEMORY_BASIC_INFORMATION mbi;
	VirtualQuery(pbImageBase, &mbi, sizeof(mbi));
	bool isMapped;
	
	if (isVersion_1())
		isMapped = (mbi.Type == MEM_IMAGE) || (mbi.Type == MEM_MAPPED); // mapped as image, save using PE headers
	else 
		isMapped = (mbi.Type == MEM_IMAGE); // v2.0.50727

	// write each section 
	for (int i=0; i<numSections; i++) {
		PBYTE buf;
		if (isMapped) {
			buf = pbImageBase + section->VirtualAddress;
		} else {
			buf = pbImageBase + section->PointerToRawData;					
		}

		//sprintf(msg, "Writing section - name: %s, buf: 0x%x, len: 0x%x, PointerToRawData: 0x%x, mbi.Type:0x%x", 
		//	section->Name, buf, section->SizeOfRawData, section->PointerToRawData, mbi.Type);
		//MessageBoxA(NULL, msg, "Info", MB_OK |MB_ICONINFORMATION);

		//printf("Writing section - name: %s, buf: 0x%x, len: 0x%x, PointerToRawData: 0x%x, mbi.Type:0x%x\n", 
		//	section->Name, buf, section->SizeOfRawData, section->PointerToRawData, mbi.Type);

		numwritten = fwrite(buf, 1, section->SizeOfRawData, stream);
		if (numwritten != section->SizeOfRawData) {
			printf("fwrite error, file: %s\n", outputfile);
		}		

		section ++;
	}

	// Close stream 
	fclose( stream ); 
	
	char fullpath[MAX_PATH+1];
	char *simpleName;
	GetFullPathNameA(outputfile, MAX_PATH, fullpath, &simpleName);

	// store info into result file
	fprintf(resultXmlFile, "<module>\r\n");
	fprintf(resultXmlFile, "<name>%s</name>\r\n", simpleName);
	fprintf(resultXmlFile, "<imagebase>0x%x</imagebase>\r\n", pbImageBase);
	fprintf(resultXmlFile, "<filename>%ls</filename>\r\n", modulename);
	fprintf(resultXmlFile, "<dump>%s</dump>\r\n", fullpath);
	fprintf(resultXmlFile, "</module>\r\n");

	sprintf(msg, "Dump module 0x%x, saved into %s", pbImageBase, outputfile);
	MessageBoxA(NULL, msg, simpleExeFileName, MB_OK |MB_ICONINFORMATION|MB_DEFAULT_DESKTOP_ONLY);
	return true;
}