from socket import *

class TCPSocket:
  """
    Socket wrapper
    Takes care of packet lengths and encryption
  """
  #Socket stuff
  _sock = None		#Our socket
  _buffer = ""		#Data buffer
  _host = ""
  _port = 0

  #Callbacks
  _ondisc = None	#Disconnect callback

  def __init__(self, host="", port=0, ondisc=None):
    """
      Sets host, port and disconnect callback.
      Connects to host
    """
    self._host = host
    self._port = port
    self._ondisc = ondisc

    if host != "" and port != 0:
      self._sock = socket(AF_INET, SOCK_STREAM)
      self._sock.setblocking(0)
      try: self._sock.connect((host, port))
      except: pass

  def Receive(self):
    """
      Returns None if there is no packet
      Else returns the text of the packet
    """
    data = ''
    try: data = self._sock.recv(512)
    except: return None

    if data == '':
      #Server disconnected
      self.Disconnect()
      return None

    self._buffer += data

    a = b = n = 0
    if len(self._buffer) < 2: return None
    a = ord(self._buffer[0:1]) - 1	#Remove 1 since we added 1 to avoid sending \o
    b = ord(self._buffer[1:2]) - 1
    n = ntohs((a << 8) | b)		#Calculate packet length

    if len(self._buffer) >= n:
      packet = self._buffer[2:n-1]
      self._buffer = self._buffer[n:]
      return packet
    else:
      return None

  def Send(self, data):
    """
      Sends data to the server handles encryption
      Sends all of the data
    """
    total = len(data) + 2
    sent = 0

    n = htons(total)
    a = ((n & 0xFF00) >> 8) + 1
    b = (n & 0xFF) + 1

    data = chr(a) + chr(b) + data

    while sent < total:
      t = self._sock.send(data[sent:])
      if t <= 0:
	self.Disconnect()
	return
      sent += t

  def Disconnect(self):
    self._sock.close()
    self._ondisc()