# -*- coding: iso-8859-1 -*-
"""TCP functions and classes

Copyright © 2005-2012, Anders Andersen, University of Tromsų, Norway.
See http://www.cs.uit.no/~aa/dist/tools/noop/COPYING (../COPYING) for
details.


Provides the `TCPsocket` class and the `tcpsend` and `tcpreceive`
functions.  Below is a small example where to processes are
communicating.  The server process first (`IPaddr` from ip.sty):

>>> a = IPaddr(node=<host>, port=4576)
>>> msg = tcpreceive(a).decode("utf-8")			# Receive request
>>> tcpsend(a, ("Received: " + msg).encode("utf-8"))	# Send response

The client process:

>>> a = IPaddr(node=<host>, port=4576)
>>> tcpsend(a, "Hello world!".encode("utf-8"))		# Send request
>>> msg = tcpreceive(a).decode("utf-8")			# Receive response

The data of the `tcpsend` and `tcpreceive` functions and the `send`
and `receive` methods are bytes.  Therefor text strings has to be
converted to and from bytes using an encoding (e.g. utf-8).

"""


# CVS/RCS information

__file__    = r"$RCSfile: tcp.py,v $"
__version__ = r"$Revision: 1.32 $"
__date__    = r"$Date: 2012-11-21 09:03:18 $"
__state__   = r"$State: Exp $"
__author__  = r"$Author: aa $"

__log__     = r"""

$Log: tcp.py,v $
Revision 1.32  2012-11-21 09:03:18  aa
Use sendall and encode() instead of bytes when encoding the TCP stream.

Revision 1.31  2012-11-20 15:33:00  aa
Removed debug messages.

Revision 1.30  2012-11-20 15:11:53  aa
Changed exception error to OSError (Python 3.3)

Revision 1.29  2012-10-09 11:03:48  aa
Updated example in doc string to use `encode()` instead of `bytes()`.

Revision 1.28  2012-10-08 20:09:32  aa
Added copyright notice.

Revision 1.27  2011-10-25 15:31:44  aa
Added signatures to functions and methods.

Revision 1.26  2011-10-25 12:56:06  aa
Added examples to the document text of the library.

Revision 1.25  2011-10-25 12:12:41  aa
Added test to check if bind and listen in the `accept` method failed.

Revision 1.24  2011-10-18 21:40:06  aa
Made `accept` more robust (manage `socket.bind` errors).

Revision 1.23  2011-10-18 20:37:45  aa
Corrected minor typo (comment).

Revision 1.22  2011-10-18 20:30:23  aa
Updated for Python 3.

Revision 1.21  2005-03-22 17:18:31  aa
Added copyright notice.

Revision 1.20  2005/03/21 12:38:23  aa
Added documentation of the arguments and return values.

Revision 1.19  2005/03/20 16:31:24  aa
Updated the code to match the style guide for python code [2].

Revision 1.18  2005/03/14 12:39:37  aa
Added a debug message.

Revision 1.17  2005/03/14 10:28:59  aa
Minor correction (moved debug statement).

Revision 1.16  2005/03/14 10:14:09  aa
Rewrote it to use new version of the debug module.

Revision 1.15  2005/03/12 16:26:32  aa
Upgraded debug functions (fdebug and mdebug, see module debug).

Revision 1.14  2005/03/09 13:39:32  aa
Added debug information printing.

Revision 1.13  2005/03/09 11:03:52  aa
Removed overloading of the `socket` attribute.

Revision 1.12  2005/03/09 01:35:28  aa
Do not use string module, and try package (noop.ip.ip) import first.

Revision 1.11  2005/03/04 15:29:28  aa
Changed tcpServerSockets to a set, and remove id from set in tcpflush.

Revision 1.10  2005/03/04 10:34:27  aa
Fixed tcpflush so that the client closes the sockets before the server [1].

Revision 1.9  2005/03/04 09:09:02  aa
Cleaned up the clean up code.

Revision 1.8  2005/03/04 06:51:57  aa
Simplified send and receive with consocket.

Revision 1.7  2005/03/03 20:56:50  aa
Fixed problem with non-tcp adresses.

Revision 1.6  2005/03/03 20:42:56  aa
Do not enforce a given address family.

Revision 1.5  2005/03/03 14:22:00  aa
Removed timeout (use non-blocking in this rare case).

Revision 1.4  2005/03/03 13:27:43  aa
Fixed two way communication.

Revision 1.3  2005/03/03 13:13:39  aa
Improved how reuse of sockets are done.

Revision 1.2  2005/03/03 12:27:47  aa
Added comments and rcs information.

Revision 1.1  2005/03/03 11:27:38  aa
First ip implementation with examples.

"""

