# Stephen Timothy Cooney
# December 2009

# PeerToPeer

#Std libs
from socket import *
from threading import *
from struct import *
from time import time
from uuid import uuid4

#Custom libs
from Messages import *
from Host import *
from Session import *
from UDPPunchthrough import *

PEER_PUNCH_TIMEOUT = 30 # 30 seconds to punch-through before forceful quit
PEER_PING_TIMEOUT = 30 # a ping not responded to for more than 30 seconds is dropped
PEER_PING_DELAY = 5.0 # ping once every 5.0 seconds

# A peerconnection represents the information
# a peer needs to know about all its neighbors
class PeerConnection:
    def __init__(self):
        self.address = ('127.0.0.1', -1)
        self.ping_timeout = time() + PEER_PING_TIMEOUT
        self.ping_next = time() # ping now!

# Peers are advanced sessions that maintain a tcp
# connection to the host/server/tracker and
# a solitary udp connection between all of
# its peer bretheren. Peer inherits most
# of the standard Session interface.
# Standard Usage:
# local_port = 2009
# host_address = ('555.555.555.555', 6900)
# peer = Peer(local_port, host_address)
# peer.start()
# -- peer will seek to connect to other peers --
# -- peer will run, derivations can define behaviour --
# peer.close()
# -- dead peer
class Peer(Session):

    def __init__(self, localPort, hostAddress):

        # remote address of the p2p host that's
        # facilitation this peer's connection.
        self._hostAddress = hostAddress

        # create and contain a specialized session
        # that tracks messages between the peer and
        # the host/server/tracker.
        self._session = None # nothing until initiate

        # create and hook up the main socket
        # that this peer will operate out of.
        udpsock = socket(AF_INET, SOCK_DGRAM)
        udpsock.bind((gethostname(), localPort))

        # the uuid is a unique identifier assigned
        # by the server. It allows peers to connect
        # and identify other peers by more than just
        # an address. This allows us to communicate
        # with other peers even if they are behind
        # the same NAT.
        self.__uuid = ''

        # a list of active punchers tupled with timeout times
        # (Puncher(), timeout_time)
        self._punchers = []

        # peers is a collection of PeerConnection
        # objects that this peer is connected to.
        self._peers = []

        # init the base session
        Session.__init__(self, udpsock, hostAddress)
        # self._sock = udpsock after this point
        # self._addr = serverAddress after this point

    def _initiate(self):
        Session._initiate(self)
        sock = socket(AF_INET, SOCK_STREAM)
        sock.connect(self._hostAddress)
        self._session = LazySession(sock)
        self._session.peer = self
        self._session.start()

    def _update(self):
        Session._update(self)

        # lock the lazysession and interpret
        # its messages here
        self._session.lock()
        for message in self._session.messageQueue:
            self._handlehostmessage(message[0], message[1], message[2])
        self._session.messageQueue = [] # clear the message queue
        self._session.release()

        # look to see if the punchers are either timed-out or successful
        dead_punchers = []
        for puncher in self._punchers:
            if puncher[0].IsPunched():
                print("Punchthrough successful to " + str(puncher[0].ResolvedAddress()))
                # create a new peerconnection
                peercon = PeerConnection()
                peercon.address = puncher[0].ResolvedAddress()
                peercon.ping_timeout = time() + PEER_PING_TIMEOUT
                peercon.ping_next = time() + PEER_PING_DELAY # ping now!
                self._peers.append(peercon) # add it to the list
                # kill the puncher
                puncher[0].close()
                dead_punchers.append(puncher)
            elif puncher[1] < time():
                print("Punchthrough has timed out.")
                # puncher has taken too long, quitting
                puncher[0].close()
                dead_punchers.append(puncher)
        # remove/erase dead punchers
        for puncher in dead_punchers:
            self._punchers.remove(puncher)

        # routine operations on our peers
        dead_peers = []
        for peer in self._peers:
            if peer.ping_timeout < time():
                # haven't received a pong for our ping in too long
                # of a time, remove the peer
                print("Timeout to peer: " + str(peer.address))
                dead_peers.append(peer)
            elif peer.ping_next < time():
                # time to send another ping to our peer to keepalive
                # and verify a connection exists.
                self._sendmessage(MESSAGE_PING, NODATA, peer.address)
                peer.ping_next = time() + PEER_PING_DELAY
        # clear references to our peer friends
        # whom are not active.
        for dead in dead_peers:
            self._peers.remove(dead)

    def _shutdown(self):
        Session._shutdown(self)
        if self._session is not None:
            self._session.close()

    def _findpeerconnection(self, address):
        for peer in self._peers:
            if peer.address == address:
                return peer # found the peerconnection, return it
        return None # didn't find it, return None
                

    # messages that are coming through the tcp socket connected
    # to the host are interpreted here.
    def _handlehostmessage(self, ident, data, address):

        print(getidentname(ident), data, address, "HOST")

        if ident is MESSAGE_ASSIGNUUID:
            # the server has given us a unique id.
            self._uuid = data[0]
            # Request peers from the Host
            self._session.sendmessage(MESSAGE_REQPEERS, NODATA)
        elif ident is MESSAGE_CONPEER:
            # create a puncher and append it to the
            # list of active punchers
            puncher = Puncher(self._sock, self._addr)
            if self._uuid is not None: puncher.AssignUUID(self._uuid)
            puncher.TargetIP(data[0])
            puncher.start()
            self._punchers.append((puncher, time()+PEER_PUNCH_TIMEOUT))
        elif ident is MESSAGE_CONPEERUUID:
            # create a puncher and append it to the
            # list of active punchers
            puncher = Puncher(self._sock, self._addr)
            if self._uuid is not None: puncher.AssignUUID(self._uuid)
            puncher.TargetUUID(data[0])
            puncher.start()
            self._punchers.append((puncher, time()+PEER_PUNCH_TIMEOUT))
            pass

    # messages that are coming through the udp server connected
    # to the other peers come in here.
    def _handlemessage(self, ident, data, address):

        # Let our parent handle messages first.
        # The default session automatically responds
        # to pings.
        Session._handlemessage(self, ident, data, address)

        if ident is MESSAGE_PING:
            # The parent Session class responds to pings
            pass
        
        elif ident is MESSAGE_PONG:
            # a peer is replying to our pong!
            # make sure to mark our reference
            # to this peer so it doesn't timeout
            peerconnection = self._findpeerconnection(address)
            if peerconnection is not None:
                # push the timeout time back
                peerconnection.ping_timeout = time() + PEER_PING_TIMEOUT
            else:
                # received a pong from  an unknown peer
                print("Received pong from an unknown peer! Ignoring.")


