flesh out the Python client more to add connect/disconnect commands, and fix various bugs

4 files changed, 152 insertions, 54 deletions

diff --git a/buffer.h b/buffer.h
index df063e9..1218809 100644
--- a/buffer.h
+++ b/buffer.h
@@ -145,6 +145,28 @@ public:
 	int32_t readS32() { int32_t v; read((char *)&v, 4); return v; }
 	int16_t readS16() { int16_t v; read((char *)&v, 2); return v; }
 	int8_t readS8() { int8_t v; read((char *)&v, 1); return v; }
+
+	void readStr(char *output, int bufferSize) {
+		uint32_t size = readU32();
+		if (!readRemains(size)) {
+			strcpy(output, "");
+			return;
+		}
+
+		// How much can we safely get?
+		int readAmount;
+		if (size < (bufferSize - 1))
+			readAmount = size;
+		else
+			readAmount = bufferSize - 1;
+
+		// Put this into the buffer
+		read(output, readAmount);
+		output[readAmount] = 0;
+
+		// In case the buffer was too small, skip over the extra source data
+		m_readPointer += (size - readAmount);
+	}
 };
 
 #endif /* BUFFER_H */
diff --git a/core.cpp b/core.cpp
index 17092c1..b4b7c82 100644
--- a/core.cpp
+++ b/core.cpp
@@ -45,9 +45,9 @@ static bool setSocketNonBlocking(int sock) {
 static bool isNullSessionKey(uint8_t *key) {
 	for (int i = 0; i < SESSION_KEY_SIZE; i++)
 		if (key[i] != 0)
-			return true;
+			return false;
 
-	return false;
+	return true;
 }
 
 static Client *findClientWithKey(uint8_t *key) {
@@ -297,14 +297,15 @@ void Client::clearCachedPackets(int maxID) {
 }
 
 
-void Client::handleLine(char *line, int size) {
+void Client::handleCommand(char *line, int size) {
 	// This is a terrible mess that will be replaced shortly
 	if (authState == AS_AUTHED) {
 		if (strncmp(line, "all ", 4) == 0) {
-			for (int i = 0; i < clientCount; i++) {
-				clients[i]->outputBuf.append(&line[4], size - 4);
-				clients[i]->outputBuf.append("\n", 1);
-			}
+			Buffer pkt;
+			pkt.writeStr(&line[4]);
+			for (int i = 0; i < clientCount; i++)
+				clients[i]->sendPacket(Packet::B2C_STATUS, pkt);
+
 		} else if (strcmp(line, "quit") == 0) {
 			quitFlag = true;
 		} else if (strncmp(line, "resolve ", 8) == 0) {
@@ -315,7 +316,12 @@ void Client::handleLine(char *line, int size) {
 			servers[serverCount]->ircPort = 1191;
 			servers[serverCount]->ircUseTls = (line[0] == 's');
 			serverCount++;
-			outputBuf.append("Your wish is my command!\n", 25);
+
+			Buffer pkt;
+			pkt.writeStr("Your wish is my command!");
+			for (int i = 0; i < clientCount; i++)
+				clients[i]->sendPacket(Packet::B2C_STATUS, pkt);
+
 		} else if (strncmp(line, "connsrv", 7) == 0) {
 			int sid = line[7] - '0';
 			servers[sid]->beginConnect();
@@ -347,34 +353,61 @@ void Client::handlePacket(Packet::Type type, char *data, int size) {
 			if (!pkt.readRemains(SESSION_KEY_SIZE))
 				error = 2;
 
-			uint8_t reqKey[SESSION_KEY_SIZE];
-			pkt.read((char *)reqKey, SESSION_KEY_SIZE);
+			// Authentication goes here at some point, too
 
-			if (!isNullSessionKey(reqKey)) {
-				Client *other = findClientWithKey(reqKey);
-				if (other && other->authState == AS_AUTHED) {
-					// Yep, we can go!
-					other->resumeSession(this, lastReceivedByClient);
-					return;
+
+			if (error != 0) {
+				// Send an error...
+				Buffer pkt;
+				pkt.writeU32(error);
+				sendPacket(Packet::B2C_OOB_LOGIN_FAILED, pkt, /*allowUnauthed=*/true);
+
+				// Would close() now but this means the login failed packet never gets sent
+				// need to figure out a fix for this. TODO FIXME etc etc.
+
+			} else {
+				// or log us in!
+				uint8_t reqKey[SESSION_KEY_SIZE];
+				pkt.read((char *)reqKey, SESSION_KEY_SIZE);
+
+				printf("[fd=%d] Client authenticating\n", sock);
+
+				if (!isNullSessionKey(reqKey)) {
+					printf("[fd=%d] Trying to resume session...\n", sock);
+
+					Client *other = findClientWithKey(reqKey);
+					printf("[fd=%d] Got client %p\n", sock, other);
+
+					if (other && other->authState == AS_AUTHED) {
+						// Yep, we can go!
+						other->resumeSession(this, lastReceivedByClient);
+						return;
+					}
 				}
-			}
 
-			// If we got here, it means we couldn't resume the session.
-			// Start over.
-			generateSessionKey();
-			authState = AS_AUTHED;
+				// If we got here, it means we couldn't resume the session.
+				// Start over.
+				printf("[fd=%d] Creating new session\n", sock);
 
-			Buffer pkt;
-			pkt.append((char *)sessionKey, SESSION_KEY_SIZE);
-			sendPacket(Packet::B2C_OOB_LOGIN_SUCCESS, pkt);
+				generateSessionKey();
+				authState = AS_AUTHED;
+
+				Buffer pkt;
+				pkt.append((char *)sessionKey, SESSION_KEY_SIZE);
+				sendPacket(Packet::B2C_OOB_LOGIN_SUCCESS, pkt);
+			}
 
 		} else {
 			printf("[fd=%d] Unrecognised packet in AS_LOGIN_WAIT authstate: type %d, size %d\n",
 					sock, type, size);
 		}
 	} else if (authState == AS_AUTHED) {
-		//if (type == Packet::) {
-		/*} else */{
+		if (type == Packet::C2B_COMMAND) {
+			char cmd[2048];
+			pkt.readStr(cmd, sizeof(cmd));
+			handleCommand(cmd, strlen(cmd));
+
+		} else {
 			printf("[fd=%d] Unrecognised packet in AS_AUTHED authstate: type %d, size %d\n",
 					sock, type, size);
 		}
diff --git a/core.h b/core.h
index a01904d..93a1e19 100644
--- a/core.h
+++ b/core.h
@@ -48,6 +48,9 @@ struct Packet {
 	enum Type {
 		T_OUT_OF_BAND_FLAG = 0x8000,
 
+		C2B_COMMAND = 1,
+		B2C_STATUS = 1,
+
 		C2B_OOB_LOGIN = 0x8001,
 
 		B2C_OOB_LOGIN_SUCCESS = 0x8001,
@@ -85,7 +88,7 @@ private:
 	int readBufPosition;
 	void processReadBuffer();
 	void handlePacket(Packet::Type type, char *data, int size);
-	void handleLine(char *line, int size);
+	void handleCommand(char *line, int size);
 
 	void generateSessionKey();
 	void resumeSession(Client *other, int lastReceivedByClient);
diff --git a/python_client.py b/python_client.py
index 21f4c55..8237988 100644
--- a/python_client.py
+++ b/python_client.py
@@ -1,10 +1,9 @@
 import socket, ssl, threading, struct
 
-basesock = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0)
-basesock.connect(('localhost', 5454))
-#sock = ssl.wrap_socket(basesock)
-sock = basesock
-
+protocolVer = 1
+sock = None
+authed = False
+sessionKey = b'\0'*16
 nextID = 1
 lastReceivedPacketID = 0
 packetCache = []
@@ -39,7 +38,7 @@ def clearCachedPackets(pid):
 			packetCache.remove(packet)
 
 def reader():
-	global lastReceivedPacketID
+	global lastReceivedPacketID, authed, sessionKey
 	readbuf = b''
 
 	print('(Connected)')
@@ -53,48 +52,89 @@ def reader():
 
 		pos = 0
 		bufsize = len(readbuf)
+		print('[bufsize: %d]' % bufsize)
 		while True:
 			if (pos + 8) > bufsize:
 				break
 
 			type, reserved, size = struct.unpack_from('<HHI', readbuf, pos)
-			pos += 8
 
 			extHeaderSize = 8 if ((type & 0x8000) == 0) else 0
-			if (pos + extHeaderSize + size) > bufsize:
+			if (pos + 8 + extHeaderSize + size) > bufsize:
 				break
 
-			if ((type & 0x8000) == 0):
-				pid, lastReceivedByServer = struct.unpack_from('<II', readbuf, pos)
-				pos += 8
+			pos += 8
+			with packetLock:
+				if ((type & 0x8000) == 0):
+					pid, lastReceivedByServer = struct.unpack_from('<II', readbuf, pos)
+					pos += 8
 
-				with packetLock:
 					lastReceivedPacketID = pid
 					clearCachedPackets(lastReceivedByServer)
 
-			packetdata = data[pos:pos+size]
-			print('0x%x : %d bytes : %s' % (type, size, packetdata))
+				packetdata = data[pos:pos+size]
+				print('0x%x : %d bytes : %s' % (type, size, packetdata))
+
+				if type == 0x8001:
+					sessionKey = packetdata
+					authed = True
+				elif type == 0x8003:
+					authed = True
+					pid = struct.unpack('<I', packetdata)
+					clearCachedPackets(pid)
+					try:
+						for packet in packetCache:
+							packet.sendOverWire()
+					except:
+						pass
+				elif type == 1:
+					print(packetdata.decode('utf-8'))
+
 			pos += size
 
-def writePacket(type, data):
-	with packetLock:
-		packet = Packet(type, data)
-		if (type & 0x8000) != 0:
-			packetCache.append(packet)
-		packet.sendOverWire()
+		print('[processed %d bytes]' % pos)
+		readbuf = readbuf[pos:]
+
+def writePacket(type, data, allowUnauthed=False):
+	packet = Packet(type, data)
+	if (type & 0x8000) == 0:
+		packetCache.append(packet)
+	try:
+		if authed or allowUnauthed:
+			packet.sendOverWire()
+	except:
+		pass
 
 
-thd = threading.Thread(None, reader)
-thd.start()
 
 while True:
 	bit = input()
 	bits = bit.split(' ', 1)
 	cmd = bits[0]
 
-	if cmd == 'login':
-		writePacket(0x8001, struct.pack('<II 16s', 0, 0, b'\0'*16))
-	elif cmd == 'cmd':
-		data = bits[1].encode('utf-8')
-		writePacket(1, struct.pack('<I', len(data)) + data)
+	with packetLock:
+		print('{')
+		if cmd == 'connect':
+			try:
+				basesock = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0)
+				basesock.connect(('localhost', 5454))
+				#sock = ssl.wrap_socket(basesock)
+				sock = basesock
+				thd = threading.Thread(None, reader)
+				thd.start()
+			except Exception as e:
+				print(e)
+		elif cmd == 'disconnect':
+			sock.shutdown(socket.SHUT_RDWR)
+			sock.close()
+			sock = None
+			authed = False
+		elif cmd == 'login':
+			writePacket(0x8001, struct.pack('<II 16s', protocolVer, lastReceivedPacketID, sessionKey), True)
+		elif cmd == 'cmd':
+			data = bits[1].encode('utf-8')
+			writePacket(1, struct.pack('<I', len(data)) + data)
+		elif cmd == 'quit':
+			break
+		print('}')