""" 
  <Description>
    This class tracks the changing state of a TCP connection.

    It'll raise errors when a application state change is
    impossible.
    
    It'll ignore any remotely requested state changes which
    are impossible.

    Otherwise, it'll change states according to the action
    and in many cases give advice on what should be done.

  <Usage>
    Object:
      TcpStateMachine() - get object
   
    Callback:
      state_machine.next(packet) - handle state changes for a packet

    Handshakers:
      state_machine.listen() - handle state changes for listen
      state_machine.connect() - handle state changes for connect
      state_machine.accept() - handle state changes for accept
      state_machine.close() - go back to closed

    Doers:
      state_machine.send() - see if send action is valid
      state_machine.recv() - see if recv action is valid
  
    Helper:
      state_machine.rollback() - go back one state because of error

  <Returns>
    The state machine gives these possible advices:
      OP_RCV_REMOTE_SYN - receive a SYN from a remote party
      OP_RCV_LOOPBACK_SYN - receive a SYN from a local party
      OP_RCV_SYNACK - receive a SYNACK from a party
      OP_RCV_PAYLOADACK - receive a payload final ACK for SYNACK
      OP_RCV_ACK - receive an ACK (maybe the final ACK of SYNACK)
      OP_RCV_PAYLOAD - receive a PAYLOAD
      OP_RCV_NONE - please ignore
      OP_RCV_FIN - receive the FIN
      OP_RCV_FINACK - receive the FINACK
      OP_RCV_ACK_FIN - receive the ACK of FIN

  <State Diagram>
    RFC 793 -- TCP

                                           Transmission Control Protocol
                                                Functional Specification

                              +---------+ ---------\      active OPEN  
                              |  CLOSED |            \    -----------  
                              +---------+<---------\   \   create TCB  
                                |     ^              \   \  snd SYN    
                   passive OPEN |     |   CLOSE        \   \           
                   ------------ |     | ----------       \   \         
                    create TCB  |     | delete TCB         \   \       
                                V     |                      \   \     
                              +---------+            CLOSE    |    \   
                              |  LISTEN |          ---------- |     |  
                              +---------+          delete TCB |     |  
                   rcv SYN      |     |     SEND              |     |  
                  -----------   |     |    -------            |     V  
 +---------+      snd SYN,ACK  /       \   snd SYN          +---------+
 |         |<-----------------           ------------------>|         |
 |   SYN   |                    rcv SYN                     |   SYN   |
 |   RCVD  |<-----------------------------------------------|   SENT  |
 |         |                    snd ACK                     |         |
 |         |------------------           -------------------|         |
 +---------+   rcv ACK of SYN  \       /  rcv SYN,ACK       +---------+
   |           --------------   |     |   -----------                  
   |                  x         |     |     snd ACK                    
   |                            V     V                                
   |  CLOSE                   +---------+                              
   | -------                  |  ESTAB  |                              
   | snd FIN                  +---------+                              
   |                   CLOSE    |     |    rcv FIN                     
   V                  -------   |     |    -------                     
 +---------+          snd FIN  /       \   snd ACK          +---------+
 |  FIN    |<-----------------           ------------------>|  CLOSE  |
 | WAIT-1  |------------------                              |   WAIT  |
 +---------+          rcv FIN  \                            +---------+
   | rcv ACK of FIN   -------   |                            CLOSE  |  
   | --------------   snd ACK   |                           ------- |  
   V        x                   V                           snd FIN V  
 +---------+                  +---------+                   +---------+
 |FINWAIT-2|                  | CLOSING |                   | LAST-ACK|
 +---------+                  +---------+                   +---------+
   |                rcv ACK of FIN |                 rcv ACK of FIN |  
   |  rcv FIN       -------------- |    Timeout=2MSL -------------- |  
   |  -------              x       V    ------------        x       V  
    \ snd ACK                 +---------+delete TCB         +---------+
     ------------------------>|TIME WAIT|------------------>| CLOSED  |
                              +---------+                   +---------+

                      TCP Connection State Diagram
                               Figure 6.
"""

include constants.repy
include errors.repy
include packet.repy 