# A P2PHost is a derivation of a Host that adds
# a udp punchthrough facilitator to enable the
# clients to do direct peer to peer style connection
# since most peers/clients are behind NATs.
# The p2p host is also responsible for reporting
# to the peer/clients which other peers they can
# connect to. In this simple instantiation, all
# peer/clients attempt to connect to one another.
class P2PHost(Host):

    def __init__(self, listenPort=HOST_PORT, sessionClass=LazySession):
        Host.__init__(self,listenPort, sessionClass)

        # The facilitator is responsible for allowing
        # clients to directly connect to each other
        # over a UDP connection in a Peer2Peer fashion.
        # Since the listener is a tcp connection and
        # the facilitator is a UDP connection, we can
        # use the same port for the host. Excellent!
        self._facilitator = Facilitator(listenPort)

    def _initiate(self):
        Host._initiate(self)
        self._facilitator.start()

    def _onconnection(self, session):
        Host._onconnection(self, session)

        # sneak a uuid onto the session
        session.uuid = uuid4().hex
        # tell the session its uuid
        session.sendmessage(MESSAGE_ASSIGNUUID, session.uuid)

    def _handlesession(self, session):
        Host._handlesession(self, session)

        # interpret the session's messages
        # sessions are expected to be LazySessions
        # or derived from them
        session.lock()
        for message in session.messageQueue:
            self._handlemessage(message[0], message[1], session)
        # clear the queue before giving it back
        session.messageQueue = []
        session.release()

    def _handlemessage(self, ident, data, session):

        # print out information about the message received
        # from the session.
        print(getidentname(ident), data, session.getremoteaddress())

        if ident is MESSAGE_REQPEERS:
            # the client/peer is requesting a list
            # of peers that it can/should connect
            # to.
            # The implementation here connects all
            # active client/peers to all other
            # client/peers
            # the facilitator maintains a list of
            # verified UDP peers
            
            # grab the UDP list from the facilitator
            for othersession in self._sessions:
                if session is not othersession:
                    session.sendmessage(MESSAGE_CONPEERUUID, othersession.uuid)
                    othersession.sendmessage(MESSAGE_CONPEERUUID, session.uuid)
                

# Stephen Timothy Cooney
# December 2009

# UDPPunchthrough
# Facilitator & UDPPuncher

#Std libs
from socket import *
from threading import *
from struct import *
from time import time

#Custom libs
from Messages import *
from Session import *

FACILITATOR_PINGDELAYTIME = 10.0 # time between pings (in seconds)
FACILITATOR_TIMEOUTTIME = 30.0 # time before the facilitator drops the connection

# A punchpeer represents udp peers that exist under
# the facilitator object.
class PunchPeer:
    def __init__(self):
        self.address = ('0.0.0.0', -1)
        self.uuid = None
        self.ping_timeout = time() + FACILITATOR_TIMEOUTTIME
        self.ping_next = time() # ping now!
        # a punch peer may have any number of requests
        # to connect to other peers/clients.
        # in order for a pucnhthrough to initiate, there
        # must be a matching/pairing request from the other
        # end of the punch
        self.ipRequests = [] # match via ip (unknown port)
        self.uuidRequests = [] # match via uuid (unknown ip and port)

