SSLClient abstraction and basic PingHandler

daemon/dist/core.py

@@ -9,87 +9,11 @@ key_path = '/home/sven/ssl/private'
 # Configuration ends here
 ################################
 
-import socket, ssl, select
 from shared.core import *
 
-def remove_from_list(ls, val):
+version = 1.0
-	return [value for value in ls if value is not val]
-
-client_list = []
-client_map = {}
-select_inputs = []
-select_outputs = []
-
-class Client:
-	buff = ""
-	channel_map = {}
-	
-	def __init__(self, connstream):
-		self.stream = connstream
-		self.channel_map[0] = Channel(self, EchoHandler())
-
-	def process_data(self, data):
-		self.buff += data
-		stack = self.buff.split(EOC)
-		self.buff = stack.pop()
-		
-		for chunk in stack:
-			self.process_chunk(chunk)
-		
-	def process_chunk(self, chunk):
-		if len(chunk) > 2:
-			channel_identifier = chunk[:2]
-			data = chunk[2:]
-			
-			channel_numeric = to_numeric(channel_identifier)
-			
-			if channel_numeric in self.channel_map:
-				self.channel_map[channel_numeric].process_chunk(data)
-			else:
-				print "WARNING: Received data on non-existent channel %d" % channel_numeric
-				
-
-
-bindsocket = socket.socket()
-bindsocket.bind(('0.0.0.0', 9151))
-bindsocket.listen(5)
-
-select_inputs = [ bindsocket ]
-
-while select_inputs:
-	readable, writable, error = select.select(select_inputs, select_outputs, select_inputs)
-	
-	for sock in readable:
-		try:
-			if sock is bindsocket:
-				newsocket, fromaddr = bindsocket.accept()
-				connstream = ssl.wrap_socket(newsocket,
-						server_side=True,
-						certfile=cert_path,
-						keyfile=key_path,
-						ssl_version=ssl.PROTOCOL_TLSv1)
-				
-				new_client = Client(connstream)
-				
-				select_inputs.append(connstream)
-				client_map[connstream.fileno()] = new_client
-				client_list.append(new_client)
-			else:
-				data = sock.recv(1024)
-				
-				if data:
-					cur_client = client_map[sock.fileno()]
-					cur_client.process_data(data)
-				else:
-					select_inputs = remove_from_list(select_inputs, sock)
-					print "NOTICE: Client disconnected"
-		except ssl.SSLError, err:
-			if err.args[0] == ssl.SSL_ERROR_WANT_READ:
-				select.select([sock], [], [])
-			elif err.args[0] == ssl.SSL_ERROR_WANT_WRITE:
-				select.select([], [sock], [])
-			else:
-				raise
 
+daemon = SSLDaemon()
+daemon.start('0.0.0.0', 9151, cert_path, key_path)
 
 print "Server socket closed, exiting..."

daemon/node/core.py

@@ -11,12 +11,33 @@ allowed_certs = '/home/sven/ssl/allowed'
 import socket, ssl, time, math
 from shared.core import *
 
-s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+select_inputs = []
-ssl_sock = ssl.wrap_socket(s, cert_reqs=ssl.CERT_REQUIRED, ca_certs=allowed_certs)
+select_outputs = []
-ssl_sock.connect(('localhost', 9151))
+client_map = {}
+
+client = SSLClient("localhost", 9151, allowed_certs)
+select_inputs.append(client.ssl_sock)
+select_outputs.append(client.ssl_sock)
+client_map[client.ssl_sock.fileno()] = client.client
+
+while True:
+	readable, writable, error = select.select(select_inputs, select_outputs, select_inputs)
+	
+	for sock in readable:
+		try:
+			data = sock.recv(1024)
+			
+			if data:
+				cur_client = client_map[sock.fileno()]
+				cur_client.process_data(data)
+			else:
+				select_inputs = remove_from_list(select_inputs, sock)
+				print "NOTICE: Client disconnected"
+		except ssl.SSLError, err:
+			if err.args[0] == ssl.SSL_ERROR_WANT_READ:
+				select.select([sock], [], [])
+			elif err.args[0] == ssl.SSL_ERROR_WANT_WRITE:
+				select.select([], [sock], [])
+			else:
+				raise
 
-ssl_sock.write(to_identifier(0) + 'test data' + EOC + to_identifier(4190) + 'SAMPLEDATA' + EOC)
-
-print ssl_sock.read()[2:]
-
-ssl_sock.close()

daemon/shared/core.py

@@ -1,4 +1,5 @@
-import math
+import math, socket, ssl, select, time, threading, random, string
+from datetime import datetime
 
 EOC = "\0"
 
@@ -8,6 +9,57 @@ def to_numeric(identifier):
 def to_identifier(numeric):
 	return chr(int(math.floor(numeric / 255) + 1)) + chr((numeric % 255) + 1)
 	
