comment %
	%

.486
.model flat,stdcall
option casemap:none

include \tasm\win32.inc
include \tasm\extrn.inc

sSEH STRUCT
	OrgEsp            DD 0
	OrgEbp            DD 0
	SaveEip           DD 0
sSEH ENDS

pesection struc		; IMAGE_SECTION_HEADER - simplified
	name		db 8 dup (?)
	vsz		dd ?
	voff		dd ?
	sz		dd ?
	off		dd ?
	junk		db 12 dup (?)
	char		dd ?
pesection ends

.DATA
about		db "PE Nightmare2b by Freddy K",0
about2		db "About PE Nightmare2",0
filter  	db 'PE-executables (*.exe)',0,'*.exe',0,0
buffer		db MAX dup (0)
header		db 'Please Pick the Target File',0
notpe		db "Chosen file is not a Win32 PE file.",13,10
		db "Please choose another file",0
cantadd		db "Win32 Object Table error - can't add.",13,10
		db "Please choose another file",0
noimagehlp	db "Can't load PE Checkum Library - skipping",0
cantopen	db "Can't open file",0
donemsg		db "Section Added !!",0
OpenStruct:	dd 76, 0
		dd 0
		dd offset filter, 0,0,0, offset buffer, MAX, 0,0,0, offset header
		dd OFN_FILEMUSTEXIST or OFN_HIDEREADONLY or OFN_EXPLORER
		dd 0, 0, 0, 0, 0
fhandle		dd 0
fsize		dd 0
buff		dd 0

library		db MAX dup (0)
imagehlp	db "\imagehlp.dll",0
mapandcheck	db "MapFileAndCheckSumA",0
libHwnd		dd 0
checksum	dd 0		; new checksum calculated

lib_err         db "A required library file is not present in the "
                db "path or is corrupt.",13,10,13,10,"Library name : "
aplib		db "aplib026.dll",0
approc		db "_aP_pack",0
apHwnd		dd 0
workmem		dd 0 		; 640*1024 dup (?)
progress_bar	dd 0
progress_var	dd 0

pMem		dd 0		; global alloc mem for file

modhand		dd 0
winHwnd		dd 0
aboutHwnd	dd 0

peHdr		dd 0		; PE hdr offset
sections	dd 0		; original
image_sz	dd 0		; original
section_vsz	dd 0		; new section size (SECTION_VSZ) rounded to FILE_
section_sz	dd 0		; new section real size (section_sz) rounded to FILE_
imagebase	dd 0
section1_diff	dd 0		; oldsection1_sz-newsection1_sz (for changing section offsets)
section1_off	dd 0
section2_off	dd 0

FILE_ALIGN	dd 0200h	; default values
SECTION_ALIGN	dd 1000h

fk		pesection <0>
win_dim         RECT <?>

.CONST
MAX			equ 255
WORK			equ 640*1024
SPEED_VAR		equ 10		; higher num==slower progress bar

IDD_DLG		 	equ 100		; Main dialog id
IDB_GO			equ 200
IDB_GETFILE		equ 201
IDB_EXIT		equ 202

IDT_box1                equ 300 	; input box
IDT_OLD			equ 301		; old size
IDT_NEW			equ 302		; new size
IDT_LIB			equ 303		; lib name

ID_PROGRESS		equ 350

ID_EXIT			equ 400

MIN_KERNEL_SEARCH_BASE	equ 070000000h
MAX_API_STRING_LENGTH	equ 150

LENGTH			equ (offset section_end - offset section)
CHAR			equ 0C0000040h

.CODE
main:
	call	InitCommonControls

	call	GetModuleHandleA, 0
	mov	[modhand], eax

	call	FindWindowA, 0, offset about
	.if eax!=0
		jmp	end
	.endif

	call	DialogBoxParamA, [modhand], IDD_DLG, 0, offset dialogproc, 0
end:
	call	ExitProcess,0

	ret

;----------------------------------------------------------------------------
dialogproc proc hwnd:DWORD, wmsg:DWORD, wparam:DWORD, lparam:DWORD
;----------------------------------------------------------------------------
	cmp     [wmsg], WM_COMMAND
        je      button
        cmp     [wmsg], WM_CLOSE
        je      exitprog
	cmp	[wmsg], WM_INITDIALOG
	je	dialoginit

	xor	eax, eax
	ret

; -------------------------------------
button:
; -------------------------------------
	cmp	[wparam], IDB_EXIT
	je	exitprog
	cmp	[wparam], IDB_GETFILE
	je	get_file
	cmp	[wparam], IDB_GO
	je	compress
	cmp	[wparam], ID_EXIT
	je	exitprog

	ret

; -------------------------------------
dialoginit:
; -------------------------------------
	call	SetWindowText, hwnd, offset about
	call	centre_win, hwnd

	call    LoadLibraryA, offset aplib
	.if eax==0
		jmp	no_lib
	.endif
	mov	[apHwnd],eax	; store the handle of the opened dll

	call	SetDlgItemText, hwnd, IDT_LIB, offset aplib

	ret
