Source Code for Module Peach.Publishers.udp

  1   
  2  ''' 
  3  Default UDP publishers. 
  4   
  5  @author: Michael Eddington 
  6  @version: $Id: Peach.Publishers.udp-pysrc.html 1138 2008-08-16 19:39:03Z meddingt $ 
  7  ''' 
  8   
  9  # 
 10  # Copyright (c) 2005-2008 Michael Eddington 
 11  # Copyright (c) 2004-2005 IOActive Inc. 
 12  # 
 13  # Permission is hereby granted, free of charge, to any person obtaining a copy  
 14  # of this software and associated documentation files (the "Software"), to deal 
 15  # in the Software without restriction, including without limitation the rights  
 16  # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell  
 17  # copies of the Software, and to permit persons to whom the Software is  
 18  # furnished to do so, subject to the following conditions: 
 19  # 
 20  # The above copyright notice and this permission notice shall be included in     
 21  # all copies or substantial portions of the Software. 
 22  # 
 23  # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR  
 24  # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,  
 25  # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE  
 26  # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER  
 27  # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 
 28  # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 
 29  # SOFTWARE. 
 30  # 
 31   
 32  # Authors: 
 33  #   Michael Eddington (mike@phed.org) 
 34   
 35  # $Id: Peach.Publishers.udp-pysrc.html 1138 2008-08-16 19:39:03Z meddingt $ 
 36   
 37  import socket, time 
 38  from Peach.publisher import * 
 39  import Peach 
 40   
 41 -def Debug(msg): 
 42          if Peach.Engine.engine.Engine.debug: 
 43                  print msg 
 44   
 45 -class Udp(Publisher): 
 46          ''' 
 47          A simple UDP publisher. 
 48          ''' 
 49           
 50 -        def __init__(self, host, port, timeout = None): 
 51                  ''' 
 52                  @type   host: string 
 53                  @param  host: Remote hostname 
 54                  @type   port: number 
 55                  @param  port: Remote port 
 56                  ''' 
 57                  self._host = host 
 58                  self._port = port 
 59                  self._timeout = timeout 
 60                   
 61                  if self._timeout == None: 
 62                          self._timeout = 2 
 63                   
 64                  else: 
 65                          self._timeout = timeout 
 66                   
 67                  self._socket = None 
 68                  self.buff = "" 
 69                  self.pos = 0 
 70           
 71 -        def stop(self): 
 72                  '''Close connection if open''' 
 73                  self.close() 
 74           
 75 -        def close(self): 
 76                  if self._socket != None: 
 77                          self._socket.close() 
 78                          self._socket = None 
 79                  self.buff = "" 
 80                  self.pos = 0 
 81           
 82 -        def connect(self): 
 83                  if self._socket != None: 
 84                          # Close out old socket first 
 85                          self._socket.close() 
 86                  self._socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) 
 87                  self._socket.connect((self._host,int(self._port))) 
 88                  self.buff = "" 
 89                  self.pos = 0 
 90           
 91 -        def send(self, data): 
 92                  ''' 
 93                  Send data via sendall. 
 94                   
 95                  @type   data: string 
 96                  @param  data: Data to send 
 97                  ''' 
 98                  self._socket.sendall(data) 
 99           
100 -        def _receiveBySize(self, size): 
101                  ''' 
102                  This is now a buffered receiver. 
103                   
104                  @rtype: string 
105                  @return: received data. 
106                  ''' 
107                   
108                  # Do we already a have it? 
109                  if size+self.pos < len(self.buff): 
110                          ret = self.buff[self.pos:self.pos+size] 
111                          self.pos += size 
112                          return ret 
113                   
114                  # Only ask for the diff of what we don't already have 
115                  diffSize = (self.pos+size) - len(self.buff) 
116                   
117                  try: 
118                          self._socket.settimeout(float(self._timeout)) 
119                          ret = self._socket.recv(65507) 
120                           
121                          Debug("udp.Udp.receive we got:") 
122                          Debug(repr(ret)) 
123                           
124                          if not ret: 
125                                  # Socket was closed 
126                                  raise PublisherSoftException("Socket is closed") 
127                           
128                          self.buff += ret 
129                           
130                  except socket.error, e: 
131                          if str(e).find('The socket operation could not complete without blocking') == -1: 
132                                  raise Timeout("Timed out waiting for data [%d:%d:%d:%d]" % (len(self.buff), (size+self.pos), size, diffSize)) 
133                           
134                          else: 
135                                  raise PublisherSoftException("recv failed: " + str(sys.exc_info()[1])) 
136                   
137                  finally: 
138                          self._socket.settimeout(None) 
139                   
140                  ret = self.buff[self.pos:] 
141                  self.pos = len(self.buff) 
142                  return ret 
143           
144 -        def _receiveByAvailable(self): 
145                  ''' 
146                  Receive as much as possible prior to timeout. 
147                   
148                  @rtype: string 
149                  @return: received data. 
150                  ''' 
151                  self._socket.settimeout(float(self._timeout)) 
152                   
153                  try: 
154                          ret = self._socket.recv(65507) 
155                           
156                          Debug("udp.Udp.receive we got:") 
157                          Debug(repr(ret)) 
158                           
159                          if not ret: 
160                                  raise PublisherSoftException("Socket is closed") 
161                           
162                          self.buff += ret 
163                           
164                  except socket.error, e: 
165                          if str(e).find('The socket operation could not complete without blocking') == -1: 
166                                  pass 
167                           
168                          else: 
169                                  raise PublisherSoftException("recv failed: " + str(sys.exc_info()[1])) 
170                   
171                  self._socket.settimeout(None) 
172                   
173                  ret = self.buff[self.pos:] 
174                  self.pos = len(self.buff) 
175                  return ret 
176           
177 -        def receive(self, size = None): 
178                   
179                  if size == None: 
180                          return self._receiveByAvailable() 
181                  else: 
182                          return self._receiveBySize(size) 
183                   
184  # end 
185