websocket.py 29 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872
  1. #!/usr/bin/env python
  2. '''
  3. Python WebSocket library with support for "wss://" encryption.
  4. Copyright 2011 Joel Martin
  5. Licensed under LGPL version 3 (see docs/LICENSE.LGPL-3)
  6. Supports following protocol versions:
  7. - http://tools.ietf.org/html/draft-hixie-thewebsocketprotocol-75
  8. - http://tools.ietf.org/html/draft-hixie-thewebsocketprotocol-76
  9. - http://tools.ietf.org/html/draft-ietf-hybi-thewebsocketprotocol-10
  10. You can make a cert/key with openssl using:
  11. openssl req -new -x509 -days 365 -nodes -out self.pem -keyout self.pem
  12. as taken from http://docs.python.org/dev/library/ssl.html#certificates
  13. '''
  14. import os, sys, time, errno, signal, socket, struct, traceback, select
  15. from cgi import parse_qsl
  16. from base64 import b64encode, b64decode
  17. # Imports that vary by python version
  18. if sys.hexversion > 0x3000000:
  19. # python >= 3.0
  20. from io import StringIO
  21. from http.server import SimpleHTTPRequestHandler
  22. from urllib.parse import urlsplit
  23. b2s = lambda buf: buf.decode('latin_1')
  24. s2b = lambda s: s.encode('latin_1')
  25. else:
  26. # python 2.X
  27. from cStringIO import StringIO
  28. from SimpleHTTPServer import SimpleHTTPRequestHandler
  29. from urlparse import urlsplit
  30. # No-ops
  31. b2s = lambda buf: buf
  32. s2b = lambda s: s
  33. if sys.hexversion >= 0x2060000:
  34. # python >= 2.6
  35. from multiprocessing import Process
  36. from hashlib import md5, sha1
  37. else:
  38. # python < 2.6
  39. Process = None
  40. from md5 import md5
  41. from sha import sha as sha1
  42. # Degraded functionality if these imports are missing
  43. for mod, sup in [('numpy', 'HyBi protocol'),
  44. ('ssl', 'TLS/SSL/wss'), ('resource', 'daemonizing')]:
  45. try:
  46. globals()[mod] = __import__(mod)
  47. except ImportError:
  48. globals()[mod] = None
  49. print("WARNING: no '%s' module, %s support disabled" % (
  50. mod, sup))
  51. class WebSocketServer(object):
  52. """
  53. WebSockets server class.
  54. Must be sub-classed with new_client method definition.
  55. """
  56. buffer_size = 65536
  57. server_handshake_hixie = """HTTP/1.1 101 Web Socket Protocol Handshake\r
  58. Upgrade: WebSocket\r
  59. Connection: Upgrade\r
  60. %sWebSocket-Origin: %s\r
  61. %sWebSocket-Location: %s://%s%s\r
  62. """
  63. server_handshake_hybi = """HTTP/1.1 101 Switching Protocols\r
  64. Upgrade: websocket\r
  65. Connection: Upgrade\r
  66. Sec-WebSocket-Accept: %s\r
  67. """
  68. GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
  69. policy_response = """<cross-domain-policy><allow-access-from domain="*" to-ports="*" /></cross-domain-policy>\n"""
  70. class EClose(Exception):
  71. pass
  72. def __init__(self, listen_host='', listen_port=None, source_is_ipv6=False,
  73. verbose=False, cert='', key='', ssl_only=None,
  74. daemon=False, record='', web=''):
  75. # settings
  76. self.verbose = verbose
  77. self.listen_host = listen_host
  78. self.listen_port = listen_port
  79. self.ssl_only = ssl_only
  80. self.daemon = daemon
  81. self.handler_id = 1
  82. # Make paths settings absolute
  83. self.cert = os.path.abspath(cert)
  84. self.key = self.web = self.record = ''
  85. if key:
  86. self.key = os.path.abspath(key)
  87. if web:
  88. self.web = os.path.abspath(web)
  89. if record:
  90. self.record = os.path.abspath(record)
  91. if self.web:
  92. os.chdir(self.web)
  93. # Sanity checks
  94. if not ssl and self.ssl_only:
  95. raise Exception("No 'ssl' module and SSL-only specified")
  96. if self.daemon and not resource:
  97. raise Exception("Module 'resource' required to daemonize")
  98. # Show configuration
  99. print("WebSocket server settings:")
  100. print(" - Listen on %s:%s" % (
  101. self.listen_host, self.listen_port))
  102. print(" - Flash security policy server")
  103. if self.web:
  104. print(" - Web server. Web root: %s" % self.web)
  105. if ssl:
  106. if os.path.exists(self.cert):
  107. print(" - SSL/TLS support")
  108. if self.ssl_only:
  109. print(" - Deny non-SSL/TLS connections")
  110. else:
  111. print(" - No SSL/TLS support (no cert file)")
  112. else:
  113. print(" - No SSL/TLS support (no 'ssl' module)")
  114. if self.daemon:
  115. print(" - Backgrounding (daemon)")
  116. if self.record:
  117. print(" - Recording to '%s.*'" % self.record)
  118. #
  119. # WebSocketServer static methods
  120. #
  121. @staticmethod
  122. def socket(host, port=None, connect=False, prefer_ipv6=False):
  123. """ Resolve a host (and optional port) to an IPv4 or IPv6
  124. address. Create a socket. Bind to it if listen is set,
  125. otherwise connect to it. Return the socket.
  126. """
  127. flags = 0
  128. if host == '':
  129. host = None
  130. if connect and not port:
  131. raise Exception("Connect mode requires a port")
  132. if not connect:
  133. flags = flags | socket.AI_PASSIVE
  134. addrs = socket.getaddrinfo(host, port, 0, socket.SOCK_STREAM,
  135. socket.IPPROTO_TCP, flags)
  136. if not addrs:
  137. raise Exception("Could resolve host '%s'" % host)
  138. addrs.sort(key=lambda x: x[0])
  139. if prefer_ipv6:
  140. addrs.reverse()
  141. sock = socket.socket(addrs[0][0], addrs[0][1])
  142. if connect:
  143. sock.connect(addrs[0][4])
  144. else:
  145. sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
  146. sock.bind(addrs[0][4])
  147. sock.listen(100)
  148. return sock
  149. @staticmethod
  150. def daemonize(keepfd=None, chdir='/'):
  151. os.umask(0)
  152. if chdir:
  153. os.chdir(chdir)
  154. else:
  155. os.chdir('/')
  156. os.setgid(os.getgid()) # relinquish elevations
  157. os.setuid(os.getuid()) # relinquish elevations
  158. # Double fork to daemonize
  159. if os.fork() > 0: os._exit(0) # Parent exits
  160. os.setsid() # Obtain new process group
  161. if os.fork() > 0: os._exit(0) # Parent exits
  162. # Signal handling
  163. def terminate(a,b): os._exit(0)
  164. signal.signal(signal.SIGTERM, terminate)
  165. signal.signal(signal.SIGINT, signal.SIG_IGN)
  166. # Close open files
  167. maxfd = resource.getrlimit(resource.RLIMIT_NOFILE)[1]
  168. if maxfd == resource.RLIM_INFINITY: maxfd = 256
  169. for fd in reversed(range(maxfd)):
  170. try:
  171. if fd != keepfd:
  172. os.close(fd)
  173. except OSError:
  174. _, exc, _ = sys.exc_info()
  175. if exc.errno != errno.EBADF: raise
  176. # Redirect I/O to /dev/null
  177. os.dup2(os.open(os.devnull, os.O_RDWR), sys.stdin.fileno())
  178. os.dup2(os.open(os.devnull, os.O_RDWR), sys.stdout.fileno())
  179. os.dup2(os.open(os.devnull, os.O_RDWR), sys.stderr.fileno())
  180. @staticmethod
  181. def encode_hybi(buf, opcode, base64=False):
  182. """ Encode a HyBi style WebSocket frame.
  183. Optional opcode:
  184. 0x0 - continuation
  185. 0x1 - text frame (base64 encode buf)
  186. 0x2 - binary frame (use raw buf)
  187. 0x8 - connection close
  188. 0x9 - ping
  189. 0xA - pong
  190. """
  191. if base64:
  192. buf = b64encode(buf)
  193. b1 = 0x80 | (opcode & 0x0f) # FIN + opcode
  194. payload_len = len(buf)
  195. if payload_len <= 125:
  196. header = struct.pack('>BB', b1, payload_len)
  197. elif payload_len > 125 and payload_len < 65536:
  198. header = struct.pack('>BBH', b1, 126, payload_len)
  199. elif payload_len >= 65536:
  200. header = struct.pack('>BBQ', b1, 127, payload_len)
  201. #print("Encoded: %s" % repr(header + buf))
  202. return header + buf, len(header), 0
  203. @staticmethod
  204. def decode_hybi(buf, base64=False):
  205. """ Decode HyBi style WebSocket packets.
  206. Returns:
  207. {'fin' : 0_or_1,
  208. 'opcode' : number,
  209. 'mask' : 32_bit_number,
  210. 'hlen' : header_bytes_number,
  211. 'length' : payload_bytes_number,
  212. 'payload' : decoded_buffer,
  213. 'left' : bytes_left_number,
  214. 'close_code' : number,
  215. 'close_reason' : string}
  216. """
  217. f = {'fin' : 0,
  218. 'opcode' : 0,
  219. 'mask' : 0,
  220. 'hlen' : 2,
  221. 'length' : 0,
  222. 'payload' : None,
  223. 'left' : 0,
  224. 'close_code' : None,
  225. 'close_reason' : None}
  226. blen = len(buf)
  227. f['left'] = blen
  228. if blen < f['hlen']:
  229. return f # Incomplete frame header
  230. b1, b2 = struct.unpack_from(">BB", buf)
  231. f['opcode'] = b1 & 0x0f
  232. f['fin'] = (b1 & 0x80) >> 7
  233. has_mask = (b2 & 0x80) >> 7
  234. f['length'] = b2 & 0x7f
  235. if f['length'] == 126:
  236. f['hlen'] = 4
  237. if blen < f['hlen']:
  238. return f # Incomplete frame header
  239. (f['length'],) = struct.unpack_from('>xxH', buf)
  240. elif f['length'] == 127:
  241. f['hlen'] = 10
  242. if blen < f['hlen']:
  243. return f # Incomplete frame header
  244. (f['length'],) = struct.unpack_from('>xxQ', buf)
  245. full_len = f['hlen'] + has_mask * 4 + f['length']
  246. if blen < full_len: # Incomplete frame
  247. return f # Incomplete frame header
  248. # Number of bytes that are part of the next frame(s)
  249. f['left'] = blen - full_len
  250. # Process 1 frame
  251. if has_mask:
  252. # unmask payload
  253. f['mask'] = buf[f['hlen']:f['hlen']+4]
  254. b = c = s2b('')
  255. if f['length'] >= 4:
  256. mask = numpy.frombuffer(buf, dtype=numpy.dtype('<u4'),
  257. offset=f['hlen'], count=1)
  258. data = numpy.frombuffer(buf, dtype=numpy.dtype('<u4'),
  259. offset=f['hlen'] + 4, count=int(f['length'] / 4))
  260. #b = numpy.bitwise_xor(data, mask).data
  261. b = numpy.bitwise_xor(data, mask).tostring()
  262. if f['length'] % 4:
  263. #print("Partial unmask")
  264. mask = numpy.frombuffer(buf, dtype=numpy.dtype('B'),
  265. offset=f['hlen'], count=(f['length'] % 4))
  266. data = numpy.frombuffer(buf, dtype=numpy.dtype('B'),
  267. offset=full_len - (f['length'] % 4),
  268. count=(f['length'] % 4))
  269. c = numpy.bitwise_xor(data, mask).tostring()
  270. f['payload'] = b + c
  271. else:
  272. print("Unmasked frame: %s" % repr(buf))
  273. f['payload'] = buf[(f['hlen'] + has_mask * 4):full_len]
  274. if base64 and f['opcode'] in [1, 2]:
  275. try:
  276. f['payload'] = b64decode(f['payload'])
  277. except:
  278. print("Exception while b64decoding buffer: %s" %
  279. repr(buf))
  280. raise
  281. if f['opcode'] == 0x08:
  282. if f['length'] >= 2:
  283. f['close_code'] = struct.unpack_from(">H", f['payload'])
  284. if f['length'] > 3:
  285. f['close_reason'] = f['payload'][2:]
  286. return f
  287. @staticmethod
  288. def encode_hixie(buf):
  289. return s2b("\x00" + b2s(b64encode(buf)) + "\xff"), 1, 1
  290. @staticmethod
  291. def decode_hixie(buf):
  292. end = buf.find(s2b('\xff'))
  293. return {'payload': b64decode(buf[1:end]),
  294. 'hlen': 1,
  295. 'length': end - 1,
  296. 'left': len(buf) - (end + 1)}
  297. @staticmethod
  298. def gen_md5(keys):
  299. """ Generate hash value for WebSockets hixie-76. """
  300. key1 = keys['Sec-WebSocket-Key1']
  301. key2 = keys['Sec-WebSocket-Key2']
  302. key3 = keys['key3']
  303. spaces1 = key1.count(" ")
  304. spaces2 = key2.count(" ")
  305. num1 = int("".join([c for c in key1 if c.isdigit()])) / spaces1
  306. num2 = int("".join([c for c in key2 if c.isdigit()])) / spaces2
  307. return b2s(md5(struct.pack('>II8s',
  308. int(num1), int(num2), key3)).digest())
  309. #
  310. # WebSocketServer logging/output functions
  311. #
  312. def traffic(self, token="."):
  313. """ Show traffic flow in verbose mode. """
  314. if self.verbose and not self.daemon:
  315. sys.stdout.write(token)
  316. sys.stdout.flush()
  317. def msg(self, msg):
  318. """ Output message with handler_id prefix. """
  319. if not self.daemon:
  320. print("% 3d: %s" % (self.handler_id, msg))
  321. def vmsg(self, msg):
  322. """ Same as msg() but only if verbose. """
  323. if self.verbose:
  324. self.msg(msg)
  325. #
  326. # Main WebSocketServer methods
  327. #
  328. def send_frames(self, bufs=None):
  329. """ Encode and send WebSocket frames. Any frames already
  330. queued will be sent first. If buf is not set then only queued
  331. frames will be sent. Returns the number of pending frames that
  332. could not be fully sent. If returned pending frames is greater
  333. than 0, then the caller should call again when the socket is
  334. ready. """
  335. tdelta = int(time.time()*1000) - self.start_time
  336. if bufs:
  337. for buf in bufs:
  338. if self.version.startswith("hybi"):
  339. if self.base64:
  340. encbuf, lenhead, lentail = self.encode_hybi(
  341. buf, opcode=1, base64=True)
  342. else:
  343. encbuf, lenhead, lentail = self.encode_hybi(
  344. buf, opcode=2, base64=False)
  345. else:
  346. encbuf, lenhead, lentail = self.encode_hixie(buf)
  347. if self.rec:
  348. self.rec.write("%s,\n" %
  349. repr("{%s{" % tdelta
  350. + encbuf[lenhead:-lentail]))
  351. self.send_parts.append(encbuf)
  352. while self.send_parts:
  353. # Send pending frames
  354. buf = self.send_parts.pop(0)
  355. sent = self.client.send(buf)
  356. if sent == len(buf):
  357. self.traffic("<")
  358. else:
  359. self.traffic("<.")
  360. self.send_parts.insert(0, buf[sent:])
  361. break
  362. return len(self.send_parts)
  363. def recv_frames(self):
  364. """ Receive and decode WebSocket frames.
  365. Returns:
  366. (bufs_list, closed_string)
  367. """
  368. closed = False
  369. bufs = []
  370. tdelta = int(time.time()*1000) - self.start_time
  371. buf = self.client.recv(self.buffer_size)
  372. if len(buf) == 0:
  373. closed = "Client closed abruptly"
  374. return bufs, closed
  375. if self.recv_part:
  376. # Add partially received frames to current read buffer
  377. buf = self.recv_part + buf
  378. self.recv_part = None
  379. while buf:
  380. if self.version.startswith("hybi"):
  381. frame = self.decode_hybi(buf, base64=self.base64)
  382. #print("Received buf: %s, frame: %s" % (repr(buf), frame))
  383. if frame['payload'] == None:
  384. # Incomplete/partial frame
  385. self.traffic("}.")
  386. if frame['left'] > 0:
  387. self.recv_part = buf[-frame['left']:]
  388. break
  389. else:
  390. if frame['opcode'] == 0x8: # connection close
  391. closed = "Client closed, reason: %s - %s" % (
  392. frame['close_code'],
  393. frame['close_reason'])
  394. break
  395. else:
  396. if buf[0:2] == '\xff\x00':
  397. closed = "Client sent orderly close frame"
  398. break
  399. elif buf[0:2] == '\x00\xff':
  400. buf = buf[2:]
  401. continue # No-op
  402. elif buf.count(s2b('\xff')) == 0:
  403. # Partial frame
  404. self.traffic("}.")
  405. self.recv_part = buf
  406. break
  407. frame = self.decode_hixie(buf)
  408. self.traffic("}")
  409. if self.rec:
  410. start = frame['hlen']
  411. end = frame['hlen'] + frame['length']
  412. self.rec.write("%s,\n" %
  413. repr("}%s}" % tdelta + buf[start:end]))
  414. bufs.append(frame['payload'])
  415. if frame['left']:
  416. buf = buf[-frame['left']:]
  417. else:
  418. buf = ''
  419. return bufs, closed
  420. def send_close(self, code=None, reason=''):
  421. """ Send a WebSocket orderly close frame. """
  422. if self.version.startswith("hybi"):
  423. msg = s2b('')
  424. if code != None:
  425. msg = struct.pack(">H%ds" % (len(reason)), code)
  426. buf, h, t = self.encode_hybi(msg, opcode=0x08, base64=False)
  427. self.client.send(buf)
  428. elif self.version == "hixie-76":
  429. buf = s2b('\xff\x00')
  430. self.client.send(buf)
  431. # No orderly close for 75
  432. def do_handshake(self, sock, address):
  433. """
  434. do_handshake does the following:
  435. - Peek at the first few bytes from the socket.
  436. - If the connection is Flash policy request then answer it,
  437. close the socket and return.
  438. - If the connection is an HTTPS/SSL/TLS connection then SSL
  439. wrap the socket.
  440. - Read from the (possibly wrapped) socket.
  441. - If we have received a HTTP GET request and the webserver
  442. functionality is enabled, answer it, close the socket and
  443. return.
  444. - Assume we have a WebSockets connection, parse the client
  445. handshake data.
  446. - Send a WebSockets handshake server response.
  447. - Return the socket for this WebSocket client.
  448. """
  449. stype = ""
  450. ready = select.select([sock], [], [], 3)[0]
  451. if not ready:
  452. raise self.EClose("ignoring socket not ready")
  453. # Peek, but do not read the data so that we have a opportunity
  454. # to SSL wrap the socket first
  455. handshake = sock.recv(1024, socket.MSG_PEEK)
  456. #self.msg("Handshake [%s]" % handshake)
  457. if handshake == "":
  458. raise self.EClose("ignoring empty handshake")
  459. elif handshake.startswith(s2b("<policy-file-request/>")):
  460. # Answer Flash policy request
  461. handshake = sock.recv(1024)
  462. sock.send(s2b(self.policy_response))
  463. raise self.EClose("Sending flash policy response")
  464. elif handshake[0] in ("\x16", "\x80"):
  465. # SSL wrap the connection
  466. if not ssl:
  467. raise self.EClose("SSL connection but no 'ssl' module")
  468. if not os.path.exists(self.cert):
  469. raise self.EClose("SSL connection but '%s' not found"
  470. % self.cert)
  471. retsock = None
  472. try:
  473. retsock = ssl.wrap_socket(
  474. sock,
  475. server_side=True,
  476. certfile=self.cert,
  477. keyfile=self.key)
  478. except ssl.SSLError:
  479. _, x, _ = sys.exc_info()
  480. if x.args[0] == ssl.SSL_ERROR_EOF:
  481. raise self.EClose("")
  482. else:
  483. raise
  484. scheme = "wss"
  485. stype = "SSL/TLS (wss://)"
  486. elif self.ssl_only:
  487. raise self.EClose("non-SSL connection received but disallowed")
  488. else:
  489. retsock = sock
  490. scheme = "ws"
  491. stype = "Plain non-SSL (ws://)"
  492. wsh = WSRequestHandler(retsock, address, not self.web)
  493. if wsh.last_code == 101:
  494. # Continue on to handle WebSocket upgrade
  495. pass
  496. elif wsh.last_code == 405:
  497. raise self.EClose("Normal web request received but disallowed")
  498. elif wsh.last_code < 200 or wsh.last_code >= 300:
  499. raise self.EClose(wsh.last_message)
  500. elif self.verbose:
  501. raise self.EClose(wsh.last_message)
  502. else:
  503. raise self.EClose("")
  504. h = self.headers = wsh.headers
  505. path = self.path = wsh.path
  506. prot = 'WebSocket-Protocol'
  507. protocols = h.get('Sec-'+prot, h.get(prot, '')).split(',')
  508. ver = h.get('Sec-WebSocket-Version')
  509. if ver:
  510. # HyBi/IETF version of the protocol
  511. if sys.hexversion < 0x2060000 or not numpy:
  512. raise self.EClose("Python >= 2.6 and numpy module is required for HyBi-07 or greater")
  513. # HyBi-07 report version 7
  514. # HyBi-08 - HyBi-12 report version 8
  515. # HyBi-13 reports version 13
  516. if ver in ['7', '8', '13']:
  517. self.version = "hybi-%02d" % int(ver)
  518. else:
  519. raise self.EClose('Unsupported protocol version %s' % ver)
  520. key = h['Sec-WebSocket-Key']
  521. # Choose binary if client supports it
  522. if 'binary' in protocols:
  523. self.base64 = False
  524. elif 'base64' in protocols:
  525. self.base64 = True
  526. else:
  527. raise self.EClose("Client must support 'binary' or 'base64' protocol")
  528. # Generate the hash value for the accept header
  529. accept = b64encode(sha1(s2b(key + self.GUID)).digest())
  530. response = self.server_handshake_hybi % b2s(accept)
  531. if self.base64:
  532. response += "Sec-WebSocket-Protocol: base64\r\n"
  533. else:
  534. response += "Sec-WebSocket-Protocol: binary\r\n"
  535. response += "\r\n"
  536. else:
  537. # Hixie version of the protocol (75 or 76)
  538. if h.get('key3'):
  539. trailer = self.gen_md5(h)
  540. pre = "Sec-"
  541. self.version = "hixie-76"
  542. else:
  543. trailer = ""
  544. pre = ""
  545. self.version = "hixie-75"
  546. # We only support base64 in Hixie era
  547. self.base64 = True
  548. response = self.server_handshake_hixie % (pre,
  549. h['Origin'], pre, scheme, h['Host'], path)
  550. if 'base64' in protocols:
  551. response += "%sWebSocket-Protocol: base64\r\n" % pre
  552. else:
  553. self.msg("Warning: client does not report 'base64' protocol support")
  554. response += "\r\n" + trailer
  555. self.msg("%s: %s WebSocket connection" % (address[0], stype))
  556. self.msg("%s: Version %s, base64: '%s'" % (address[0],
  557. self.version, self.base64))
  558. # Send server WebSockets handshake response
  559. #self.msg("sending response [%s]" % response)
  560. retsock.send(s2b(response))
  561. # Return the WebSockets socket which may be SSL wrapped
  562. return retsock
  563. #
  564. # Events that can/should be overridden in sub-classes
  565. #
  566. def started(self):
  567. """ Called after WebSockets startup """
  568. self.vmsg("WebSockets server started")
  569. def poll(self):
  570. """ Run periodically while waiting for connections. """
  571. #self.vmsg("Running poll()")
  572. pass
  573. def fallback_SIGCHLD(self, sig, stack):
  574. # Reap zombies when using os.fork() (python 2.4)
  575. self.vmsg("Got SIGCHLD, reaping zombies")
  576. try:
  577. result = os.waitpid(-1, os.WNOHANG)
  578. while result[0]:
  579. self.vmsg("Reaped child process %s" % result[0])
  580. result = os.waitpid(-1, os.WNOHANG)
  581. except (OSError):
  582. pass
  583. def do_SIGINT(self, sig, stack):
  584. self.msg("Got SIGINT, exiting")
  585. sys.exit(0)
  586. def top_new_client(self, startsock, address):
  587. """ Do something with a WebSockets client connection. """
  588. # Initialize per client settings
  589. self.send_parts = []
  590. self.recv_part = None
  591. self.base64 = False
  592. self.rec = None
  593. self.start_time = int(time.time()*1000)
  594. # handler process
  595. try:
  596. try:
  597. self.client = self.do_handshake(startsock, address)
  598. if self.record:
  599. # Record raw frame data as JavaScript array
  600. fname = "%s.%s" % (self.record,
  601. self.handler_id)
  602. self.msg("opening record file: %s" % fname)
  603. self.rec = open(fname, 'w+')
  604. self.rec.write("var VNC_frame_data = [\n")
  605. self.new_client()
  606. except self.EClose:
  607. _, exc, _ = sys.exc_info()
  608. # Connection was not a WebSockets connection
  609. if exc.args[0]:
  610. self.msg("%s: %s" % (address[0], exc.args[0]))
  611. except Exception:
  612. _, exc, _ = sys.exc_info()
  613. self.msg("handler exception: %s" % str(exc))
  614. if self.verbose:
  615. self.msg(traceback.format_exc())
  616. finally:
  617. if self.rec:
  618. self.rec.write("'EOF']\n")
  619. self.rec.close()
  620. if self.client and self.client != startsock:
  621. self.client.close()
  622. def new_client(self):
  623. """ Do something with a WebSockets client connection. """
  624. raise("WebSocketServer.new_client() must be overloaded")
  625. def start_server(self):
  626. """
  627. Daemonize if requested. Listen for for connections. Run
  628. do_handshake() method for each connection. If the connection
  629. is a WebSockets client then call new_client() method (which must
  630. be overridden) for each new client connection.
  631. """
  632. lsock = self.socket(self.listen_host, self.listen_port)
  633. if self.daemon:
  634. self.daemonize(keepfd=lsock.fileno(), chdir=self.web)
  635. self.started() # Some things need to happen after daemonizing
  636. # Allow override of SIGINT
  637. signal.signal(signal.SIGINT, self.do_SIGINT)
  638. if not Process:
  639. # os.fork() (python 2.4) child reaper
  640. signal.signal(signal.SIGCHLD, self.fallback_SIGCHLD)
  641. while True:
  642. try:
  643. try:
  644. self.client = None
  645. startsock = None
  646. pid = err = 0
  647. try:
  648. self.poll()
  649. ready = select.select([lsock], [], [], 1)[0]
  650. if lsock in ready:
  651. startsock, address = lsock.accept()
  652. else:
  653. continue
  654. except Exception:
  655. _, exc, _ = sys.exc_info()
  656. if hasattr(exc, 'errno'):
  657. err = exc.errno
  658. elif hasattr(exc, 'args'):
  659. err = exc.args[0]
  660. else:
  661. err = exc[0]
  662. if err == errno.EINTR:
  663. self.vmsg("Ignoring interrupted syscall")
  664. continue
  665. else:
  666. raise
  667. if Process:
  668. self.vmsg('%s: new handler Process' % address[0])
  669. p = Process(target=self.top_new_client,
  670. args=(startsock, address))
  671. p.start()
  672. # child will not return
  673. else:
  674. # python 2.4
  675. self.vmsg('%s: forking handler' % address[0])
  676. pid = os.fork()
  677. if pid == 0:
  678. # child handler process
  679. self.top_new_client(startsock, address)
  680. break # child process exits
  681. # parent process
  682. self.handler_id += 1
  683. except KeyboardInterrupt:
  684. _, exc, _ = sys.exc_info()
  685. print("In KeyboardInterrupt")
  686. pass
  687. except SystemExit:
  688. _, exc, _ = sys.exc_info()
  689. print("In SystemExit")
  690. break
  691. except Exception:
  692. _, exc, _ = sys.exc_info()
  693. self.msg("handler exception: %s" % str(exc))
  694. if self.verbose:
  695. self.msg(traceback.format_exc())
  696. finally:
  697. if startsock:
  698. startsock.close()
  699. # HTTP handler with WebSocket upgrade support
  700. class WSRequestHandler(SimpleHTTPRequestHandler):
  701. def __init__(self, req, addr, only_upgrade=False):
  702. self.only_upgrade = only_upgrade # only allow upgrades
  703. SimpleHTTPRequestHandler.__init__(self, req, addr, object())
  704. def do_GET(self):
  705. if (self.headers.get('upgrade') and
  706. self.headers.get('upgrade').lower() == 'websocket'):
  707. if (self.headers.get('sec-websocket-key1') or
  708. self.headers.get('websocket-key1')):
  709. # For Hixie-76 read out the key hash
  710. self.headers.__setitem__('key3', self.rfile.read(8))
  711. # Just indicate that an WebSocket upgrade is needed
  712. self.last_code = 101
  713. self.last_message = "101 Switching Protocols"
  714. elif self.only_upgrade:
  715. # Normal web request responses are disabled
  716. self.last_code = 405
  717. self.last_message = "405 Method Not Allowed"
  718. else:
  719. SimpleHTTPRequestHandler.do_GET(self)
  720. def send_response(self, code, message=None):
  721. # Save the status code
  722. self.last_code = code
  723. SimpleHTTPRequestHandler.send_response(self, code, message)
  724. def log_message(self, f, *args):
  725. # Save instead of printing
  726. self.last_message = f % args