.386
.model flat,stdcall
option casemap:none

include windows.inc
include kernel32.inc
include user32.inc
include comdlg32.inc

includelib user32.lib
includelib kernel32.lib
includelib comdlg32.lib

DialogProc proto :HWND, :UINT, :WPARAM, :LPARAM
ListLog proto :DWORD

  
szText MACRO Name, Text:VARARG
	LOCAL lbl
	jmp lbl
	Name db Text,0
	lbl:
ENDM

.const
IDD_DLG1 		equ 1000
IDC_EDT1		equ 1001
IDC_OPENFILE	equ 1002
IDC_ADDSECTION 	equ 1003
IDC_EXIT 		equ 1004
IDC_LST1		equ 1005
OPENFILE_FLAGS 	equ OFN_PATHMUSTEXIST+OFN_FILEMUSTEXIST+OFN_LONGNAMES+OFN_HIDEREADONLY
NewSectionSize	equ	350h
Characteristics equ IMAGE_SCN_MEM_EXECUTE+\
					IMAGE_SCN_MEM_READ+\
					IMAGE_SCN_MEM_WRITE+\
					IMAGE_SCN_CNT_CODE+\
					IMAGE_SCN_CNT_INITIALIZED_DATA

.data?
hInst HINSTANCE ?
hFile HANDLE ?
hMapObject HANDLE ?
Ofn OPENFILENAME <>
pMapView dd ?
szFileName db MAX_PATH dup(?)
szBuffer db MAX_PATH dup(?)
dwSize dd ?
g_hwndDlg HWND ?
pSectionName DWORD ?


.data
ValidFile BOOL FALSE


szSectionName db ".pickle",0
szBackup db ".bak",0
lpstrFilter db "Executable File",0,"*.exe",0,"Dynamic Link Library",0,"*.dll",0,0
szErrorFile db "Select a proper file!",0
szCaptionError db "Error",0
szSucess_MapFile db "Sucesfully mapped file to RVA: 0x%08X",0
szSucess_ValidPE db "Valid file detected",0
szSucess_CreateFile db "File ready to be mapped",0
szFail_MapFile db "Failed to map file!",0
szFail_CreateFile db "File could not be opened!",0
szFail_Generic db "Operation Failed",0
szFail_ValidPE db "Not a valid Executable or Dynamic Link Library!",0
szFail_CopyFile db "Failed to copy file",0
szFail_AddSection db "Failed to allocate a new section!",0
szSucess_FileLoaded db "File: %s loaded",0
szSucess_AddSection db "Section added Sucesfully!",0

.code
start:

invoke GetModuleHandle,NULL
mov hInst,eax
invoke DialogBoxParam,hInst,IDD_DLG1,NULL,addr DialogProc,NULL
invoke ExitProcess,0

DialogProc proc uses edx ecx hwndDlg:HWND, uMsg:UINT, wParam:WPARAM, lParam:LPARAM

	mov eax,hwndDlg
	mov g_hwndDlg,eax
	
	.if uMsg == WM_COMMAND
		mov eax,wParam
		.if eax == IDC_EXIT
			invoke SendMessage,hwndDlg,WM_CLOSE,NULL,NULL
		.elseif eax == IDC_ADDSECTION
			.if ValidFile == TRUE
				push offset szFileName
				call AddSection
				.if eax == TRUE
					invoke ListLog,addr szSucess_AddSection
					call CloseMappedFile
				.else
					invoke ListLog,addr szFail_AddSection
				.endif
			.else
				invoke ListLog,addr szErrorFile
			.endif
		.elseif eax == IDC_OPENFILE
			mov Ofn.lStructSize,sizeof OPENFILENAME
			mov eax,g_hwndDlg
			mov Ofn.hwndOwner,eax
			mov eax, hInst
			mov Ofn.hInstance,eax
			mov Ofn.lpstrFilter,offset lpstrFilter
			mov Ofn.lpstrFile,offset szFileName
			mov Ofn.nMaxFile,MAX_PATH
			mov Ofn.nFilterIndex,1
			mov Ofn.Flags, OPENFILE_FLAGS
			invoke GetOpenFileName,addr Ofn
			.if eax == TRUE
				invoke wsprintf,addr szBuffer,addr szSucess_FileLoaded,Ofn.lpstrFile
				invoke ListLog,addr szBuffer
				invoke SetDlgItemText,hwndDlg,IDC_EDT1,Ofn.lpstrFile
								
				push 0
				push FILE_MAP_READ
				push PAGE_READONLY
				push FILE_ATTRIBUTE_NORMAL
				push GENERIC_READ   
				push offset szFileName
				call MapFileToRVA
				mov pMapView,eax
				.if pMapView != NULL
					invoke wsprintf,addr szBuffer,addr szSucess_MapFile,pMapView
					invoke ListLog,addr szBuffer
					push pMapView		
					call ValidPEFile
					.if eax == TRUE
						mov ValidFile,eax
						invoke ListLog,addr szSucess_ValidPE
					.elseif eax == FALSE
						invoke ListLog,addr szFail_ValidPE
					.else
						invoke ListLog,addr szFail_Generic
					.endif
				.else
					invoke ListLog,addr szFail_Generic
				.endif
				call CloseMappedFile
			.endif
		.endif
	.elseif uMsg == WM_CLOSE
		invoke EndDialog,hwndDlg,0
	.endif
	
	xor eax,eax
	ret
	