no_lib:
	call	MessageBoxA, 0, offset lib_err, offset about, MB_OK+MB_ICONSTOP
	jmp	end

; -------------------------------------
get_file:
; -------------------------------------
	call 	GetOpenFileNameA, offset OpenStruct
	.if eax==0
		ret
	.endif
       	call    SetDlgItemTextA, [hwnd], IDT_box1, offset buffer

	ret

; -----------------------------------------
compress:
; -----------------------------------------
        call 	CreateFileA, offset buffer, GENERIC_READ+GENERIC_WRITE,\
		0,0, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL,0
	.if eax==INVALID_HANDLE_VALUE
		jmp	cant_open
	.endif
	mov	fhandle, eax
	call	GetFileSize, fhandle, 0

	call	roundup, eax, FILE_ALIGN
	mov	fsize, eax
	call	SetDlgItemInt, hwnd, IDT_OLD, eax, 0 
	call	roundup, LENGTH, SECTION_ALIGN
	mov	section_vsz, eax			; store virtual vsz

	add	eax, fsize
	add	eax, SECTION_ALIGN			; = FILESIZE+SECTION(ROUND)+BUFFER
	call	GlobalAlloc, GMEM_FIXED OR GMEM_ZEROINIT, eax
	mov	pMem, eax

	call	ReadFile, fhandle, pMem, fsize, offset buff, 0
	mov	esi, pMem				; esi now points to file
	.if word ptr [esi] != IMAGE_DOS_SIGNATURE
		jmp	not_pe
	.endif
	add	esi, 03ch
	mov	eax, dword ptr [esi]
	mov	peHdr, eax
	sub	eax, 03ch
	add	esi, eax				; now points to PE HDR
	.if dword ptr [esi] != IMAGE_NT_SIGNATURE
		jmp	not_pe
	.endif
	call	CloseHandle, fhandle			; close to create_new later

; ----------- Adjust old PE Header ------------------------------
	assume	esi:ptr IMAGE_NT_HEADERS
	mov	eax, [esi].OptionalHeader.SectionAlignment
	mov	SECTION_ALIGN, eax			; use progs section align (voff)

	mov	eax, FILE_ALIGN				; use OUR FILE align (off)
	mov	[esi].OptionalHeader.FileAlignment, eax
	call	roundup, LENGTH, FILE_ALIGN
	mov	section_sz, eax				; store real sz

	mov	eax, [esi].OptionalHeader.AddressOfEntryPoint	; gets Entry point
	mov	ecx, [esi].OptionalHeader.ImageBase		; OEP == base+EP
	mov	imagebase, ecx
	add	eax, ecx
	mov	oep, eax					; --!!! .CODE VARIABLE !!!--

	mov	edi, pMem
	mov	edx, peHdr				; now points to PE HDR
	add	edx, 0f8h				; now 1st section offset				
	add	edi, edx

	assume	edi:ptr pesection
	mov	eax, [edi].sz
	mov	section1_orig, eax			; for decompress later
	mov	eax, [edi].off
	mov	section1_off, eax
	mov	eax, [edi].voff
	mov	section1_voff, eax
	mov	eax, imagebase
	add	section1_voff, eax			; store for decompress offset
	mov	[edi].char, CHAR			; make first section writable
	add	edi, size pesection
	mov	eax, [edi].off
	mov	section2_off, eax			; store second section offset
	sub	edi, size pesection			; adjust back for next routines
	assume	edi:

	xor	eax, eax
	mov	ax, [esi].FileHeader.NumberOfSections
	mov	sections, eax
	mov	ecx, size pesection			; each section size=28h bytes
	imul	ecx

	inc	[esi].FileHeader.NumberOfSections	; sections++
	add	edi, eax				; now points to NEW last section
	add	eax, edx				; eax == offset new last section
	
	.if eax > [esi].OptionalHeader.SizeOfHeaders || byte ptr [edi] != NULL
		jmp	cant_add
	.endif

	mov	eax, [esi].OptionalHeader.SizeOfImage
	mov	[esi].OptionalHeader.AddressOfEntryPoint, eax	; NEW OEP = OLD EOF (image_sz)

	mov	eax, [esi].OptionalHeader.SizeOfImage
	mov	image_sz, eax
	mov	eax, section_vsz
	add	[esi].OptionalHeader.SizeOfImage, eax		; image size w/new section added
	assume	esi:
; ---------------------------------------------------------------

; create our section hdr
	mov	dword ptr [fk.name], "erf."
	mov	dword ptr [fk.name+4], "kydd"
	mov	eax, section_vsz
	mov	fk.vsz, eax				; virtual size
	mov	eax, image_sz
	mov	fk.voff, eax				; virtual offset
	mov	eax, section_sz
	mov	fk.sz, eax				; real size
	mov	eax, fsize
	mov	fk.off, eax				; real offset

; copy our section hdr to file
	mov	esi, offset fk
	mov	ecx, size pesection
	rep movsb