# The facilitator is responsible for maintaining an open
# udp connection so that clients behind NATs can have a
# route to connect to other clients behind NATs. Basically
# both clients make a request to connect to each other
# from the facilitator. When the facilitator receives a
# request to connect to each partner, the facilitator then
# looks up their public IPs and PORTS and sends them to the
# opposite peers. The peers then attempt to connect to each
# other directly.
# Standard usage:
# fac = Facilitator(6900)
# fac.start()
# # and your good... will automatically route connections
class Facilitator(Session):

    def __init__(self, facilitatorPort):
        
        # Create a UDP port for the purpose of facilitating
        # UDP connections between peers behind NATs.
        udpsock = socket(AF_INET, SOCK_DGRAM)
        udpsock.bind((gethostname(), facilitatorPort))
        
        # Fill in the parent blanks
        Session.__init__(self, udpsock, None)

        # The facilitator mantains a collection
        # of addrs of previously resolved ip/port combos.
        # maps to the last time ping was sent (for keepalive)
        # PunchPeer objects!!
        self._addrs = []

    def _update(self):
        Session._update(self)

        # loop through each of the stored addresses
        # and send a keepalive ping (since these ARE
        # UDP connections)
        dead_addrs = []
        for peer in self._addrs:

            # Has the client timed out?
            if peer.ping_timeout < time():
                # The peer/client has been away far too
                # long. Remove them from the list.
                print("Client/Peer timeout on " + str(peer.address))
                dead_addrs.append(peer)
                
            else:
                # Do routine pings to keepalive the udp hole
                if peer.ping_next < time():
                    # send a ping (expect a pong)
                    self._sendmessage(MESSAGE_PING, NODATA, peer.address)
                    # reset the timer for the address/peer
                    peer.ping_next = time() + FACILITATOR_PINGDELAYTIME

        # go through and clear out the dead addresses
        for addr in dead_addrs:
            self._addrs.remove(addr)

        # compile a list of matching requests
        # that can now begin a punchthrough.
        # matches is a list of tuples of two
        # addresses. (address, address)
        matches = self._compilematches()
        for match in matches:
            # POW, begin punchthrough
            self._sendmessage(MESSAGE_EXCPUNCH, match[0], match[1])
            self._sendmessage(MESSAGE_EXCPUNCH, match[1], match[0])
        

    # searches through all of the requests and
    # compiles a list of address pairs. If matches
    # are made, the requests are automatically removed
    # from the peers.
    def _compilematches(self):
        matches = []

        # for every peer,
        for i in range(0, len(self._addrs)):
            peerA = self._addrs[i]
            # look at the other peers,
            for j in range(i+1, len(self._addrs)):
                peerB = self._addrs[j]

                # IP MATCHING

                # dead requests are requests that
                # have found matches
                dead_arequests = []
                dead_brequests = []
                
                # compare this peers ip requests,
                for reqA in peerA.ipRequests:
                    # to their ip requests,
                    for reqB in peerB.ipRequests:
                        # and check for a match
                        if reqA == peerB.address[0] and reqB == peerA.address[0]:
                            matches.append((peerA.address, peerB.address))
                            dead_arequests.append(reqA)
                            dead_brequests.append(reqB)

                # erase matched requests
                for req in dead_arequests:
                    peerA.ipRequests.remove(req)
                for req in dead_brequests:
                    peerB.ipRequests.remove(req)

                # UUID MATCHING

                # reset the dead list for uuid requests
                dead_arequests = []
                dead_brequests = []

                # compare this peers uuid requests,
                for reqA in peerA.uuidRequests:
                    # to their uuid requests,
                    for reqB in peerB.uuidRequests:
                        # and check for a match
                        if reqA == peerB.uuid and reqB == peerA.uuid:
                            matches.append((peerA.address, peerB.address))
                            dead_arequests.append(reqA)
                            dead_brequests.append(reqB)

                # erase matched requests
                for req in dead_arequests:
                    peerA.uuidRequests.remove(req)
                for req in dead_brequests:
                    peerB.uuidRequests.remove(req)

        return matches

    def _findaddress(self, address):
        for addr in self._addrs:
            if addr.address == address:
                return addr
        return None

    def _addaddress(self, address):
        newaddress = PunchPeer()
        newaddress.address = address
        newaddress.ping_next = time() # pingback immediately (for server/client punchthrough)
        newaddress.ping_timeout = time() + FACILITATOR_TIMEOUTTIME
        self._addrs.append(newaddress)
        return newaddress
        
    def _handlemessage(self, ident, data, address):

        if ident is MESSAGE_REGISTERUUID:

            # The client can register a uuid to be identified
            # by. The UUID can be used for punchthrough
            # matching when two clients are behind the same
            # public address.
            addr = self._findaddress(address)
            if addr is None:
                addr = self._addaddress(address)
            # set the registered peer's uuid
            addr.uuid = data[0]
            
        elif ident is MESSAGE_PING:

            # Generally as a facilitator we should
            # only receive a ping as the first communication
            # in which a client does not immediately
            # want to connect but wants to be known.
            addr = self._findaddress(address)
            if addr is None:
                # unkown client, add them
                self._addaddress(address)
                # pong back
                self._sendmessage(MESSAGE_PONG, NODATA, address)
            
        elif ident is MESSAGE_PONG:

            # We received a pong!
            # set the incoming addresses ping to
            # FACILITATOR_PINGDELAYTIME seconds in
            # the future and mark the pong time
            # as the current time.
            addr = self._findaddress(address)
            if addr is not None:
                addr.ping_next = time() + FACILITATOR_PINGDELAYTIME
                addr.ping_timeout = time() + FACILITATOR_TIMEOUTTIME
                
        elif ident is MESSAGE_REQPUNCH or ident is MESSAGE_REQPUNCHUUID:

            # Received a request from a client to connect
            # via puncthrough to another client.
            print("Received request: " + str(address) + " to " + str(data[0]))

            # See if the address requesting exists.
            # If it doesnt, add a new punchpeer
            addr = self._findaddress(address)
            if addr is None:
                addr = self._addaddress(address)

            # queue a punch-request
            if ident is MESSAGE_REQPUNCH:
                addr.ipRequests.append(data[0])
            elif ident is MESSAGE_REQPUNCHUUID:
                addr.uuidRequests.append(data[0])
                
        else:

            # We didn't handle the message so let our
            # parent class handle it.
            Session._handlemessage(self, ident, data, address)
            

PUNCHER_PINGDELAY = 0.1 # seconds between pings
PUNCHER_MAXPINGS = 100 # number of pings before we give up!