DialogProc endp

AddSection proc uses ecx ebx edx edi esi pFileName:DWORD
	
	LOCAL RetVal:DWORD
	
	invoke lstrcpy,addr szBuffer,pFileName
	invoke lstrcat,addr szBuffer,addr szBackup
	invoke CopyFile,pFileName,addr szBuffer,FALSE
	
	.if eax != TRUE
		invoke ListLog,addr szFail_CopyFile	
		mov eax,FALSE
		jmp @exit
	.endif
	
	
	;Align Section properly
	mov eax,NewSectionSize
	mov ecx,100h
	xor edx,edx
	div ecx
	.if edx != 0
		inc eax
	.endif
	mul ecx
	
	push eax
	push FILE_MAP_ALL_ACCESS
	push PAGE_READWRITE
	push FILE_ATTRIBUTE_NORMAL
	push GENERIC_READ+GENERIC_WRITE  
	push pFileName
	call MapFileToRVA
	mov pMapView,eax
	
	.if eax == NULL
		invoke ListLog,addr szFail_Generic
		call CloseMappedFile
		mov eax,FALSE
		jmp @exit
	.endif
	
	
	;Populate new section
	mov esi,pMapView
	add esi,[esi.IMAGE_DOS_HEADER.e_lfanew]
	movsx edx,[esi.IMAGE_NT_HEADERS.FileHeader.NumberOfSections]
	mov eax,sizeof IMAGE_SECTION_HEADER
	imul eax,edx
	add eax,sizeof IMAGE_NT_HEADERS
	
	
	;ESI - pNtHeader
	;EBX - pNewSection
	;EDI - pPrevSection
	
	
	;Write new section info
	add eax,esi 
	mov ebx,eax
	
	;Calc size of raw data
	mov eax,NewSectionSize
	mov [ebx.IMAGE_SECTION_HEADER.Misc.VirtualSize],eax
	
	xor edx,edx		;edx=0
	mov ecx,100h	;filealignment
	div ecx			;div eax 300 / ecx 100
	.if edx != 0	;remainder?
		inc eax		;inc eax
	.endif
	mul ecx			;mul eax,ecx
	mov [ebx.IMAGE_SECTION_HEADER.SizeOfRawData],eax
	
	;Align new section RVA
	;VirtualSize aligned to Section + VirtualAddress = new Section RVA
	mov edi,ebx	
	sub edi,sizeof IMAGE_SECTION_HEADER
	mov eax,[edi.IMAGE_SECTION_HEADER.Misc.VirtualSize]
	mov ecx,1000h ;SectionAlignment
	xor edx,edx
	div ecx
	.if edx !=0
		inc eax
	.endif		
	mul ecx
	add eax,[edi.IMAGE_SECTION_HEADER.VirtualAddress]
	mov [ebx.IMAGE_SECTION_HEADER.VirtualAddress],eax
	
	;Align new section raw address
	mov eax,[edi.IMAGE_SECTION_HEADER.SizeOfRawData]
	add eax,[edi.IMAGE_SECTION_HEADER.PointerToRawData]
	mov [ebx.IMAGE_SECTION_HEADER.PointerToRawData],eax
	
	;Write Characteristics
	mov eax,Characteristics
	mov [ebx.IMAGE_SECTION_HEADER.Characteristics],eax
	
	;Update PE Header
	inc [esi.IMAGE_NT_HEADERS.FileHeader.NumberOfSections]
	add [esi.IMAGE_NT_HEADERS.OptionalHeader.SizeOfHeaders], sizeof IMAGE_SECTION_HEADER
	
	;Re-Align SizeOfHeaders!
	mov eax,[esi.IMAGE_NT_HEADERS.OptionalHeader.SizeOfHeaders]
	mov ecx,100h
	xor edx,edx
	div ecx
	.if edx != 0
		inc eax
	.endif
	mul ecx
	mov [esi.IMAGE_NT_HEADERS.OptionalHeader.SizeOfHeaders],eax
	
	;Re-Align SizeOfImage!	
	mov eax,[ebx.IMAGE_SECTION_HEADER.Misc.VirtualSize]
	mov ecx,1000h
	xor edx,edx
	div ecx
	.if edx != 0
		inc eax
	.endif
	mul ecx
	add [esi.IMAGE_NT_HEADERS.OptionalHeader.SizeOfImage],eax
	
	invoke lstrcpy,ebx,addr szSectionName
	
	mov eax,TRUE
	
	@exit:
	mov RetVal,eax
	ret