# [1] http://hea-www.harvard.edu/~fine/Tech/addrinuse.html
# [2] http://www.python.org/peps/pep-0008.html

# Standard Python libraries
from socket import *
from select import select
import sys


# Includes the IPaddr class (see addr arguments below)
try:
    from noop.ip.ip import *
except ImportError:
    from ip import *


# Sockets currently in use
_tcpsockets = {}

# Server sockets (close connection after client)
_tcpserversockets = set()


class TCPsocket:
    """An easy to use interface to TCP sockets

    An easy to use interface for TCP connetions using sockets.  The
    client side uses the constructor (providing an `IPaddr` address
    from the `ip` module) to create a socket (actually a list of
    sockets matching the provided address), the `connect` method to
    connect to the given address (the socket provinding the first
    successfull connection is used), and the `send` method to send
    data to the remote receiver.  The server side uses the constructor
    (providing an `IPaddr` address from the `ip` module) to create a
    socket (actually a list of sockets matching the provided address),
    the `accept` method to block and wait for a connection, and the
    `receive` to actually receive the data.  `send` annotates the data
    with the size actually sent, and `receive` will read exactly the
    data sent by a single `send`.

    """
    

    @signature
    def __init__(self, addr: IPaddr):
        """Create a socket and save the address

        Creates a socket (actually a list of sockets matching the
        provided address) and saves the address.  Only TCP
        (SOCK_STREAM) addresses are accepted.

        addr -- the address to bind or connect to (see ip.IPaddr)

        """

        # Save address
        self.addr = addr

        # Create sockets
        self.socket = {}
        for id, args in self.addr.socket_args():
            if args[TYPE] == SOCK_STREAM:
                self.socket[id] = socket(*args)

        # Any matching sockets?
        assert len(self.socket) > 0, \
               "TCPsocket: no SOCK_STREAM address provided"

        # Move to status UNUSED
        self.status = UNUSED
        #debugmsg(str(self.addr) + " (UNUSED)", self.__init__)


    @signature
    def connect(self):
        """Connect socket to address

        Tries to connect to connect to the remote address.  First
        successfull connection is used.
        
        """

        # The status has to be UNUSED
        assert self.status == UNUSED, "TCPsocket: connect: socket used"

        # The IPaddr address provides the arguments for each socket
        for id, args in self.addr.connect_args():
            if id in self.socket:
                
                # Try to connect using one of the sockets 
                try:
                    self.socket[id].connect(*args)

                # If no success, try next (but save error message)
                except OSError:
                    etype, value = sys.exc_info()[:2]
                    continue

                # If success, select this socket and move to status CONNECTED
                else:
                    self.status = CONNECTED
                    self.consocket = self.socket[id]
                    break

        # If no connections were accepted, raise an error
        if not hasattr(self, "consocket"):
            raise etype("TCPsocket: connect: unable to connect to socket: %s" \
                  % (value,))

        # Remove unused sockets
        #debugmsg(id + " (CONNECTED)", self.connect)
        self._clean(id)
        

    @signature
    def send(self, data: bytes):
        """Sends the data

        Sends the data using the connected socket.  The data are
        annotated with the size.

        data -- the data to send (a byte stream)

        """

        # Can only send it we have a connection
        assert self.status == CONNECTED, "TCPsocket: send: socket not connected"
            
        # Send the actual data annotated with the size
        self.consocket.sendall(str(len(data)).encode() + ":".encode() + data)
        #debugmsg(str(self.consocket.getsockname()) + " " + str(data), self.send)


    @signature
    def accept(self, num: opt(int) = 1):
        """Accept a connection

        Accept a connection on one socket.  The other sockets matching
        the address will never be used.

        num -- the argument to listen (default 1)

        """

        # Status has to be UNUSED
        assert self.status == UNUSED, "TCPsocket: accept: socket unused"

        # Bind a listen to every socket on the given address
        socketbound = False
        for id, args in self.addr.bind_args():
            if id in self.socket:
                try:
                    self.socket[id].bind(*args)
                    self.socket[id].listen(num)
                except OSError:
                    self.socket[id].close()
                    del self.socket[id]
                    continue
                else:
                    socketbound = True

        # Are we bound (and do we listen)?
        if not socketbound:
            raise OSError("No socket bound")

        # Wait for input on the sockets
        (isockets, o, e) = select(self.socket.values(), [], [])

        # We will only use one socket, so I expect this always to be true
        if len(isockets) == 1:

            # Accept the connection (now, receive can be used to fetch the data)
            recvsocket, addr = isockets[0].accept()
            self.consocket = recvsocket

        # If we have input on more sockets, we select one (NEVER HAPPENS?)
        else:

            # I do not like this; print a warning
            sys.stderr.write(
                "WARNING: only selecting one of many sockets with input.\n")
            sys.stderr.flush()

            # We will select one, but this will prevent accept from blocking
            for isocket in isockets:
                isocket.setblocking(0)

            # Try all sockets that listen says have input
            for isocket in isockets:

                # Is this alive
                try:
                    recvsocket, addr = isocket.accept()

                # Hmm problem, try next
                except timeout:
                    continue

                # Success, use this socket
                else:
                    self.consocket = recvsocket
                    break

            # If all had problems, raise an error
            if not hasattr(self, "consocket"):
                raise OSError("TCPsocket: accept: unable to accept connection.")

        # Move status to CONNECTED and remove unused socket
        self.status = CONNECTED
        _tcpserversockets.add(str(self.addr))
        #debugmsg(str(self.addr) + " (CONNECTED)", self.accept)
        self._clean()


    @signature
    def receive(self) -> bytes:
        """Receive the data

        Receive the data on the accepted connection.  This is actually
        more complex than you might think, since we only want to
        receive data from one `send`.  We use the size annotation
        added by `send` the get the right amount of data.  If we read
        to much, this is saved to next `receive`.

        returns --  the received data (a byte stream)

        """

        # Status has to be CONNECTED before we can receive anything
        assert self.status == CONNECTED, "TCPsocket: receive: socket unaccepted"

        # Read first chunk (either from a previous recv or from socket)
        if hasattr(self, "buffer"):
            msg = self.buffer
            del self.buffer
        else:
            msg = self.consocket.recv(BUFSIZ)

        # Read until we have the size tag (until first colon)
        while 1:
            try:
                (sz, pk) = msg.split(b":", 1)
            except ValueError:
                msg += self.consocket.recv(BUFSIZ)
            else:
                break

        # The size of the message
        size = int(sz)

        # Receive the rest of the message
        while len(pk) < size:
            pk = pk + self.consocket.recv(BUFSIZ)

        # Return data for this message, and save the rest
        #debugmsg(str(self.consocket.getsockname()) + " " + str(pk[:size]),
        #    self.receive)
        if len(pk) > size:
            self.buffer = pk[size:]
        return pk[:size]


    #@signature
    def _clean(self, uid: opt(str) = None):
        """Remove unused sockets

        This is used after a connection has been established to close
        and remove unused sockets.

        uid -- socket in use (used id), do not delete this (default None)

        """
        for id in list(self.socket.keys()):
            if id != uid:
                self.socket[id].close()
                del self.socket[id]
                #debugmsg("socket %s removed" % (id,), self._clean)
        self.socket = {}


    #@signature
    def __del__(self):
        """Clean up

        Close and remove the all sockets

        """
        self._clean()
        if hasattr(self, "consocket"):
            #debugmsg("socket %s removed" % (str(self.consocket.getsockname()),),
            #         self.__del__)
            self.consocket.close()
            del self.consocket
            pass
        #debugmsg(str(self.addr), self.__del__)


