source: python-simple-websocket-server/simple-websocket-server/SimpleWebSocketServer/SimpleWebSocketServer.py @ 1309

Last change on this file since 1309 was 1309, checked in by joamuran, 5 years ago

initial upload
File size: 21.9 KB
Line 
1'''
2The MIT License (MIT)
3Copyright (c) 2013 Dave P.
4'''
5import sys
6VER = sys.version_info[0]
7if VER >= 3:
8    import socketserver
9    from http.server import BaseHTTPRequestHandler
10    from io import StringIO, BytesIO
11else:
12    import SocketServer
13    from BaseHTTPServer import BaseHTTPRequestHandler
14    from StringIO import StringIO
15
16import hashlib
17import base64
18import socket
19import struct
20import ssl
21import errno
22import codecs
23from collections import deque
24from select import select
25
26__all__ = ['WebSocket',
27            'SimpleWebSocketServer',
28            'SimpleSSLWebSocketServer']
29
30def _check_unicode(val):
31    if VER >= 3:
32        return isinstance(val, str)
33    else:
34        return isinstance(val, unicode)
35
36class HTTPRequest(BaseHTTPRequestHandler):
37   def __init__(self, request_text):
38      if VER >= 3:
39          self.rfile = BytesIO(request_text)
40      else:
41          self.rfile = StringIO(request_text)
42      self.raw_requestline = self.rfile.readline()
43      self.error_code = self.error_message = None
44      self.parse_request()
45
46_VALID_STATUS_CODES = [1000, 1001, 1002, 1003, 1007, 1008,
47                        1009, 1010, 1011, 3000, 3999, 4000, 4999]
48
49HANDSHAKE_STR = (
50   "HTTP/1.1 101 Switching Protocols\r\n"
51   "Upgrade: WebSocket\r\n"
52   "Connection: Upgrade\r\n"
53   "Sec-WebSocket-Accept: %(acceptstr)s\r\n\r\n"
54)
55
56GUID_STR = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11'
57
58STREAM = 0x0
59TEXT = 0x1
60BINARY = 0x2
61CLOSE = 0x8
62PING = 0x9
63PONG = 0xA
64
65HEADERB1 = 1
66HEADERB2 = 3
67LENGTHSHORT = 4
68LENGTHLONG = 5
69MASK = 6
70PAYLOAD = 7
71
72MAXHEADER = 65536
73MAXPAYLOAD = 33554432
74
75class WebSocket(object):
76
77   def __init__(self, server, sock, address):
78      self.server = server
79      self.client = sock
80      self.address = address
81
82      self.handshaked = False
83      self.headerbuffer = bytearray()
84      self.headertoread = 2048
85
86      self.fin = 0
87      self.data = bytearray()
88      self.opcode = 0
89      self.hasmask = 0
90      self.maskarray = None
91      self.length = 0
92      self.lengtharray = None
93      self.index = 0
94      self.request = None
95      self.usingssl = False
96
97      self.frag_start = False
98      self.frag_type = BINARY
99      self.frag_buffer = None
100      self.frag_decoder = codecs.getincrementaldecoder('utf-8')(errors='strict')
101      self.closed = False
102      self.sendq = deque()
103
104      self.state = HEADERB1
105
106      # restrict the size of header and payload for security reasons
107      self.maxheader = MAXHEADER
108      self.maxpayload = MAXPAYLOAD
109
110   def handleMessage(self):
111      """
112          Called when websocket frame is received.
113          To access the frame data call self.data.
114
115          If the frame is Text then self.data is a unicode object.
116          If the frame is Binary then self.data is a bytearray object.
117      """
118      pass
119
120   def handleConnected(self):
121      """
122          Called when a websocket client connects to the server.
123      """
124      pass
125
126   def handleClose(self):
127      """
128          Called when a websocket server gets a Close frame from a client.
129      """
130      pass
131
132   def _handlePacket(self):
133      if self.opcode == CLOSE:
134         pass
135      elif self.opcode == STREAM:
136         pass
137      elif self.opcode == TEXT:
138         pass
139      elif self.opcode == BINARY:
140         pass
141      elif self.opcode == PONG or self.opcode == PING:
142         if len(self.data) > 125:
143            raise Exception('control frame length can not be > 125')
144      else:
145          # unknown or reserved opcode so just close
146         raise Exception('unknown opcode')
147
148      if self.opcode == CLOSE:
149         status = 1000
150         reason = u''
151         length = len(self.data)
152
153         if length == 0:
154            pass
155         elif length >= 2:
156            status = struct.unpack_from('!H', self.data[:2])[0]
157            reason = self.data[2:]
158
159            if status not in _VALID_STATUS_CODES:
160                status = 1002
161
162            if len(reason) > 0:
163                try:
164                    reason = reason.decode('utf8', errors='strict')
165                except:
166                    status = 1002
167         else:
168            status = 1002
169
170         self.close(status, reason)
171         return
172
173      elif self.fin == 0:
174          if self.opcode != STREAM:
175              if self.opcode == PING or self.opcode == PONG:
176                  raise Exception('control messages can not be fragmented')
177
178              self.frag_type = self.opcode
179              self.frag_start = True
180              self.frag_decoder.reset()
181
182              if self.frag_type == TEXT:
183                  self.frag_buffer = []
184                  utf_str = self.frag_decoder.decode(self.data, final = False)
185                  if utf_str:
186                      self.frag_buffer.append(utf_str)
187              else:
188                  self.frag_buffer = bytearray()
189                  self.frag_buffer.extend(self.data)
190
191          else:
192              if self.frag_start is False:
193                  raise Exception('fragmentation protocol error')
194
195              if self.frag_type == TEXT:
196                  utf_str = self.frag_decoder.decode(self.data, final = False)
197                  if utf_str:
198                      self.frag_buffer.append(utf_str)
199              else:
200                  self.frag_buffer.extend(self.data)
201
202      else:
203          if self.opcode == STREAM:
204              if self.frag_start is False:
205                  raise Exception('fragmentation protocol error')
206
207              if self.frag_type == TEXT:
208                  utf_str = self.frag_decoder.decode(self.data, final = True)
209                  self.frag_buffer.append(utf_str)
210                  self.data = u''.join(self.frag_buffer)
211              else:
212                  self.frag_buffer.extend(self.data)
213                  self.data = self.frag_buffer
214
215              self.handleMessage()
216
217              self.frag_decoder.reset()
218              self.frag_type = BINARY
219              self.frag_start = False
220              self.frag_buffer = None
221
222          elif self.opcode == PING:
223              self._sendMessage(False, PONG, self.data)
224
225          elif self.opcode == PONG:
226              pass
227
228          else:
229              if self.frag_start is True:
230                  raise Exception('fragmentation protocol error')
231
232              if self.opcode == TEXT:
233                  try:
234                      self.data = self.data.decode('utf8', errors='strict')
235                  except Exception as exp:
236                      raise Exception('invalid utf-8 payload')
237
238              self.handleMessage()
239
240
241   def _handleData(self):
242      # do the HTTP header and handshake
243      if self.handshaked is False:
244
245         data = self.client.recv(self.headertoread)
246         if not data:
247            raise Exception("remote socket closed")
248
249         else:
250            # accumulate
251            self.headerbuffer.extend(data)
252
253            if len(self.headerbuffer) >= self.maxheader:
254               raise Exception('header exceeded allowable size')
255
256            # indicates end of HTTP header
257            if b'\r\n\r\n' in self.headerbuffer:
258               self.request = HTTPRequest(self.headerbuffer)
259
260               # handshake rfc 6455
261               try:
262                  key = self.request.headers['Sec-WebSocket-Key']
263                  k = key.encode('ascii') + GUID_STR.encode('ascii')
264                  k_s = base64.b64encode(hashlib.sha1(k).digest()).decode('ascii')
265                  hStr = HANDSHAKE_STR % {'acceptstr': k_s}
266                  self.sendq.append((BINARY, hStr.encode('ascii')))
267                  self.handshaked = True
268                  self.handleConnected()
269               except Exception as e:
270                  raise Exception('handshake failed: %s', str(e))
271
272      # else do normal data
273      else:
274         data = self.client.recv(8192)
275         if not data:
276            raise Exception("remote socket closed")
277
278         if VER >= 3:
279             for d in data:
280                 self._parseMessage(d)
281         else:
282             for d in data:
283                 self._parseMessage(ord(d))
284
285   def close(self, status = 1000, reason = u''):
286       """
287          Send Close frame to the client. The underlying socket is only closed
288          when the client acknowledges the Close frame.
289
290          status is the closing identifier.
291          reason is the reason for the close.
292        """
293       try:
294          if self.closed is False:
295            close_msg = bytearray()
296            close_msg.extend(struct.pack("!H", status))
297            if _check_unicode(reason):
298                close_msg.extend(reason.encode('utf-8'))
299            else:
300                close_msg.extend(reason)
301
302            self._sendMessage(False, CLOSE, close_msg)
303
304       finally:
305            self.closed = True
306
307
308   def _sendBuffer(self, buff):
309      size = len(buff)
310      tosend = size
311      already_sent = 0
312
313      while tosend > 0:
314         try:
315            # i should be able to send a bytearray
316            sent = self.client.send(buff[already_sent:])
317            if sent == 0:
318               raise RuntimeError("socket connection broken")
319
320            already_sent += sent
321            tosend -= sent
322
323         except socket.error as e:
324            # if we have full buffers then wait for them to drain and try again
325            if e.errno in [errno.EAGAIN, errno.EWOULDBLOCK]:
326               return buff[already_sent:]
327            else:
328               raise e
329
330      return None
331
332   def sendFragmentStart(self, data):
333      """
334          Send the start of a data fragment stream to a websocket client.
335          Subsequent data should be sent using sendFragment().
336          A fragment stream is completed when sendFragmentEnd() is called.
337
338          If data is a unicode object then the frame is sent as Text.
339          If the data is a bytearray object then the frame is sent as Binary.
340      """
341      opcode = BINARY
342      if _check_unicode(data):
343         opcode = TEXT
344      self._sendMessage(True, opcode, data)
345
346   def sendFragment(self, data):
347      """
348          see sendFragmentStart()
349
350          If data is a unicode object then the frame is sent as Text.
351          If the data is a bytearray object then the frame is sent as Binary.
352      """
353      self._sendMessage(True, STREAM, data)
354
355   def sendFragmentEnd(self, data):
356      """
357          see sendFragmentEnd()
358
359          If data is a unicode object then the frame is sent as Text.
360          If the data is a bytearray object then the frame is sent as Binary.
361      """
362      self._sendMessage(False, STREAM, data)
363
364   def sendMessage(self, data):
365      """
366          Send websocket data frame to the client.
367
368          If data is a unicode object then the frame is sent as Text.
369          If the data is a bytearray object then the frame is sent as Binary.
370      """
371      opcode = BINARY
372      if _check_unicode(data):
373         opcode = TEXT
374      self._sendMessage(False, opcode, data)
375
376
377   def _sendMessage(self, fin, opcode, data):
378
379        payload = bytearray()
380
381        b1 = 0
382        b2 = 0
383        if fin is False:
384           b1 |= 0x80
385        b1 |= opcode
386
387        if _check_unicode(data):
388           data = data.encode('utf-8')
389
390        length = len(data)
391        payload.append(b1)
392
393        if length <= 125:
394           b2 |= length
395           payload.append(b2)
396
397        elif length >= 126 and length <= 65535:
398           b2 |= 126
399           payload.append(b2)
400           payload.extend(struct.pack("!H", length))
401
402        else:
403           b2 |= 127
404           payload.append(b2)
405           payload.extend(struct.pack("!Q", length))
406
407        if length > 0:
408           payload.extend(data)
409
410        self.sendq.append((opcode, payload))
411
412
413   def _parseMessage(self, byte):
414      # read in the header
415      if self.state == HEADERB1:
416
417         self.fin = byte & 0x80
418         self.opcode = byte & 0x0F
419         self.state = HEADERB2
420
421         self.index = 0
422         self.length = 0
423         self.lengtharray = bytearray()
424         self.data = bytearray()
425
426         rsv = byte & 0x70
427         if rsv != 0:
428            raise Exception('RSV bit must be 0')
429
430      elif self.state == HEADERB2:
431         mask = byte & 0x80
432         length = byte & 0x7F
433
434         if self.opcode == PING and length > 125:
435             raise Exception('ping packet is too large')
436
437         if mask == 128:
438            self.hasmask = True
439         else:
440            self.hasmask = False
441
442         if length <= 125:
443            self.length = length
444
445            # if we have a mask we must read it
446            if self.hasmask is True:
447               self.maskarray = bytearray()
448               self.state = MASK
449            else:
450               # if there is no mask and no payload we are done
451               if self.length <= 0:
452                  try:
453                     self._handlePacket()
454                  finally:
455                     self.state = self.HEADERB1
456                     self.data = bytearray()
457
458               # we have no mask and some payload
459               else:
460                  #self.index = 0
461                  self.data = bytearray()
462                  self.state = PAYLOAD
463
464         elif length == 126:
465            self.lengtharray = bytearray()
466            self.state = LENGTHSHORT
467
468         elif length == 127:
469            self.lengtharray = bytearray()
470            self.state = LENGTHLONG
471
472
473      elif self.state == LENGTHSHORT:
474         self.lengtharray.append(byte)
475
476         if len(self.lengtharray) > 2:
477            raise Exception('short length exceeded allowable size')
478
479         if len(self.lengtharray) == 2:
480            self.length = struct.unpack_from('!H', self.lengtharray)[0]
481
482            if self.hasmask is True:
483               self.maskarray = bytearray()
484               self.state = MASK
485            else:
486               # if there is no mask and no payload we are done
487               if self.length <= 0:
488                  try:
489                     self._handlePacket()
490                  finally:
491                     self.state = HEADERB1
492                     self.data = bytearray()
493
494               # we have no mask and some payload
495               else:
496                  #self.index = 0
497                  self.data = bytearray()
498                  self.state = PAYLOAD
499
500      elif self.state == LENGTHLONG:
501
502         self.lengtharray.append(byte)
503
504         if len(self.lengtharray) > 8:
505            raise Exception('long length exceeded allowable size')
506
507         if len(self.lengtharray) == 8:
508            self.length = struct.unpack_from('!Q', self.lengtharray)[0]
509
510            if self.hasmask is True:
511               self.maskarray = bytearray()
512               self.state = MASK
513            else:
514               # if there is no mask and no payload we are done
515               if self.length <= 0:
516                  try:
517                     self._handlePacket()
518                  finally:
519                     self.state = HEADERB1
520                     self.data = bytearray()
521
522               # we have no mask and some payload
523               else:
524                  #self.index = 0
525                  self.data = bytearray()
526                  self.state = PAYLOAD
527
528      # MASK STATE
529      elif self.state == MASK:
530         self.maskarray.append(byte)
531
532         if len(self.maskarray) > 4:
533            raise Exception('mask exceeded allowable size')
534
535         if len(self.maskarray) == 4:
536            # if there is no mask and no payload we are done
537            if self.length <= 0:
538               try:
539                  self._handlePacket()
540               finally:
541                  self.state = HEADERB1
542                  self.data = bytearray()
543
544            # we have no mask and some payload
545            else:
546               #self.index = 0
547               self.data = bytearray()
548               self.state = PAYLOAD
549
550      # PAYLOAD STATE
551      elif self.state == PAYLOAD:
552         if self.hasmask is True:
553            self.data.append( byte ^ self.maskarray[self.index % 4] )
554         else:
555            self.data.append( byte )
556
557         # if length exceeds allowable size then we except and remove the connection
558         if len(self.data) >= self.maxpayload:
559            raise Exception('payload exceeded allowable size')
560
561         # check if we have processed length bytes; if so we are done
562         if (self.index+1) == self.length:
563            try:
564               self._handlePacket()
565            finally:
566               #self.index = 0
567               self.state = HEADERB1
568               self.data = bytearray()
569         else:
570            self.index += 1
571
572
573class SimpleWebSocketServer(object):
574   def __init__(self, host, port, websocketclass, selectInterval = 0.1):
575      self.websocketclass = websocketclass
576      self.serversocket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
577      self.serversocket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
578      self.serversocket.bind((host, port))
579      self.serversocket.listen(5)
580      self.selectInterval = selectInterval
581      self.connections = {}
582      self.listeners = [self.serversocket]
583
584   def _decorateSocket(self, sock):
585      return sock
586
587   def _constructWebSocket(self, sock, address):
588      return self.websocketclass(self, sock, address)
589
590   def close(self):
591      self.serversocket.close()
592
593      for desc, conn in self.connections.items():
594         conn.close()
595         try:
596            conn.handleClose()
597         except:
598            pass
599
600   def serveforever(self):
601      while True:
602         writers = []
603         for fileno in self.listeners:
604            try:
605               client = self.connections[fileno]
606               if client.sendq:
607                  writers.append(fileno)
608            except Exception as n:
609               pass
610
611         if self.selectInterval:
612            rList, wList, xList = select(self.listeners, writers, self.listeners, self.selectInterval)
613         else:
614            rList, wList, xList = select(self.listeners, writers, self.listeners)
615
616         for ready in wList:
617            client = None
618            try:
619               client = self.connections[ready]
620               while client.sendq:
621                  opcode, payload = client.sendq.popleft()
622                  remaining = client._sendBuffer(payload)
623                  if remaining is not None:
624                      client.sendq.appendleft((opcode, remaining))
625                      break
626                  else:
627                      if opcode == CLOSE:
628                         raise Exception("received client close")
629
630            except Exception as n:
631
632               if client:
633                  client.client.close()
634
635               try:
636                  if client:
637                     client.handleClose()
638               except:
639                  pass
640
641               try:
642                  del self.connections[ready]
643               except:
644                  pass
645
646               try:
647                  self.listeners.remove(ready)
648               except:
649                  pass
650
651         for ready in rList:
652            if ready == self.serversocket:
653               try:
654                  sock, address = self.serversocket.accept()
655                  newsock = self._decorateSocket(sock)
656                  newsock.setblocking(0)
657                  fileno = newsock.fileno()
658                  self.listeners.append(fileno)
659                  self.connections[fileno] = self._constructWebSocket(newsock, address)
660               except Exception as n:
661                  if sock is not None:
662                     sock.close()
663            else:
664               client = None
665               try:
666                  client = self.connections[ready]
667                  client._handleData()
668               except Exception as n:
669                  if client:
670                     client.client.close()
671
672                  try:
673                     if client:
674                        client.handleClose()
675                  except:
676                     pass
677
678                  try:
679                     del self.connections[ready]
680                  except:
681                     pass
682
683                  try:
684                     self.listeners.remove(ready)
685                  except:
686                     pass
687
688         for failed in xList:
689            if failed == self.serversocket:
690               self.close()
691               raise Exception("server socket failed")
692            else:
693               client = None
694               try:
695                   client = self.connections[failed]
696                   client.client.close()
697
698                   try:
699                      client.handleClose()
700                   except:
701                      pass
702
703                   try:
704                      self.listeners.remove(failed)
705                   except:
706                      pass
707
708               except:
709                  pass
710
711               finally:
712                  if client:
713                     del self.connections[failed]
714
715class SimpleSSLWebSocketServer(SimpleWebSocketServer):
716
717   def __init__(self, host, port, websocketclass, certfile,
718                keyfile, version = ssl.PROTOCOL_TLSv1, selectInterval = 0.1):
719
720      SimpleWebSocketServer.__init__(self, host, port,
721                                        websocketclass, selectInterval)
722
723      self.context = ssl.SSLContext(version)
724      self.context.load_cert_chain(certfile, keyfile)
725
726   def close(self):
727      super(SimpleSSLWebSocketServer, self).close()
728
729   def _decorateSocket(self, sock):
730      sslsock = self.context.wrap_socket(sock, server_side=True)
731      return sslsock
732
733   def _constructWebSocket(self, sock, address):
734      ws = self.websocketclass(self, sock, address)
735      ws.usingssl = True
736      return ws
737
738   def serveforever(self):
739      super(SimpleSSLWebSocketServer, self).serveforever()
Note: See TracBrowser for help on using the repository browser.