# The puncher session encapsulates procedures for
# connecting to a remote session behind a NAT via
# a facilitator that is accessible at a public
# address. The puncher constructor requires:
# sock - UDP socket that is already constructed and bound
# destIP - public Ip address (minus port) that we want to connect to
# facilitatorAddress - (ip,port) tuple that maps the facilitators location
#
# Standard usage:
# udpsock = //some_sock//
# destIP = //some_address// ex: "75.94.122.9"
# facilitatorAddress = //facilitator// ("88.99.55.33", 6900)
# puncher = Puncher(udpsock, facilitatorAddress)
# puncher.TargetIP(destIP)
# puncher.start()
# timeout_time = time() + 60 # one minute to try
# while timeout_time < time():
#   if puncher.IsPunched():
#       # awesome UDP socket is punched!
#       break;
# puncher.close()
# puncher.join()
# # udpsock is now connected and usable!
class Puncher(Session):

    def __init__(self, sock, facilitatorAddress):
        Session.__init__(self, sock)

        # destIP is a link to the IP address that we
        # want to connect to. We don't know the port
        # of this address however. (the point of punching)
        self._destIP = ''
        self._destPort = -1 # unknown!

        # The uuid is intended to be a unique number that
        # clients can identify each other with ignorant of
        # public address
        self._uuid = None

        # the puncher could seek connection to a peer
        # id instead of a peer IP. If no UUID is specified
        # then the puncher will seek an IP.
        self._destUUID =None
        
        # represents the address of the facilitator
        # that enables the UDP punchthrough.
        self._facilitator = facilitatorAddress

        # while executing punch, we send out
        # pings repeatedly until we receive
        # a pong from the destination IP
        self._executingpunch = False
        # the time() that we execute the next ping
        self._nextping = 0
        # the number of pings that have been executed
        self._pingcounter = 0

        # marked to true when a punch has been acheived
        self._successful = False

        # need a ping and pong message to be pucnehd
        self._gotpong = False
        self._gotping = False

    def TargetIP(self, destinationIP):
        self._destIP = destinationIP

    def AssignUUID(self, myUUID):
        self._uuid = myUUID

    def TargetUUID(self, destinationUUID):
        self._destUUID = destinationUUID

    def IsPunching(self):
        return self._executingpunch

    def IsPunched(self):
        return self._successful

    def ResolvedAddress(self):
        return (self._destIP, self._destPort)

    def _initiate(self):
        Session._initiate(self)

        # register the custom uuid
        if self._uuid is not None:
            self._sendmessage(MESSAGE_REGISTERUUID, (self._uuid,), self._facilitator)

        # If the puncher was given a target uuid via Puncher.TargetUUID(uuid)
        # then try to connect via a uuid connection.
        if self._destUUID is not None:
            # start off by sending the facilitator server the uuid of the client you want to connect to
            self._sendmessage(MESSAGE_REQPUNCHUUID, (self._destUUID,), self._facilitator)
        else:
            # start off by sending the facilitator server where you want to connect
            self._sendmessage(MESSAGE_REQPUNCH, (self._destIP,), self._facilitator)

    def _update(self):
        Session._update(self)

        if self._executingpunch is True and self._successful is not True:
            # Ping at interval rates,
            # check to see if it's time to ping
            # again.
            if self._nextping <= time():
                self._sendmessage(MESSAGE_PING, NODATA, (self._destIP, self._destPort))
                # iterate
                self._nextping = time() + PUNCHER_PINGDELAY
                self._pingcounter += 1
            # Don't punch forever.
            # check to make sure that we haven't
            # tried too much
            if self._pingcounter >= PUNCHER_MAXPINGS:
                self._executingpunch = False
                print("Punchthrough has received no response. Quitting.")

    # overridding the Session._shutdown
    # The session implementation kills
    # the socket but we want the puncher
    # to exist temporarily as a means of
    # opening ports. We don't kill the
    # socket here
    def _shutdown(self):
        self._sock = None
        self._desireshutdown = True # probably redundant, but maybe not

    def _handlemessage(self, ident, data, address):

        if ident is MESSAGE_EXCPUNCH:
            #do punchthrough!
            if self._successful is True or self._executingpunch is True:
                # already punched? Ignore
                print("Received request to execute punch yet we already have, ignoring")
            else:
                self._destIP = data[0]
                self._destPort = data[1]
                self._nextping = time()
                self._executingpunch = True
                print("Executing punchthrough to " + str((data[0], data[1])) + "!")
        elif ident is MESSAGE_PONG:
            # a pong is a response to our ping.
            # A successful punch if this address
            # is from the destIP
            if address[0] == self._destIP:
                print("PONG! Punchthrough successful!")
                # respond to a pong with a punch acknowledgement message
                self._sendmessage(MESSAGE_PUNCHACK, NODATA, address)
                self.flushoutgoing()
                # success
                self._successful = True
                self._desireshutdown = True
            else:
                print("Pong, but from wrong address, still punching.")
        elif ident is MESSAGE_PUNCHACK:
            if address[0] == self._destIP:
                print("PUNCH ACK! Punchthrough is successful!")
                self.flushoutgoing()
                # success!
                self._successful = True
                self._desireshutdown = True
        else:
            # Puncher didn't handle the message so let
            # our parent class handle it.
            Session._handlemessage(self, ident, data, address)

from struct import *

# global message list is a collection
# of tuplets connecting message IDs to
# message formatting functions. The idea
# here is to abstract and make messages
# more portable.
defaultMessageFormats = []

