;simple TCP/IP client, written in fasm, displays menu which allows user to connect to server, send/recieve data, and close connection
format PE console
INCLUDE 'win32a.inc'

section '' code readable executable
start:
pop [exit] ;get exit address
push msg_startinfo
call [puts]
add esp, 4
sub esp, 512
mov eax, esp
push eax
push 0x202
call [WSAStartup] ;init winsock
add esp, 512 ;clean up stack
test eax, eax
jnz error ;error
xor ebx, ebx
cmdrtn:
push dword '>'
call [putchar]
sub esp, 508 ;512 byte buffer, 4 bytes on stack
mov esi, esp ;save cmdline ptr to esi
mov eax, esp
mov ecx, 511
call getline
mov eax, [esp]
cmp al, 'c' ;c, connect
je connect_server
cmp al, 'd' ;d, disconnect
je disconnect_server
cmp al, 'r' ;r, recv
je recv_server
cmp al, 's' ;s, send
je send_server
cmp al, 'q' ;q, quit
je terminate_program
cmp al, '?' ;?, help
je disp_help
push msg_invcmd
call [puts]
add esp, 4
nextcmd:
add esp, 512 ;restore stack
jmp cmdrtn
terminate_program:
push msg_cleanup
call [puts]
add esp, 4
call [WSACleanup]
jmp [exit]

;fatal error handler
error:
push msg_fatal_error
call [puts]
add esp, 4
jmp [exit]

;connect to server, with hostname/ip in eax
connect_server:
cmp [connected], 0
jz @f
push msg_already_connected
call [puts]
add esp, 4
jmp nextcmd
@@:
mov eax, ebx
push 6
push 1
push 2
call [socket] ;create socket
cmp eax, not 0
jnz @f
push msg_socket_failed
call [puts]
add esp, 4
push ebx
call [closesocket]
xor ebx, ebx
jmp nextcmd
@@:
mov ebx, eax ;save connection in ebx
push dword ' %d'
push dword 'c %s'
mov eax, esp
push port
push destname
push eax
push esi
call [sscanf]
add esp, 20 ;clean up stack 
push msg_get_inet_addr
call [puts]
sub esp, 12 ;SOCKADDR_IN, 16 bytes with 4 already present
mov ebp, esp
push destname
call [inet_addr]
cmp eax, 0xffffffff
jne addrknown
push destname
call [gethostbyname]
test eax, eax
jnz @f
push destname
push msg_hostname_error
call [printf]
add esp, 24
xor ebx, ebx
jmp nextcmd
@@:
mov eax, [eax+12]
mov eax, [eax]
mov eax, [eax] 
addrknown:
mov word [ebp], 2
mov ecx, [port]
xchg cl, ch ;network byte order
mov [ebp+2], cx
mov [ebp+4], eax
xor eax, eax
mov [ebp+8], eax
mov [ebp+12], eax
push msg_connect
call [puts]
add esp, 4
push 16
push ebp
push ebx
call [connect] ;connect to server
cmp eax, -1
jnz @f
push msg_connect_failed
call [puts]
add esp, 4
push ebx
call [closesocket]
xor ebx, ebx
add esp, 16
jmp nextcmd
@@:
add esp, 16 ;clean up stack
mov [connected], 1
jmp nextcmd

;send message to server
send_server:
test ebx, ebx
jz noconnect
sub esp, 512 ;512 byte buffer
mov ebp, esp
mov eax, esp
mov ecx, 511
call getline
push 0
push eax
push ebp
push ebx
call [send]
add esp, 512
jmp nextcmd

;get data from server
recv_server:
test ebx, ebx
jz noconnect
sub esp, 512 ;512 byte buffer
mov ebp, esp
push 0
push 511
push ebp
push ebx
call [recv]
mov ecx, ebp
add ecx, eax
xor al, al
mov [ecx], al
push ebp
call [puts]
add esp, 516 ;restore stack
jmp nextcmd

;disconnect from server
disconnect_server:
test ebx, ebx
jz noconnect
push 2
push ebx
call [shutdown]
push msg_closesocket
call [puts]
add esp, 4
push ebx
call [closesocket] ;disconnect
xor ebx, ebx
mov [connected], 0
jmp nextcmd

;no connection exist
noconnect:
push msg_no_connection
call [puts]
add esp, 4
jmp nextcmd

;disp help
disp_help:
push msg_help
call [puts]
add esp, 4
jmp nextcmd

;read line from stdin (terminated by \n), pointer to string in eax, maximum length of string in ecx, length of string returned in eax
getline:
push ebp esi edi
cld
mov esi, ecx
mov edi, eax
xor ebp, ebp
push [_iob] ;address of stdin
call [fflush]
add esp, 4
.next:
test esi, esi
jz .done
call [getchar]
stosb
dec esi
inc ebp
cmp al, 10
jnz .next
.done:
xor al, al
stosb
mov eax, ebp
pop edi esi ebp
ret

section '' data readable writable
port dd 0
exit dd 0
connected dd 0
msg_fatal_error db 'Error occured', 0
msg_get_inet_addr db 'Determining internet address...', 0
msg_connect db 'Connect to server...', 0
msg_closesocket db 'Closing connection to server...', 0
msg_startinfo db 'TCP client v0.1', 13, 10, 'Enter ? to get help', 0
msg_cleanup db 'Close winsock...', 0
msg_invcmd db 'Invalid command', 0
msg_hostname_error db 'Could not resolve hostname %s', 13, 10, 0
msg_connect_failed db 'Could not connect to server', 0
msg_no_connection db 'Not connected to server', 0
msg_socket_failed db 'socket create failed', 0
msg_already_connected db 'Already connected', 0
msg_help db 'TCP client v0.1', 13, 10, 'c server port, connect to server (hostname or ip accepted) server on port port', 13, 10, 's, send typed line', 13, 10, 'r, get data from server', 13, 10, 'd, disconnect', 13, 10, 'q, quit', 13, 10, '?, help', 0
destname rb 256

section '' import readable
library kernel32, 'kernel32.dll', ws2_32, 'ws2_32.dll', msvcrt, 'msvcrt.dll'
import kernel32, GetLastError, 'GetLastError', SetConsoleCtrlHandler, 'SetConsoleCtrlHandler', SetConsoleTextAttribute, 'SetConsoleTextAttribute', GetStdHandle, 'GetStdHandle'
import ws2_32, WSAStartup, 'WSAStartup', socket, 'socket', gethostbyaddr, 'gethostbyaddr', gethostbyname, 'gethostbyname', inet_addr, 'inet_addr', connect, 'connect',\
 send, 'send', recv, 'recv', shutdown, 'shutdown', closesocket, 'closesocket', WSACleanup, 'WSACleanup'
import msvcrt, printf, 'printf', puts, 'puts', _getch, '_getch', getchar, 'getchar', putchar, 'putchar', fflush, 'fflush', _iob, '_iob', system, 'system', sscanf, 'sscanf'

section '' fixups discardable