Cleanup, and correctly handle connection closing

3 files changed, 42 insertions, 84 deletions

diff --git a/console.c b/console.c
index 558184f..2ac2068 100644
--- a/console.c
+++ b/console.c
@@ -15,7 +15,5 @@ int getTerminalSize(unsigned short *width, unsigned short *height) {
 	*width = ws.ws_col;
 	*height = ws.ws_row;
 
-	printf("Console width: %d, height: %d\n", *width, *height);
-
 	return 1;
 }
diff --git a/main.c b/main.c
index 5bc0180..bde22c1 100644
--- a/main.c
+++ b/main.c
@@ -16,6 +16,7 @@
     with this program; if not, write to the Free Software Foundation, Inc.,
     51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 */
+#include <stdlib.h>
 #include <stdio.h>
 #include <arpa/inet.h>
 #include <netinet/in.h>
@@ -35,11 +36,38 @@ int outcounter=0;
 int sessionkey=0;
 unsigned char *src = "00:e0:81:b5:ac:8e";
 unsigned char dstmem[] = "00:0c:42:43:58:a4";
-unsigned char *dst = &dstmem;
+unsigned char *dst = dstmem;
 unsigned char encryptionkey[128];
 unsigned char username[255];
 unsigned char password[255];
 
+void sendAuthData(unsigned char *username, unsigned char *password) {
+	unsigned char data[1500];
+	unsigned char *terminal = "linux";
+	int userLen = strlen(username);
+	int terminalLen = strlen(terminal);
+	unsigned short width = 0;
+	unsigned short height = 0;
+	int result;
+	int plen;
+	int databytes;
+
+	plen = initPacket(data, MT_PTYPE_DATA, src, dst, sessionkey, outcounter);
+	databytes = plen;
+	plen += addControlPacket(data + plen, MT_CPTYPE_PASSWORD, password, 17);
+	plen += addControlPacket(data + plen, MT_CPTYPE_USERNAME, username, userLen);
+	plen += addControlPacket(data + plen, MT_CPTYPE_TERM_TYPE, terminal, terminalLen);
+
+	if (getTerminalSize(&width, &height) > 0) {
+		plen += addControlPacket(data + plen, MT_CPTYPE_TERM_WIDTH, &width, 2);
+		plen += addControlPacket(data + plen, MT_CPTYPE_TERM_HEIGHT, &height, 2);
+	}
+
+	outcounter += plen - databytes;
+
+	result = sendCustomUDP(sockfd, src, dst, "213.236.240.252", 20561, "255.255.255.255", 20561, data, plen);
+}
+
 void handlePacket(unsigned char *data, int data_len) {
 	struct mt_mactelnet_hdr pkthdr;
 	struct mt_mactelnet_control_hdr cpkthdr;
@@ -96,33 +124,19 @@ void handlePacket(unsigned char *data, int data_len) {
 			}
 		}
 	}