NODATA = b''
PACK = True
UNPACK = False
# either packs or unpacks the data relative
# to an ident and returns it in either a bytes object
# or a tuple of sorts. Expects the reverse depending on
# whether or not you're packing.
def formatmessage(ident, data, packing, formatList=defaultMessageFormats):
    for formatter in formatList:
        if formatter[0] == ident:
            return formatter[1](data, packing)

    print("Unknown message type: " + str(data))

    #no handler, return nothing
    if packing:
        return b''
    else:
        return ()

# Return a string name that pairs with each message id. This
# is useful for debugging purposes.
def getidentname(ident, formatList=defaultMessageFormats):
    # Loop through the list of formatters and
    # return the name of the matching id.
    for formatter in formatList:
        if formatter[0] == ident:
            return formatter[2]

    # The id is unknown
    return "Unknown"
#________
# UTILITY

# strip the identifier out of a bytes object
# and return a pair: identifier, bytes_data
def stripident(rawdata):
    if len(rawdata) == 2:
        return  unpack("!H", rawdata)[0], b''
    elif len(rawdata) > 2:
        return unpack("!H", rawdata[0:2])[0], rawdata[2:]
    else:
        return 0

# clipstring is a utility function used for
# taking oversized strings (due to fixed length
# network  messages) and shrinking them down
# to their actual size
def clipstring(oversized):
    for i in range(len(oversized)):
        if oversized[i] is 0 or oversized[i] == '\x00':
            return oversized[0:i]
    return oversized

#_________
# MESSAGES

id_iterator = 0
def AllocID():
    global id_iterator
    out = id_iterator
    id_iterator += 1
    return out

# No message
MESSAGE_NONE = AllocID()

# PING ID
MESSAGE_PING = AllocID()
# PING FORMATTER
def fm_ping(data, packing):
    if packing:
        return b''
    else:
        return ()
# APPENDING PING
defaultMessageFormats.append((MESSAGE_PING, fm_ping, "Ping"))

# PONG ID (reply to ping)
MESSAGE_PONG = AllocID()
# PONG FORMATTER
def fm_pong(data, packing):
    if packing:
        return b''
    else:
        return ()
# APPENDING PONG
defaultMessageFormats.append((MESSAGE_PONG, fm_pong, "Pong"))

#__________________________
# UDP PUNCHTHROUGH MESSAGES

# PUNCH ACK
# An acknowledgement of a full punchthrough
# The flow is as such:
# Ping -> Pong -> Punch Ack
MESSAGE_PUNCHACK = AllocID()
# REGISTER PUNCH ACK FORMATTER
def fm_punchack(data, packing):
    if packing:
        return b''
    else:
        return ()
# APPENDING PUNCHACK
defaultMessageFormats.append((MESSAGE_PUNCHACK, fm_punchack, "Punch Acknowledge"))

# REGISTER UUID
# From the client to facilitator registering
# a uuid with this IP address. The UUID allows
# clients to recognize each other even if they
# are behind the same public address.
MESSAGE_REGISTERUUID = AllocID()
# REGISTER UUID FORMATTER
def fm_registeruuid(data, packing):
    if packing:
        # as a one length tuple it could accidentally be passed
        # in as a string. ("string",) is proper for one length tuples.
        if type(data) is str:
            data = (data,)
        # data tuple expected as (ip,)
        return pack("!32s", data[0])
    else:
        unpacked = unpack("!32s", data)
        cleaned = (clipstring(unpacked[0]).decode('ascii'),)
        return cleaned
# APPENDING UREGISTER UID FORMATTER
defaultMessageFormats.append((MESSAGE_REGISTERUUID, fm_registeruuid, "Register UUID"))

# REQUEST PUNCHTHROUGH
# From client to facilitator requesting a
# connection to a remote ip.
MESSAGE_REQPUNCH = AllocID()
# REQUEST PUNCHTHROUGH FORMATTER
def fm_reqpunch(data, packing):
    if packing:
        # as a one length tuple it could accidentally be passed
        # in as a string. ("string",) is proper for one length tuples.
        if type(data) is str: 
            data = (data,)
        # data tuple expected as (ip,)
        return pack("!16s", data[0])
    else:
        unpacked = unpack("!16s", data)
        cleaned = (clipstring(unpacked[0]).decode('ascii'),)
        return cleaned
# APPENDING REQUEST PUNCH
defaultMessageFormats.append((MESSAGE_REQPUNCH, fm_reqpunch, "Request IP Punchthrough"))

# REQUEST PUNCTHROUGH VIA UUID
# From the client to the facilitator requesting
# a connection to a remote peer via their
# uuid as registered on the facilitator.
MESSAGE_REQPUNCHUUID = AllocID()
# REQUEST PUNCHTHROUGH VIA UUID FORMATTER
def fm_reqpunchuuid(data, packing):
    if packing:
        # as a one length tuple it could accidentally be passed
        # in as a string. ("string",) is proper for one length tuples.
        if type(data) is str: 
            data = (data,)
        # data tuple expected as (ip,)
        return pack("!32s", data[0])
    else:
        unpacked = unpack("!32s", data)
        cleaned = (clipstring(unpacked[0]).decode('ascii'),)
        return cleaned
# APPENDING PUNCHTHROUGH VIA UUID
defaultMessageFormats.append((MESSAGE_REQPUNCHUUID, fm_reqpunchuuid, "Request UUID Punchthrough"))
        