AddSection endp

MapFileToRVA proc \
		pFileName:DWORD,\
		dwDesiredAccessFile:DWORD,\
		dwFlagAndAttribFile:DWORD,\
		flProtect:DWORD,\
		dwDesiredAccessMap:DWORD,\
		dwSectionSize:DWORD
	
	LOCAL RetVal:DWORD
	LOCAL NewSize:DWORD
	
	invoke CreateFile,pFileName,dwDesiredAccessFile,0,0,OPEN_EXISTING,dwFlagAndAttribFile,NULL
	mov hFile,eax
	.if eax != INVALID_HANDLE_VALUE
		.if dwSectionSize != 0
			invoke GetFileSize,hFile,NULL
			add eax,dwSectionSize
			mov NewSize,eax
		.else
			xor eax,eax
			mov NewSize,eax
		.endif
		invoke CreateFileMapping,hFile,NULL,flProtect,0,NewSize,NULL
		mov hMapObject,eax
		invoke GetLastError
		.if hMapObject != NULL
			invoke MapViewOfFile,hMapObject,dwDesiredAccessMap,0,0,NewSize
			mov RetVal,eax
		.else
			invoke ListLog,addr szFail_MapFile
			mov RetVal,NULL
			jmp @exit
		.endif
	.else
		invoke ListLog,addr szFail_CreateFile
		mov RetVal,NULL
		jmp @exit
	.endif
	
	@exit:
	mov eax,RetVal
	ret

MapFileToRVA endp

CloseMappedFile proc 
	
	invoke UnmapViewOfFile,pMapView
	invoke CloseHandle,hMapObject
	invoke CloseHandle,hFile
	mov pMapView,NULL
	mov hMapObject,NULL
	mov hFile,NULL
	ret

CloseMappedFile endp

ValidPEFile proc pMap:DWORD
	
	LOCAL RetVal:DWORD
	
	mov eax,pMap
	
	.if eax != NULL
		.if [eax.IMAGE_DOS_HEADER.e_magic] == IMAGE_DOS_SIGNATURE
			add eax,[eax.IMAGE_DOS_HEADER.e_lfanew]
			.if [eax.IMAGE_NT_HEADERS.Signature] == IMAGE_NT_SIGNATURE
				mov RetVal,TRUE
			.else
				mov RetVal,FALSE
				jmp @exit	
			.endif
		.else
			mov RetVal,FALSE
			jmp @exit
		.endif
	.else
		mov RetVal,FALSE
	.endif
	
	@exit:
	invoke CloseMappedFile
	mov eax,RetVal
	ret

ValidPEFile endp

ListLog proc szLBString:DWORD

	invoke SendDlgItemMessage,g_hwndDlg,IDC_LST1,LB_ADDSTRING,NULL,szLBString
	ret

ListLog endp

end start