websocket.py 31 KB

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