+def remove_from_list(ls, val):
+	return [value for value in ls if value is not val]
+
+class PingThread(threading.Thread):
+	def __init__(self, channel):
+		self.channel = channel
+		threading.Thread.__init__(self)
+		self.daemon = True
+		
+	def run(self):
+		while True:
+			self.pingkey = ''.join(random.choice(string.letters + string.digits) for i in xrange(5))
+			self.pingtime = datetime.time(datetime.now())
+			self.channel.send("PING %s" % self.pingkey)
+			time.sleep(10)
+
+class Client:
+	buff = ""
+	channel_map = {}
+	handshake_status = 0
+	
+	version = 1.0
+	
+	def __init__(self, connstream):
+		self.stream = connstream
+		self.channel_map[0] = Channel(self, ControlHandler(self))
+		
+	def start_handshake(self):
+		self.handshake_status = 1
+		self.channel_map[0].send("OHAI")
+
+	def process_data(self, data):
+		self.buff += data
+		stack = self.buff.split(EOC)
+		self.buff = stack.pop()
+		
+		for chunk in stack:
+			self.process_chunk(chunk)
+		
+	def process_chunk(self, chunk):
+		if len(chunk) > 2:
+			channel_identifier = chunk[:2]
+			data = chunk[2:]
+			
+			channel_numeric = to_numeric(channel_identifier)
+			
+			if channel_numeric in self.channel_map:
+				self.channel_map[channel_numeric].process_chunk(data)
+			else:
+				print "WARNING: Received data on non-existent channel %d" % channel_numeric
+
 class Channel:
 	numeric = 0
 	binary = False
@@ -26,9 +78,134 @@ class Channel:
 		self.client.stream.send(to_identifier(self.numeric) + data + EOC)
 		
 class Handler:
+	client = None
+	channel = None
+	
+	def __init__(self, client):
+		self.client = client
+	
 	def process(self, chunk):
 		pass
 		
+class ControlHandler(Handler):
+	pingthread = None
+	
+	def process(self, chunk):
+		target = self.client.channel_map[0]
+		
+		if chunk == "OHAI":
+			self.client.handshake_status = 1
+			target.send("HITHAR")
+		elif chunk == "HITHAR":
+			self.client.handshake_status = 2
+			target.send("VERSION %d" % self.client.version)
+		elif chunk[:7] == "VERSION":
+			if self.client.handshake_status == 1:
+				# Version received, return own version
+				self.client.handshake_status = 3
+				target.send("VERSION %d" % self.client.version)
+			elif self.client.handshake_status == 2:
+				# Version received and already sent, version exchange complete
+				self.client.handshake_status = 4
+				target.send("KGO")
+				print "Handshake complete!"
+				self.pingthread = PingThread(target)
+				self.pingthread.start()
+			else:
+				raise Exception("VERSION received outside version exchange.")
+		elif chunk == "KGO":
+			if self.client.handshake_status == 3:
+				# Handshake complete
+				self.client.handshake_status = 4
+				print "Handshake complete!"
+				self.pingthread = PingThread(target)
+				self.pingthread.start()
+			else:
+				raise Exception("KGO received before handshake finalization.")
+		elif chunk[:4] == "PING":
+			command, pingkey = chunk.split(' ')
+			target.send("PONG %s" % pingkey)
+		elif chunk[:4] == "PONG":
+			command, pingkey = chunk.split(' ')
+			microseconds = (self.pingthread.pingtime.microsecond * 1.0) / 1000000
+			seconds_start = self.pingthread.pingtime.second + microseconds
+			current_time = datetime.time(datetime.now())
+			microseconds = (current_time.microsecond * 1.0) / 1000000
+			seconds_end = current_time.second + microseconds
+			
+			print "Latency: %f seconds" % (seconds_end - seconds_start)
+		
 class EchoHandler(Handler):
 	def process(self, chunk):
 		print "Received %s" % chunk
+
+class SSLClient:
+	hostname = ""
+	port = 0
+	client = None
+	controlchannel = None
+	controlhandler = None
+	ssl_sock = None
+	
+	def __init__(self, hostname, port, allowed_certs):
+		sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+		self.ssl_sock = ssl.wrap_socket(sock, cert_reqs=ssl.CERT_REQUIRED, ca_certs=allowed_certs)
+		self.ssl_sock.connect((hostname, port))
+		
+		self.client = Client(self.ssl_sock)
+		
+		self.client.start_handshake()
+		
+
+class SSLDaemon:
+	client_list = []
+	client_map = {}
+	select_inputs = []
+	select_outputs = []
+	
+	def __init__(self):
+		pass
+		
+	def start(self, interface, port, cert_path, key_path):
+		bindsocket = socket.socket()
+		bindsocket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+		bindsocket.bind((interface, port))
+		bindsocket.listen(5)
+
+		self.select_inputs = [ bindsocket ]
+
+		while self.select_inputs:
+			readable, writable, error = select.select(self.select_inputs, self.select_outputs, self.select_inputs)
+			
+			for sock in readable:
+				try:
+					if sock is bindsocket:
+						newsocket, fromaddr = bindsocket.accept()
+						connstream = ssl.wrap_socket(newsocket,
+								server_side=True,
+								certfile=cert_path,
+								keyfile=key_path,
+								ssl_version=ssl.PROTOCOL_TLSv1)
+						
+						new_client = Client(connstream)
+						
+						self.select_inputs.append(connstream)
+						self.select_outputs.append(connstream)
+						self.client_map[connstream.fileno()] = new_client
+						self.client_list.append(new_client)
+					else:
+						data = sock.recv(1024)
+						
+						if data:
+							cur_client = self.client_map[sock.fileno()]
+							cur_client.process_data(data)
+						else:
+							self.select_inputs = remove_from_list(self.select_inputs, sock)
+							print "NOTICE: Client disconnected"
+				except ssl.SSLError, err:
+					if err.args[0] == ssl.SSL_ERROR_WANT_READ:
+						select.select([sock], [], [])
+					elif err.args[0] == ssl.SSL_ERROR_WANT_WRITE:
+						select.select([], [sock], [])
+					else:
+						raise