websocket.py 34 KB

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