websocket.py 31 KB

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