# EXECUTE PUNCHTHROUGH
# From the server to the clients containing
# an IP and PORT to connect to. The clients
# are expected to execute the punch routine.
MESSAGE_EXCPUNCH = AllocID()
# EXECUTE PUNCH FORMATTER
def fm_excpunch(data, packing):
    if packing:
        # data tuple expected as (ip, port)
        return pack("!16sI", data[0], data[1])
    else:
        unpacked = unpack("!16sI", data)
        cleaned = (clipstring(unpacked[0]).decode('ascii'), unpacked[1])
        return cleaned
# APPENDING EXECUTE PUNCH
defaultMessageFormats.append((MESSAGE_EXCPUNCH, fm_excpunch, "Execute Punchthrough"))

#_____________
# P2P MESSAGES

# REQUEST PEERS
# From the peer/client to the host
# Request a list of peers to connect to.
# This is generic and leaves the host
# the power to decide who should connect
# to whom.
MESSAGE_REQPEERS = AllocID()
# REQPEER FORMATTER
def fm_reqpeers(data, packing):
    if packing:
        return b''
    else:
        return ()
# APPEND REQPEERS
defaultMessageFormats.append((MESSAGE_REQPEERS, fm_reqpeers, "Request Peers"))

# CONNECT PEER
# From the host to the peer/client
# Sends a peer ip to connect to.
# As a note, these should be sent
# to both end, or a punchthrough
# won't work. Sends the ip of
# a peer that is registered on the
# facilitator. The client still
# needs to interface with the facilitator
# to actually connect.
MESSAGE_CONPEER = AllocID()
# CONPEER FORMATTER
def fm_conpeer(data, packing):
    if packing:
        # as a one length tuple it could accidentally be passed
        # in as a string. ("string",) is proper for one length tuples.
        if type(data) is str:
            data = (data,)
        return pack("!16s", data[0])
    else:
        return (unpack("!16s", data)[0].decode('ascii'),)
# APPENDING CONPEER
defaultMessageFormats.append((MESSAGE_CONPEER, fm_conpeer, "Connect Peer IP"))

# CONNECT PEER UUID
# From the host to the peer/client
# Tells the peer the uuid of a peer
# that the client shoudl try to connect
# to that is registered on the facilitator.
MESSAGE_CONPEERUUID = AllocID()
# CONPEERUUID FORMATTER
def fm_conpeeruuid(data, packing):
    if packing:
        # as a one length tuple it could accidentally be passed
        # in as a string. ("string",) is proper for one length tuples.
        if type(data) is str:
            data = (data,)
        return pack("!32s", data[0])
    else:
        return (unpack("!32s", data)[0].decode('ascii'),)
# APPENDING CONPEERUUID
defaultMessageFormats.append((MESSAGE_CONPEERUUID, fm_conpeeruuid, "Connect Peer UUID"))

# ASSIGN UUID
# the host assigns each peer a uuid so they may be mapped
# even if they are on the same NAT showing the same public
# ip address. This message is also used when the peer registers
# itself on the facilitator.
MESSAGE_ASSIGNUUID = AllocID()
# ASSIGN UUID FORMATTER
def fm_assignuuid(data, packing):
    if packing:
        # as a one length tuple it could accidentally be passed
        # in as a string. ("string",) is proper for one length tuples.
        if type(data) is str: 
            data = (data,)
        return pack("!32s", data[0])
    else:
        return (unpack("!32s", data)[0].decode('ascii'),)
# APPENDING ASSIGN UUID
defaultMessageFormats.append((MESSAGE_ASSIGNUUID, fm_assignuuid, "Assign UUID"))

# Stephen Timothy Cooney
# December 2009

# Session
# LazySession

# Std libs
from socket import *
import select
from threading import *
from struct import *

# Custom libs
from Messages import *

# Helper funcs
# These functions are here to assist
# in the creation/testing of sessions
#
# Accept Session
# Start up a listen server on a port
# and listen for the incoming session.
# return a session with the new sock.
def AcceptSession(listenPort):
    lis = socket(AF_INET, SOCK_STREAM)
    lis.bind((gethostname(), listenPort))
    lis.listen(1)
    sock, addr = lis.accept()
    lis.close()
    sess = Session(sock)
    sess.start()
    return sess
#
# ConnectSession
# Directly connect to a remote listen
# server and then return the resultant
# Session with connection sock.
def ConnectSession(address):
    sock = socket(AF_INET, SOCK_STREAM)
    sock.connect(address)
    sess = Session(sock)
    sess.start()
    return sess
#
# UDPSession
# Starts up a UDP session that listens on
# localPort and connects to remoteAddress.
# remoteAddress is a (IP,Port) tuple
def UDPSession(localPort, remoteAddress):
    sock = socket(AF_INET, SOCK_DGRAM)
    sock.bind((gethostname(), localPort))
    sess = Session(sock, remoteAddress)
    sess.start()
    return sess

SESSION_MSG_SIZE = 4096 # max allowed message size