-}
-
-void sendAuthData(unsigned char *username, unsigned char *password) {
-	unsigned char data[1500];
-	unsigned char *terminal = "linux";
-	int userLen = strlen(username);
-	int terminalLen = strlen(terminal);
-	unsigned short width = 0;
-	unsigned short height = 0;
-	int result;
-	int plen;
-	int databytes;
-
-	plen = initPacket(data, MT_PTYPE_DATA, src, dst, sessionkey, outcounter);
-	databytes = plen;
-	plen += addControlPacket(data + plen, MT_CPTYPE_PASSWORD, password, 17);
-	plen += addControlPacket(data + plen, MT_CPTYPE_USERNAME, username, userLen);
-	plen += addControlPacket(data + plen, MT_CPTYPE_TERM_TYPE, terminal, terminalLen);
-
-	if (getTerminalSize(&width, &height) > 0) {
-		plen += addControlPacket(data + plen, MT_CPTYPE_TERM_WIDTH, &width, 2);
-		plen += addControlPacket(data + plen, MT_CPTYPE_TERM_HEIGHT, &height, 2);
+	else if (pkthdr.ptype == MT_PTYPE_ACK) {
+		// TODO: If we were resubmitting lost messages, stop resubmitting here if received counter is correct.
+	}
+	else if (pkthdr.ptype == MT_PTYPE_END) {
+		char odata[200];
+		int plen=0,result=0;
+		plen = initPacket(odata, MT_PTYPE_END, src, dst, pkthdr.seskey, 0);
+		result = sendCustomUDP(sockfd, src, dst, "213.236.240.252", 20561, "255.255.255.255", 20561, odata, plen);
+		fprintf(stderr, "Connection closed.\n");
+		exit(0);
+	} else {
+		fprintf(stderr, "Unhandeled packet type: %d received from server %s\n", pkthdr.ptype, dst);
 	}
-
-	outcounter += plen - databytes;
-
-	result = sendCustomUDP(sockfd, src, dst, "213.236.240.252", 20561, "255.255.255.255", 20561, data, plen);
 }
 
 int main (int argc, char **argv) {
@@ -199,19 +213,12 @@ int main (int argc, char **argv) {
 	result = recvfrom(insockfd, buff, 1500, 0, 0, 0);
 	handlePacket(buff, result);
 
-	memset(buff, 0, 1500);
-	result = recvfrom(insockfd, buff, 1500, 0, 0, 0);
-	handlePacket(buff, result);
-
-	memset(buff, 0, 1500);
-	result = recvfrom(insockfd, buff, 1500, 0, 0, 0);
-	handlePacket(buff, result);
-
 while(1) {
 	memset(buff, 0, 1500);
 	result = recvfrom(insockfd, buff, 1500, 0, 0, 0);
 	handlePacket(buff, result);
 }
+	
 	close(sockfd);
 	close(insockfd);
 
diff --git a/udp.c b/udp.c
index 96fd906..2881f89 100644
--- a/udp.c
+++ b/udp.c
@@ -10,52 +10,6 @@
 #include <linux/if_arp.h>
 #include <stdio.h>
 
-unsigned short udp_sum_calc(unsigned short len_udp, unsigned char src_addr[],unsigned char dest_addr[], unsigned char buff[]) {
-	unsigned short prot_udp=17;
-	unsigned short padd=0;
-	unsigned short word16;
-	unsigned long sum;
-	int i;	
-	
-	// Find out if the length of data is even or odd number. If odd,
-	// add a padding byte = 0 at the end of packet
-	if (len_udp % 2 == 1){
-		padd=1;
-		buff[len_udp]=0;
-	}
-	
-	//initialize sum to zero
-	sum=0;
-	
-	// make 16 bit words out of every two adjacent 8 bit words and 
-	// calculate the sum of all 16 vit words
-	for (i=0;i<len_udp+padd;i=i+2){
-		word16 =((buff[i]<<8)&0xFF00)+(buff[i+1]&0xFF);
-		sum = sum + (unsigned long)word16;
-	}	
-	// add the UDP pseudo header which contains the IP source and destinationn addresses
-	for (i=0;i<4;i=i+2){
-		word16 =((src_addr[i]<<8)&0xFF00)+(src_addr[i+1]&0xFF);
-		sum=sum+word16;	
-	}
-	for (i=0;i<4;i=i+2){
-		word16 =((dest_addr[i]<<8)&0xFF00)+(dest_addr[i+1]&0xFF);
-		sum=sum+word16; 	
-	}
-	// the protocol number and the length of the UDP packet
-	sum = sum + prot_udp + len_udp;
-
-	// keep only the last 16 bits of the 32 bit calculated sum and add the carries
-    	while (sum>>16)
-		sum = (sum & 0xFFFF)+(sum >> 16);
-		
-	// Take the one's complement of sum
-	sum = ~sum;
-
-	return ((unsigned short) sum);
-}
-
-//#define ETH_FRAME_LEN 1518
 unsigned short in_cksum(unsigned short *addr, int len)
 {
 	int nleft = len;
@@ -161,7 +115,6 @@ int sendCustomUDP(const int socket, const char *sourcemac, const char *destmac,
 	udp->dest = htons(20561);
 	udp->check = 0;
 	udp->len = htons(sizeof(struct udphdr) + datalen);
-	//udp->check = udp_sum_calc(datalen+8, (unsigned char *)&(ip->saddr), (unsigned char *)&(ip->daddr), (unsigned char *)udp);
 
 	memcpy(resten, data, datalen);