@signature
def tcpsend_new(addr: IPaddr, data: bytes):
    """Send to address with a new socket

    Create and connect to a socket, send the message, and remove the
    socket.

    addr -- the address to send data to (see ip.IPaddr)
    data -- the data to send (a byte stream)

    """
    tcpflush(addr)
    sendsocket = TCPsocket(addr)
    sendsocket.connect()
    sendsocket.send(data)


@signature
def tcpreceive_new(addr: IPaddr) -> bytes:
    """Receive from address with a new socket

    Create and accept connection to a socket, receive the message, and
    remove the socket.

    addr -- the address excpect data on (see ip.IPaddr)
    returns --  the received data (a byte stream)

    """
    tcpflush(addr)
    recvsocket = TCPsocket(addr)
    recvsocket.accept()
    return recvsocket.receive()


@signature
def tcpsend_reuse(addr: IPaddr, data: bytes):
    """Send to address and reuse socket

    If a socket mathing the address exits, reuse it.  A new socket is
    only created if a mathing socket doesn't exists.

    addr -- the address to send data to (see ip.IPaddr)
    data -- the data to send (a byte stream)

    """

    # Try to reuse
    if str(addr) in _tcpsockets:
        try:
            _tcpsockets[str(addr)].send(data)
            return
        except:
            tcpflush(addr)

    # Create new
    sendsocket = TCPsocket(addr)
    _tcpsockets[str(addr)] = sendsocket
    sendsocket.connect()
    sendsocket.send(data)