# A session encapsulates a connection.
# Sessions have the ability to pack and unpack
# messages. Messages/Data can be queued and
# are sent when appropriate.
# for derivation:
# _initiate(self) - executed on thread start before update loop
# _update(self) - executed in a loop
# _shutdown(self) - executed just before thread death
# _handlemessage(self) - called on receipt of message
# public funcs:
# close() - kill the session
# isdead() - is the session dead?
class Session(Thread):

    def __init__(self, socket, dstaddr=None): # must pass dstaddr for UDP
        Thread.__init__(self)
        # Retain a connection to the socket/address
        # combo that this session will monitor
        self._sock = socket
        # Outgoing is a collection of message data
        # that we want to send to the remote socket.
        self._outgoing = []
        # Will hopefull run the socket connection until
        # shutdown is set to True.
        self._desireshutdown = False
        
        # Retain the destination address
        if self._sock.type is SOCK_STREAM:
            self._addr = self._sock.getpeername()
        elif self._sock.type is SOCK_DGRAM:
            # For UDP connections we need to collect
            # a destination address. This is manual.
            self._addr = dstaddr

    def run(self):
        self._initiate()
        while self._desireshutdown is not True:
            self._update()
        self._shutdown()

    # kill the session
    def close(self):
        try:
            # Try to flush anything we have stored in the
            # outgoing queue.
            self.flushoutgoing()
        except:
            # Might have a bad socket (which may be the reason
            # we're closing) so we do nothing!
            pass
        self._desireshutdown = True # Mark the thread for destruction

    # check to see if the session is considered dead
    def isdead(self):
        return self._desireshutdown

    # To be overridden
    # Provides a place for the session to be initialized
    # before it is in it's main loop but after the thread
    # has been asked to execute.
    def _initiate(self):
        pass

    # To be overridden
    # Runs in a loop and is the root of consistent
    # session/socket functions.
    def _update(self):
        import sys
        try:
            # poll options
            rlist, wlist, xlist = select.select([self._sock],[self._sock],[self._sock])
            # what can we do?
            if xlist.count(self._sock) > 0: #ERRORS
                print("(Select) Unknown socket error! Shutting down.")
                self._desireshutdown = True
                return
            elif rlist.count(self._sock) > 0: #INCOMING DATA
                buffer, address = self._sock.recvfrom(SESSION_MSG_SIZE) # receive the data
                if address is None:
                    # most likely tcp (not udp)
                    address = self._sock.getsockname()
                if buffer == b'': # empty data means lost connection to socket
                    print("Session lost connection to the socket. Shutting down.")
                    self._desireshutdown = True # mark for death
                    return
                self._handlebuffer(buffer, address) # interpret and handle the packet buffer
            elif wlist.count(self._sock) > 0: #OUTGOING DATA
                self.flushoutgoing() # send them all!
        except:
            print("Unknown socket error! Shutting down.")
            print(sys.exc_info())
            self._desireshutdown = True

    # To be overridden
    # Last chance to cleanup any initialization before the
    # thread ultimately closes.
    def _shutdown(self):
        self._sock.close()
        self._desireshutdown = True # just in case

    # Handle the incoming buffer.
    # Identifies the message and formats the buffer.
    # Then, calls _handlemessage(ident,data) for
    # final interpretation
    def _handlebuffer(self, buffer, address):
        import sys
        # Interpret and handle the data from the buffer
        try:
            ident, data = stripident(buffer) # seperate the message identifier from the raw data
            data = formatmessage(ident, data, UNPACK) # interpret the data
        except:
            # Bad message, ignore
            print("Error interpreting buffer! Ignoring.")
            print(str(sys.exc_info()))
        try:
            self._handlemessage(ident, data, address) # handle the data
        except:
            print("Error handling message! Ignoring.")
            print(str(sys.exc_info()))
    
    # To be overridden
    # Plug for handling messages from the socket.
    # At this point ident has been broken into
    # an identifier(short) and data is split into
    # a tuple.
    def _handlemessage(self, ident, data, address):

        # print out information about the message
        print(getidentname(ident), data, address)
        
        if ident is MESSAGE_PING:
            # respond to ping with a pong
            self._sendmessage(MESSAGE_PONG, NODATA, address)
        elif ident is MESSAGE_PONG:
            # do nothing!
            # pong is a response to our ping.
            pass

    # Packs the message using the identifier
    # data paremeter. Data is expected to
    # be a tuple that maps to the format handler
    # specified by the identifier.
    def _sendmessage(self, ident, data, address):
        print("=> Sending", getidentname(ident), data, address)
        data = formatmessage(ident, data, PACK) # pack the message into a datastream
        data = pack("!H", ident) + data # merge the msg id with the message data
        self._outgoing.append((address,data)) # queue the datastream

    # Queues the data parameter. Assumes that the data
    # can be interpreted by the remote message handler.
    def _senddata(self, data, address):
        self._outgoing.append((address,data)) # queue the datastream outright
        
    def sendmessage(self, ident, data, address=None):
        if address:
            self._sendmessage(ident,data,address)
        else:
            self._sendmessage(ident,data,self._addr)

    def flushoutgoing(self):
        for data in self._outgoing:
            self._sock.sendto(data[1], data[0]) # send data
        self._outgoing = [] # clear the outgoing list

    def getlocaladdress(self):
        return self._sock.getsockname()
    
    def getremoteaddress(self):
        return self._sock.getpeername()


# Lazy session does just about the
# same things as any other session
# except technically these don't really
# do anything(contradictory, yes).
# Instead of processing messages,
# the lazy session just keeps a big
# list of them. It is the owners
# responsibility to process and clear
# the messages.
# Introduces the lock() and release()
# calls to safely access the messageQueue
class LazySession(Session):

    def __init__(self, sock, dstAddr=None):
        Session.__init__(self,sock,dstAddr)

        # Asynchronous thread, keep data from getting
        # poorly mingled.
        self._queuelock = Lock()
        # messageQueue is a public list of
        # messages. These should not be accessed
        # without locking the session-thread first
        self.messageQueue = []

    def lock(self):
        self._queuelock.acquire()

    def release(self):
        self._queuelock.release()

    def _handlemessage(self, ident, data, address):
        # appending the message as a tuple
        self.lock()
        self.messageQueue.append((ident,data,address))
        self.release()