class TcpStateMachine:
  """
  <Purpose>
    The state machine organizes the changes in a relationship
    between a pair of connected parties.
    
  <Example Use>
    sm = TcpStateMachine()
    def callback(ip, port, mess, ch):
      packet = unpack(mess)
      sm.next(packet)
    recvmess(getmyip(), 12345, callback)

  """
  
  def __init__(self):
    """
    <Purpose>
      Initialize the state machine's variables. Sets state
      to be a closed connection.
    
    <Arguments>
      None
    
    <Exceptions>
      None
    
    <Side Effects>
      Call getlock().
      Sets state to be STATE_CLOSED
     
    <Returns>
      The new state machine object.
    """
    (self.STATE_CLOSED, self.STATE_LISTEN, self.STATE_SYN_SENT, \
     self.STATE_SYN_RCVD, self.STATE_ESTAB, self.STATE_FIN_WAIT_1, \
     self.STATE_FIN_WAIT_2, self.STATE_CLOSING, self.STATE_CLOSE_WAIT, \
     self.STATE_LAST_ACK, self.STATE_TIME_WAIT) = range(0, 11)

    # not initialized yet
    self.state = self.STATE_CLOSED
    self.last_state = None # for rolling back

    # state changes must be atomic
    self.lock = getlock()

  #######################
  # Callback
  #####################

  def next(self, packet):
    """
    <Purpose>
      Handles all callbacks.  This function is the engine
      of a state machine.
    
    <Arguments>
      packet - The packet to change on.
    
    <Exceptions>
      None
    
    <Side Effects>
      None

    <Returns>
      An opcode specifying the exact transition
      made.  It is one of the advice constants.
    """
    self.lock.acquire()       
    advice = OP_RCV_NONE
    
    self.last_state = self.state
    if(self.state == self.STATE_CLOSED): 
      pass # ignore

    elif(self.state == self.STATE_LISTEN):
      if packet.is_syn():
        self.state = self.STATE_SYN_RCVD
        advice = OP_RCV_REMOTE_SYN

    elif(self.state == self.STATE_SYN_SENT): 
      # finish the handshake with an ACK...
      if packet.is_synack():
        self.state = self.STATE_ESTAB
        advice = OP_RCV_SYNACK
      # loopback special case: recv own SYN!
      elif packet.is_syn():
        self.state = self.STATE_SYN_RCVD
        advice = OP_RCV_LOOPBACK_SYN
      else:
	pass # ignore

    elif(self.state == self.STATE_SYN_RCVD):
      # skip the ack then...
      if packet.is_payload():
        self.state = self.STATE_ESTAB
        advice = OP_RCV_PAYLOADACK
      # ok ready
      elif packet.is_ack():
        self.state = self.STATE_ESTAB
	advice = OP_RCV_ACK
      else:
	pass

    #####################
    # Established
    #################

    elif(self.state == self.STATE_ESTAB):
      # count packet or ignore old ack
      if packet.is_ack():
        advice = OP_RCV_ACK
      # Send back an ack or reack.
      elif packet.is_payload():
        advice = OP_RCV_PAYLOAD
      # old handshake? how about old syn?
      elif packet.is_synack():
        advice = OP_RCV_SYNACK
      # close
      elif packet.is_fin():
        self.state = self.STATE_CLOSE_WAIT
        advice = OP_RCV_FIN
      else:
	pass

    ######################
    # Closing Handshake States
    #######################

    elif(self.state == self.STATE_FIN_WAIT_1):
      # rcvd ACK of fin
      if packet.is_ack():
	self.state = self.STATE_FIN_WAIT_2
        advice = OP_RCV_ACK_FIN
      # send ACK of FIN
      elif packet.is_fin():
        self.state = self.STATE_CLOSING
        advice = OP_RCV_FIN
      # send ACK of FIN and close
      elif packet.is_finack():
        self.state = self.STATE_TIME_WAIT
        advice = OP_RCV_FINACK
        # self.state = self.STATE_CLOSED # need to timeout wait
      else:
	pass

    elif(self.state == self.STATE_FIN_WAIT_2):
      # send ACK of FIN
      if packet.is_fin():
        self.state = self.STATE_TIME_WAIT
        advice = OP_RCV_FIN
 	# self.state = self.STATE_CLOSED # need to timeout wait
      elif packet.is_payload(): # Half-open estab
        advice = OP_RCV_PAYLOAD
      else:
        pass

    elif(self.state == self.STATE_CLOSING):
      # rcv ACK of FIN
      if packet.is_ack():
        advice = OP_RCV_ACK_FIN
	# self.state = self.STATE_TIME_WAIT
        self.state = self.STATE_CLOSED # Do I need to timeout wait?
      if packet.is_fin(): # reack fin
        advice = OP_RCV_FIN
      else:
	pass

    elif(self.state == self.STATE_CLOSE_WAIT):
      if packet.is_ack(): # Half-open estab
        advice = OP_RCV_ACK
      else: 
        pass

    elif(self.state == self.STATE_LAST_ACK):
      # rcv ACK of FIN
      if packet.is_ack():
        advice = OP_RCV_ACK_FIN
        self.state = self.STATE_CLOSED
      else:
	pass

    elif(self.state == self.STATE_TIME_WAIT):
      if packet.is_fin(): # reack fin
        advice = OP_RCV_FIN
      elif packet.is_fin(): # reack fin
        advice = OP_RCV_FIN
      else:
        pass

    else: # garbled packet
      pass # ignore it

    self.lock.release()
    return advice

  ##########################3
  # Handshakers
  ############################

  def listen(self):
    """
    <Purpose>   
      Do any state changes to listen.

    <Arguments>
      None
          
    <Exceptions>
      raise AlreadyConnectedError("cannot listen because the connection is already in use")
    
    <Side Effects>
      Sets the state to STATE_LISTEN.

    <Returns>
      None
    """
    self.lock.acquire()
    if self.state != self.STATE_CLOSED:
      self.lock.release()
      raise AlreadyConnectedError("cannot listen because the connection is already in use")
    self.state = self.STATE_LISTEN 
    self.lock.release()

  def connect(self):
    """
    <Purpose>   
      Do any state changes to start connecting.

    <Arguments>
      None
          
    <Exceptions>
      raise AlreadyConnectedError("cannot connect because the connection is already in use")
    
    <Side Effects>
      Changes state to STATE_SYN_SENT.     

    <Returns>
      None
    """

    self.lock.acquire()
    if self.state != self.STATE_CLOSED:
      self.lock.release()
      raise AlreadyConnectedError("cannot connect because the connection is already in use")
    self.state = self.STATE_SYN_SENT
    self.lock.release()

  def accept(self):
    """
    <Purpose>   
      Do any state changes for a recv.

    <Arguments>
      None
          
    <Exceptions>
      raise NotRequestedError("cannot accept without a requested connection")
      NotRequestedError if not in STATE_SYN_RCVD from listen()
      returning.
    
    <Side Effects>
      None
     
    <Returns>
      None
    """
    self.lock.acquire()
    if self.state != self.STATE_SYN_RCVD:
      self.lock.release()
      raise NotRequestedError("cannot accept without a requested connection")
    self.lock.release()

  def close(self):
    """
    <Purpose>
      Does what is needed to close a connection.
    
    <Arguments>
      None
    
    <Exceptions>
      None
    
    <Side Effects>
      Set the state to STATE_CLOSED
     
    <Returns>
      True if the conn should send a FIN. False to
      recommend the client to take no action.
    """
    self.lock.acquire()
    advice = True

    if self.state == self.STATE_FIN_WAIT_1 \
       or self.state == self.STATE_FIN_WAIT_2 \
       or self.state == self.STATE_LAST_ACK \
       or self.state == self.STATE_TIME_WAIT \
       or self.state == self.STATE_CLOSING:
      self.lock.release()
      raise AlreadyClosingError("cannot close an already closing connection")

    elif self.state == self.STATE_CLOSED \
       or self.state == self.STATE_LISTEN \
       or self.state == self.STATE_SYN_SENT:
      self.state = self.STATE_CLOSED
      advice = False # don't need to send a FIN

    elif self.state == self.STATE_ESTAB \
       or self.state == self.STATE_SYN_RCVD:
      self.state = self.STATE_FIN_WAIT_1

    else: # CLOSE_WAIT
      assert self.state == self.STATE_CLOSE_WAIT
      self.state = self.STATE_LAST_ACK

    self.lock.release()
    return advice
   
  ##########################
  # Doers
  ###########################

  def send(self):
    """
    <Purpose>
      Do any state changes for a send.
    
    <Arguments>
      None
    
    <Exceptions>
      raise NotEstablishedError("must be established to send")
    
    <Side Effects>
      None
     
    <Returns>
      None
    """
    self.lock.acquire()
    if self.state != self.STATE_ESTAB:
      self.lock.release()
      raise NotEstablishedError("must be established to send")
    self.lock.release()

  def recv(self):
    """
    <Purpose>
      Do any state changes for a recv.
    
    <Arguments>
      None
    
    <Exceptions>
      raise NotEstablishedError("must be established to recv")
    
    <Side Effects>
      None
     
    <Returns>
      None
    """
    self.lock.acquire()
    if self.state != self.STATE_ESTAB:
      self.lock.release()
      raise NotEstablishedError("must be established to recv")
    self.lock.release()


  ##########################
  # Helpers
  ##########################

  def rollback(self):
    """
    <Purpose>
      Reverts back one state in case of an error 
      after a state machine.

    <Side Effects>
      Sets last_state to None
    """
    self.state = self.last_state
    self.last_state = None

  def timeout_close(self):
    """Wait a little to maybe re-ACK final FIN"""
    self.lock.acquire()
    if self.state != self.STATE_TIME_WAIT:
      print self.state
      raise Exception("need to be closing to do final timewait")
    self.lock.release()

    sleep(.002)

    self.lock.acquire()
    self.state = self.state = self.STATE_CLOSED
    self.lock.release()

  def time_waiting(self):
    return self.state == self.STATE_TIME_WAIT