@signature
def tcpreceive_reuse(addr: IPaddr) -> bytes:
    """Reeceive from address and reuse socket

    If a socket mathing the address exits, reuse it.  A new socket is
    only created if a matching socket doesn't exists.

    addr -- the address excpect data on (see ip.IPaddr)
    returns --  the received data (a byte stream)
        
    """

    # Try to reuse
    if str(addr) in _tcpsockets:
        try:
            return _tcpsockets[str(addr)].receive()
        except:
            tcpflush(addr)

    # Create new
    recvsocket = TCPsocket(addr)
    _tcpsockets[str(addr)] = recvsocket
    recvsocket.accept()
    return recvsocket.receive()


@signature
def tcpflush(addr: opt(IPaddr) = None):
    """Remove sockets

    Remove a given socket or all sockets.  This is done gracefully,
    and we try to avoid the 'address already in use' problem [1].

    addr -- the address to flush (default None)

    """

    # Flush a single address?
    if addr:
        if str(addr) in _tcpsockets:
            ids = [str(addr)]
        else:
            ids = []

    # Or flush all?
    else:
        ids = list(_tcpsockets.keys())

    # Flush it (meaning remove them, meaning close the sockets and so on)
    for id in ids:

        # Client side should close before server side (see TIME_WAIT [1])
        if id in _tcpserversockets:

            # Wait for eof from client side (but do not wait forever)
            _tcpsockets[id].consocket.settimeout(CLOSETIMEOUT)
            s = _tcpsockets[id].consocket.recv(BUFSIZ)
            _tcpserversockets.remove(id)

        # Remove socket (see __del__ in TCPsocket)
        del _tcpsockets[id]


# The default is to reuse sockets
tcpsend = tcpsend_reuse
tcpreceive = tcpreceive_reuse


# Uncomment these two lines if you want a new socket for every send/receive
#tcpsend = tcpsend_new
#tcpreceive = tcpreceive_new