from threading import *
from socket import *

#________________________________________
#////////////////////////////////////////
# LISTENER (TCP)

LISTEN_PORT = 2009 #default

# Simple class that just listens for incoming
# connections and stores them in a list. It's
# expected that an external user takes and
# clears this list when it is done using them
# See lock and relese for dealing the the list.
#
# The listener class isn't smart enough to drop
# connections on sockets that have been closed.
# Whomever owns the listener is responsible for
# pulling off the connections and handling them
# normally else the listener will get muddled.
#
# Typical Usage:
# l = Listener()
# l.start(some_port_number)
# ...
# do some stuff, in a regular update loop:
# ...
# l.lock()
# for connection in l.connections:
#   self.clients.append(connection) #clients is possible member array of a more complex networking class
# l.connections = [] #empty the list since we own the socket now
# l.release()
# ...
# do some more stuff, when ready for shutdown:
# ...
# l.close()
# # you can also do an optional l.join()
# # if you want to make sure the listener is
# # dead before moving on
class Listener(Thread):

    def __init__(self, listenPort=LISTEN_PORT):
        Thread.__init__(self)

        # A list of connections (clientsocket, address)
        # This is the only public member variable
        # use functions to interface with the rest.
        self.connections = []

        # what port are we listening through?
        self.__listenPort = listenPort

        # A custom lock for dealing with connections
        self.__threadlock = Lock() #wrapped around connections

        # Create a listening socket
        self.__listenSocket = socket(AF_INET, SOCK_STREAM);

        # Will loop until we are marked for shutdown!
        self.__shutdown = False

    def run(self):

        # start up the socket
        self.__listenSocket.bind((gethostname(), self.__listenPort))
        self.__listenSocket.listen(5)

        #loop until a shutdown request
        while not self.__shutdown:
            #accept and add new connections
            try:
                connection = self.__listenSocket.accept()
                self.__threadlock.acquire()
                self.connections.append(connection)
                self.__threadlock.release()
            except:
                self.__shutdown = True

        # Do shutdown procedures
        # disconnect any unhandled connections
        self.closeconnections()
        self.close()       

    # Call lock before dealing with the
    # self.connections array
    def lock(self):
        self.__threadlock.acquire()

    # Remember to free after dealing with
    # the self.connections.array
    def release(self):
        self.__threadlock.release()

    # Empty the list of connections
    # Does not close these connections,
    # this must be done elsewhere!
    def clear(self):
        self.__threadlock.acquire()
        self.connections = []
        self.__threadlock.release()

    # Shtudowns all the sockets that are
    # monitor and then removes them from
    # the list.
    def closeconnections(self):
        self.__threadlock.acquire()
        for con in self.connections:
            con[0].close()
        self.connections = []
        self.__threadlock.release()

    def close(self):
        self.__shutdown = True
        self.__listenSocket.close()
        

# Stephen Timothy Cooney
# December 2009

# Host

#Std libs
from socket import *
from threading import *
from struct import *

#Custom libs
from Session import *
from Listener import *
from UDPPunchthrough import Facilitator

HOST_PORT = 6900 # default

# The host provides an interface for automatically
# accepting new connections and creating sessions
# for them.
class Host(Thread):
    
    def __init__(self, listenPort=HOST_PORT, sessionClass=Session):
        Thread.__init__(self)

        # The Listener is a custom module that
        # basically automatically accepts tcp
        # connections and sticks them in a big
        # list. This collects all the peers that
        # want to log in.
        self._listener = Listener(listenPort)

        # This is a list of sessions that the
        # server is responsible for. These are
        # peers that want to talk to their friends.
        self._sessions = []
        self._sessionClass = sessionClass

        # internal. The host thread will keep
        # running until it is requested to shutdown.
        self._desireshutdown = False

    def run(self):
        self._initiate() # initiate the host (start up the listener)
        while not self._desireshutdown: # run until _shutdown == True or Host.close() is called
            self._update() # main host loop
        self._shutdown() # shutdown the host

    def close(self):
        # Kill the host!
        self._desireshutdown = True

    def _acceptsessions(self):
        # first, take any connections on the
        # listener and add the connection to
        # a list of sessions
        self._listener.lock()
        for connection in self._listener.connections:
            session = self._sessionClass(connection[0], connection[1])
            session.start()
            self._onconnection(session)
            self._sessions.append(session)
        self._listener.connections = []
        self._listener.release()
        
    def _handlesessions(self):
        deadsessions = []
        for session in self._sessions:
            if session.isdead() is True:
                deadsessions.append(session)
                self._ondisconnection(session)
            else:
                self._handlesession(session)
        for dead in deadsessions:
            self._sessions.remove(dead)

    def _handlesession(self, session):
        pass

    def _initiate(self):
        self._listener.start() # start the Listener module socket

    def _update(self):
            self._acceptsessions() # create new sessions from the listener
            self._handlesessions() # update session stuff
    
    def _shutdown(self):
        self._listener.close() # kill the Listener module
        self._listener.join() # wait for the thread to close
        #kill all the child sessions.
        for session in self._sessions:
            session.close()
        self._sessions = []

    # called on a session whenever it's created
    def _onconnection(self, session):
        print("New session from: " + str(session.getremoteaddress()))
        pass

    # called on a session whenever it is destroyed
    def _ondisconnection(self, session):
        pass