socket class

"""
This module handles the classes reponsible for creating a TCP socket server.
It should be resuable for any telnet-like application, and should not contain
any game logic or formatting code.

TODO: register event handlers so that our engine can hook in to server events
TODO: needs asynchronicity 
"""

import socket
import select

class Connection(object):
    
    def on_connect(self): pass
    def on_close(self): pass
    def on_write(self, data): return data
    def on_read(self): pass
    def on_server_stop(self): pass
    
    def fileno(self):
        return self.socket.fileno()
        
    def __init__(self, socket, address):
        self.socket = socket
        self.address = address
        self.send_buffer = ""
        self.recv_buffer = ""
        self.challenge = ""
        self.on_connect()
        
    def write(self, data):
        data = self.on_write(data)
        self.send_buffer += data
        
    def read():
        data = self.recv_buffer
        self.recv_buffer = ''
        self.on_read(data)
        return data
                  
class Server(object):
    
    def on_start(self): pass
    def on_update(self): pass
    def on_stop(self): pass
    def on_read(self, connection, data): pass
    def on_write(self, connection, data): return data
    def on_new_connection(self, connection): pass
    def on_dropped_connection(self, connection): pass
            
    def __init__(self, address, port):        
        self.incoming = []
        self.connections = []
        self.dropped_connections = []
        self.control = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        
        self.control.setblocking(0)
        self.control.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.control.setsockopt(socket.SOL_SOCKET, socket.TCP_NODELAY, 0)
        self.control.bind((address, port))
        self.control.listen(5)
                                        
    def poll(self):
        """
        This is the main communication function of the server. Meant to be called at a regular schedule,
        this function will:
        A. check for new connections, and accept them if they exist
        B. check for incoming data from each active connection, and write this data to a buffer
        C. check if any data is scheduled to be sent to each connection, and send it if true
        """
                
        socket_list = [self.control]
                        
        for c in self.connections:
            # Add each connection to the socket list
            # TODO: check if these sockets should be read from or written to
            socket_list.append(c)
            
        # non-blocking select
        in_set, out_set, exc_set = select.select(socket_list, socket_list, [], 0)
                
        if self.control in in_set:
            # new connection
            (socket, address) = self.control.accept()
            c = Connection(socket, address)
            self.connections.append(c)
            self.on_new_connection(c)
            in_set.remove(self.control)
        
        for c in in_set:
            c.recv_buffer = ''
            c.recv_buffer = c.socket.recv(1024)
            
            if len(c.recv_buffer) > 0:
                self.on_read(c, c.recv_buffer)
            else:
                # closed connection
                self.on_dropped_connection(c)
                c.socket.close()
                self.connections.remove(c)
                continue
                                                        
        for c in out_set:
            # If there is data in the send buffer send it
            if len(c.send_buffer) > 0:
                data = self.on_write(c, c.send_buffer)
                nbytes = c.socket.send(data)                
                c.send_buffer = ''
                            
        return True
        

    def close_sockets(self):
        self.control.close()
                    
    def send_all(self, data):
        for c in self.connections:
            c.send_buffer += data
            
    def send_all_except(self, connection, data):
        for c in self.connections:
            if c != connection: c.send_buffer += data