import math, socket, ssl, select, time, threading, random, string from datetime import datetime EOC = "\0" def to_numeric(identifier): return ((ord(identifier[:1]) - 1) * 255) + (ord(identifier[1:]) - 1) 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): aborted = False def __init__(self, channel): self.channel = channel threading.Thread.__init__(self) self.daemon = True def run(self): while self.aborted is False: 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) def abort(self): self.aborted = True print "Aborted thread." 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 def end(self): for key, chan in self.channel_map.iteritems(): print "Destructing handler for channel %d" % chan.numeric chan.handler.destruct() class Channel: numeric = 0 binary = False handler = None client = None def __init__(self, client, handler, binary=False): self.handler = handler self.binary = binary self.client = client def process_chunk(self, chunk): self.handler.process(chunk) def send(self, data): 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 def destruct(self): 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) def destruct(self): self.pingthread.abort() 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) if len(readable) > 0: 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) cur_client = self.client_map[sock.fileno()] if data: cur_client.process_data(data) else: cur_client.end() 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 else: time.sleep(0.010)