; ------------ COMPRESS SECTION1 ---------------------------------------
	call	GlobalAlloc, GMEM_FIXED OR GMEM_ZEROINIT, WORK
	mov	workmem, eax
	call	GlobalAlloc, GMEM_FIXED OR GMEM_ZEROINIT, section1_orig
	mov	pMemRd, eax	

	call	GetProcAddress, apHwnd, offset approc
	push	esi
	mov	esi, pMem
	add	esi, section1_off

	mov	progress_bar, 0
	mov	edx, hwnd
	mov	winHwnd, edx		; enable function to get GLOBAL var
	call	eax, esi, pMemRd, section1_orig, workmem, offset progress
	mov	progress_bar, 0

	call	roundup, eax, FILE_ALIGN
	mov	section1_comp, eax		; new section size
	mov	ecx, section1_orig
	sub	ecx, eax
	mov	section1_diff, ecx 

	mov	esi, pMemRd
	mov	edi, pMem
	add	edi, section1_off
	mov	ecx, section1_orig			; copies compressed section->file
	rep movsb					; and buffers rest w/ zeros	
	pop	esi

	call	GlobalFree, pMemRd
	call	GlobalFree, workmem

	mov	eax, imagebase
	add	section1_off, eax
; --------------------------------------------------------------------

add_section:
	mov	edi, pMem
	add	edi, fsize				; edi now points to EOF
	lea	eax, section
	xchg	esi, eax				; esi == new section code to add in 
	mov	ecx, LENGTH
	rep movsb

; ---------- adjust section info - WRITABLE and NEW off ---------------
	mov	esi, pMem
	add	esi, peHdr
	add	esi, 0f8h				; points to first section HDR

	assume	esi:ptr pesection
	call	roundup, [esi].vsz, SECTION_ALIGN
	mov	[esi].vsz, eax				; make section voff aligned
	mov	eax, section1_comp
	mov	[esi].sz, eax
	mov	eax, section1_diff
	mov	ecx, sections
section_offsets:
	add	esi, size pesection
	sub	[esi].off, eax				; adjust to new offset
	mov	[esi].char, CHAR			; make WRITABLE
	push	eax
	call	roundup, [esi].vsz, SECTION_ALIGN
	mov	[esi].vsz, eax				; make section voff aligned
	pop	eax
	loop	section_offsets
	assume	esi:
; ----------------------------------------------------------------------

; Make backup HERE if req'd
;
        call 	CreateFileA, offset buffer, GENERIC_READ or GENERIC_WRITE,\
		FILE_SHARE_READ+FILE_SHARE_WRITE,0, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL,0
	.if eax==INVALID_HANDLE_VALUE
		jmp	cant_open
	.endif
	mov	fhandle, eax

	mov	eax, section1_off
	add	eax, section1_comp
	sub	eax, imagebase
	mov	checksum, eax				; just using a spare var
	call 	WriteFile, [fhandle], pMem, eax, offset buff, 0		; write hdr+section1

	mov	eax, fsize
	add	eax, section_sz
	sub	eax, section2_off					; write section2++
	mov	esi, pMem
	add	esi, section2_off
	add	checksum, eax
	call 	WriteFile, [fhandle], esi, eax, offset buff, 0		; write the rest

	call	SetDlgItemInt, hwnd, IDT_NEW, checksum, 0

; ----- UPDATE CHECKSUM ---------------------------------------------
 	call	GetSystemDirectory, offset library, MAX
	call	lstrcat, offset library, offset imagehlp
	call    LoadLibraryA, offset library			; imagehlp.dll
	.if eax==0
		jmp	no_chk_lib
	.endif
	mov	libHwnd, eax
	call    GetProcAddress, eax, offset mapandcheck		; mapandchecksum
	call	eax, offset buffer, offset buff, offset checksum

	mov	esi, pMem
	mov	eax, peHdr
	add	eax, 58h
	add	esi, eax				; now points to checksum offset
	mov	ecx, checksum
	mov	dword ptr [esi], ecx
	call	SetFilePointer, fhandle, eax,0, FILE_BEGIN
	call 	WriteFile, [fhandle], offset checksum, 4, offset buff, 0
	call	FreeLibrary, libHwnd
; --------------------------------------------------------------------

	call	MessageBoxA, 0, offset donemsg, offset about, MB_OK

close_file:
	call	GlobalFree, pMem
	call	CloseHandle, fhandle

	ret
not_pe:
	call	MessageBoxA, 0, offset notpe, offset about, MB_OK+MB_ICONSTOP
	jmp	close_file
cant_add:
	call	MessageBoxA, 0, offset cantadd, offset about, MB_OK+MB_ICONSTOP
	jmp	close_file
no_chk_lib:
	call	MessageBoxA, 0, offset noimagehlp, offset about, MB_OK+MB_ICONSTOP
	jmp	close_file
cant_open:
	call	MessageBoxA, 0, offset cantopen, offset about, MB_OK+MB_ICONSTOP
	jmp	close_file

; ----------------------------
exitprog:
; ----------------------------
	call	FreeLibrary, apHwnd
	call	EndDialog, [hwnd], 0

	ret

endp dialogproc

; --------------------------------------------------
progress:	; edx=orig, ecx=new
; --------------------------------------------------
	pushad