websocket.py 33 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965
  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, f):
  207. pstart = f['hlen'] + 4
  208. pend = pstart + f['length']
  209. if numpy:
  210. b = c = s2b('')
  211. if f['length'] >= 4:
  212. mask = numpy.frombuffer(buf, dtype=numpy.dtype('<u4'),
  213. offset=f['hlen'], count=1)
  214. data = numpy.frombuffer(buf, dtype=numpy.dtype('<u4'),
  215. offset=pstart, count=int(f['length'] / 4))
  216. #b = numpy.bitwise_xor(data, mask).data
  217. b = numpy.bitwise_xor(data, mask).tostring()
  218. if f['length'] % 4:
  219. #print("Partial unmask")
  220. mask = numpy.frombuffer(buf, dtype=numpy.dtype('B'),
  221. offset=f['hlen'], count=(f['length'] % 4))
  222. data = numpy.frombuffer(buf, dtype=numpy.dtype('B'),
  223. offset=pend - (f['length'] % 4),
  224. count=(f['length'] % 4))
  225. c = numpy.bitwise_xor(data, mask).tostring()
  226. return b + c
  227. else:
  228. # Slower fallback
  229. data = array.array('B')
  230. mask = s2a(f['mask'])
  231. data.fromstring(buf[pstart:pend])
  232. for i in range(len(data)):
  233. data[i] ^= mask[i % 4]
  234. return data.tostring()
  235. @staticmethod
  236. def encode_hybi(buf, opcode, base64=False):
  237. """ Encode a HyBi style WebSocket frame.
  238. Optional opcode:
  239. 0x0 - continuation
  240. 0x1 - text frame (base64 encode buf)
  241. 0x2 - binary frame (use raw buf)
  242. 0x8 - connection close
  243. 0x9 - ping
  244. 0xA - pong
  245. """
  246. if base64:
  247. buf = b64encode(buf)
  248. b1 = 0x80 | (opcode & 0x0f) # FIN + opcode
  249. payload_len = len(buf)
  250. if payload_len <= 125:
  251. header = pack('>BB', b1, payload_len)
  252. elif payload_len > 125 and payload_len < 65536:
  253. header = pack('>BBH', b1, 126, payload_len)
  254. elif payload_len >= 65536:
  255. header = pack('>BBQ', b1, 127, payload_len)
  256. #print("Encoded: %s" % repr(header + buf))
  257. return header + buf, len(header), 0
  258. @staticmethod
  259. def decode_hybi(buf, base64=False):
  260. """ Decode HyBi style WebSocket packets.
  261. Returns:
  262. {'fin' : 0_or_1,
  263. 'opcode' : number,
  264. 'mask' : 32_bit_number,
  265. 'hlen' : header_bytes_number,
  266. 'length' : payload_bytes_number,
  267. 'payload' : decoded_buffer,
  268. 'left' : bytes_left_number,
  269. 'close_code' : number,
  270. 'close_reason' : string}
  271. """
  272. f = {'fin' : 0,
  273. 'opcode' : 0,
  274. 'mask' : 0,
  275. 'hlen' : 2,
  276. 'length' : 0,
  277. 'payload' : None,
  278. 'left' : 0,
  279. 'close_code' : 1000,
  280. 'close_reason' : ''}
  281. blen = len(buf)
  282. f['left'] = blen
  283. if blen < f['hlen']:
  284. return f # Incomplete frame header
  285. b1, b2 = unpack_from(">BB", buf)
  286. f['opcode'] = b1 & 0x0f
  287. f['fin'] = (b1 & 0x80) >> 7
  288. has_mask = (b2 & 0x80) >> 7
  289. f['length'] = b2 & 0x7f
  290. if f['length'] == 126:
  291. f['hlen'] = 4
  292. if blen < f['hlen']:
  293. return f # Incomplete frame header
  294. (f['length'],) = unpack_from('>xxH', buf)
  295. elif f['length'] == 127:
  296. f['hlen'] = 10
  297. if blen < f['hlen']:
  298. return f # Incomplete frame header
  299. (f['length'],) = unpack_from('>xxQ', buf)
  300. full_len = f['hlen'] + has_mask * 4 + f['length']
  301. if blen < full_len: # Incomplete frame
  302. return f # Incomplete frame header
  303. # Number of bytes that are part of the next frame(s)
  304. f['left'] = blen - full_len
  305. # Process 1 frame
  306. if has_mask:
  307. # unmask payload
  308. f['mask'] = buf[f['hlen']:f['hlen']+4]
  309. f['payload'] = WebSocketServer.unmask(buf, f)
  310. else:
  311. print("Unmasked frame: %s" % repr(buf))
  312. f['payload'] = buf[(f['hlen'] + has_mask * 4):full_len]
  313. if base64 and f['opcode'] in [1, 2]:
  314. try:
  315. f['payload'] = b64decode(f['payload'])
  316. except:
  317. print("Exception while b64decoding buffer: %s" %
  318. repr(buf))
  319. raise
  320. if f['opcode'] == 0x08:
  321. if f['length'] >= 2:
  322. f['close_code'] = unpack_from(">H", f['payload'])[0]
  323. if f['length'] > 3:
  324. f['close_reason'] = f['payload'][2:]
  325. return f
  326. @staticmethod
  327. def encode_hixie(buf):
  328. return s2b("\x00" + b2s(b64encode(buf)) + "\xff"), 1, 1
  329. @staticmethod
  330. def decode_hixie(buf):
  331. end = buf.find(s2b('\xff'))
  332. return {'payload': b64decode(buf[1:end]),
  333. 'hlen': 1,
  334. 'length': end - 1,
  335. 'left': len(buf) - (end + 1)}
  336. @staticmethod
  337. def gen_md5(keys):
  338. """ Generate hash value for WebSockets hixie-76. """
  339. key1 = keys['Sec-WebSocket-Key1']
  340. key2 = keys['Sec-WebSocket-Key2']
  341. key3 = keys['key3']
  342. spaces1 = key1.count(" ")
  343. spaces2 = key2.count(" ")
  344. num1 = int("".join([c for c in key1 if c.isdigit()])) / spaces1
  345. num2 = int("".join([c for c in key2 if c.isdigit()])) / spaces2
  346. return b2s(md5(pack('>II8s',
  347. int(num1), int(num2), key3)).digest())
  348. #
  349. # WebSocketServer logging/output functions
  350. #
  351. def traffic(self, token="."):
  352. """ Show traffic flow in verbose mode. """
  353. if self.verbose and not self.daemon:
  354. sys.stdout.write(token)
  355. sys.stdout.flush()
  356. def msg(self, msg):
  357. """ Output message with handler_id prefix. """
  358. if not self.daemon:
  359. print("% 3d: %s" % (self.handler_id, msg))
  360. def vmsg(self, msg):
  361. """ Same as msg() but only if verbose. """
  362. if self.verbose:
  363. self.msg(msg)
  364. #
  365. # Main WebSocketServer methods
  366. #
  367. def send_frames(self, bufs=None):
  368. """ Encode and send WebSocket frames. Any frames already
  369. queued will be sent first. If buf is not set then only queued
  370. frames will be sent. Returns the number of pending frames that
  371. could not be fully sent. If returned pending frames is greater
  372. than 0, then the caller should call again when the socket is
  373. ready. """
  374. tdelta = int(time.time()*1000) - self.start_time
  375. if bufs:
  376. for buf in bufs:
  377. if self.version.startswith("hybi"):
  378. if self.base64:
  379. encbuf, lenhead, lentail = self.encode_hybi(
  380. buf, opcode=1, base64=True)
  381. else:
  382. encbuf, lenhead, lentail = self.encode_hybi(
  383. buf, opcode=2, base64=False)
  384. else:
  385. encbuf, lenhead, lentail = self.encode_hixie(buf)
  386. if self.rec:
  387. self.rec.write("%s,\n" %
  388. repr("{%s{" % tdelta
  389. + encbuf[lenhead:-lentail]))
  390. self.send_parts.append(encbuf)
  391. while self.send_parts:
  392. # Send pending frames
  393. buf = self.send_parts.pop(0)
  394. sent = self.client.send(buf)
  395. if sent == len(buf):
  396. self.traffic("<")
  397. else:
  398. self.traffic("<.")
  399. self.send_parts.insert(0, buf[sent:])
  400. break
  401. return len(self.send_parts)
  402. def recv_frames(self):
  403. """ Receive and decode WebSocket frames.
  404. Returns:
  405. (bufs_list, closed_string)
  406. """
  407. closed = False
  408. bufs = []
  409. tdelta = int(time.time()*1000) - self.start_time
  410. buf = self.client.recv(self.buffer_size)
  411. if len(buf) == 0:
  412. closed = {'code': 1000, 'reason': "Client closed abruptly"}
  413. return bufs, closed
  414. if self.recv_part:
  415. # Add partially received frames to current read buffer
  416. buf = self.recv_part + buf
  417. self.recv_part = None
  418. while buf:
  419. if self.version.startswith("hybi"):
  420. frame = self.decode_hybi(buf, base64=self.base64)
  421. #print("Received buf: %s, frame: %s" % (repr(buf), frame))
  422. if frame['payload'] == None:
  423. # Incomplete/partial frame
  424. self.traffic("}.")
  425. if frame['left'] > 0:
  426. self.recv_part = buf[-frame['left']:]
  427. break
  428. else:
  429. if frame['opcode'] == 0x8: # connection close
  430. closed = {'code': frame['close_code'],
  431. 'reason': frame['close_reason']}
  432. break
  433. else:
  434. if buf[0:2] == s2b('\xff\x00'):
  435. closed = {'code': 1000,
  436. 'reason': "Client sent orderly close frame"}
  437. break
  438. elif buf[0:2] == s2b('\x00\xff'):
  439. buf = buf[2:]
  440. continue # No-op
  441. elif buf.count(s2b('\xff')) == 0:
  442. # Partial frame
  443. self.traffic("}.")
  444. self.recv_part = buf
  445. break
  446. frame = self.decode_hixie(buf)
  447. self.traffic("}")
  448. if self.rec:
  449. start = frame['hlen']
  450. end = frame['hlen'] + frame['length']
  451. self.rec.write("%s,\n" %
  452. repr("}%s}" % tdelta + buf[start:end]))
  453. bufs.append(frame['payload'])
  454. if frame['left']:
  455. buf = buf[-frame['left']:]
  456. else:
  457. buf = ''
  458. return bufs, closed
  459. def send_close(self, code=1000, reason=''):
  460. """ Send a WebSocket orderly close frame. """
  461. if self.version.startswith("hybi"):
  462. msg = pack(">H%ds" % len(reason), code, reason)
  463. buf, h, t = self.encode_hybi(msg, opcode=0x08, base64=False)
  464. self.client.send(buf)
  465. elif self.version == "hixie-76":
  466. buf = s2b('\xff\x00')
  467. self.client.send(buf)
  468. # No orderly close for 75
  469. def do_websocket_handshake(self, headers, path):
  470. h = self.headers = headers
  471. self.path = path
  472. prot = 'WebSocket-Protocol'
  473. protocols = h.get('Sec-'+prot, h.get(prot, '')).split(',')
  474. ver = h.get('Sec-WebSocket-Version')
  475. if ver:
  476. # HyBi/IETF version of the protocol
  477. # HyBi-07 report version 7
  478. # HyBi-08 - HyBi-12 report version 8
  479. # HyBi-13 reports version 13
  480. if ver in ['7', '8', '13']:
  481. self.version = "hybi-%02d" % int(ver)
  482. else:
  483. raise self.EClose('Unsupported protocol version %s' % ver)
  484. key = h['Sec-WebSocket-Key']
  485. # Choose binary if client supports it
  486. if 'binary' in protocols:
  487. self.base64 = False
  488. elif 'base64' in protocols:
  489. self.base64 = True
  490. else:
  491. raise self.EClose("Client must support 'binary' or 'base64' protocol")
  492. # Generate the hash value for the accept header
  493. accept = b64encode(sha1(s2b(key + self.GUID)).digest())
  494. response = self.server_handshake_hybi % b2s(accept)
  495. if self.base64:
  496. response += "Sec-WebSocket-Protocol: base64\r\n"
  497. else:
  498. response += "Sec-WebSocket-Protocol: binary\r\n"
  499. response += "\r\n"
  500. else:
  501. # Hixie version of the protocol (75 or 76)
  502. if h.get('key3'):
  503. trailer = self.gen_md5(h)
  504. pre = "Sec-"
  505. self.version = "hixie-76"
  506. else:
  507. trailer = ""
  508. pre = ""
  509. self.version = "hixie-75"
  510. # We only support base64 in Hixie era
  511. self.base64 = True
  512. response = self.server_handshake_hixie % (pre,
  513. h['Origin'], pre, self.scheme, h['Host'], path)
  514. if 'base64' in protocols:
  515. response += "%sWebSocket-Protocol: base64\r\n" % pre
  516. else:
  517. self.msg("Warning: client does not report 'base64' protocol support")
  518. response += "\r\n" + trailer
  519. return response
  520. def do_handshake(self, sock, address):
  521. """
  522. do_handshake does the following:
  523. - Peek at the first few bytes from the socket.
  524. - If the connection is Flash policy request then answer it,
  525. close the socket and return.
  526. - If the connection is an HTTPS/SSL/TLS connection then SSL
  527. wrap the socket.
  528. - Read from the (possibly wrapped) socket.
  529. - If we have received a HTTP GET request and the webserver
  530. functionality is enabled, answer it, close the socket and
  531. return.
  532. - Assume we have a WebSockets connection, parse the client
  533. handshake data.
  534. - Send a WebSockets handshake server response.
  535. - Return the socket for this WebSocket client.
  536. """
  537. stype = ""
  538. ready = select.select([sock], [], [], 3)[0]
  539. if not ready:
  540. raise self.EClose("ignoring socket not ready")
  541. # Peek, but do not read the data so that we have a opportunity
  542. # to SSL wrap the socket first
  543. handshake = sock.recv(1024, socket.MSG_PEEK)
  544. #self.msg("Handshake [%s]" % handshake)
  545. if handshake == "":
  546. raise self.EClose("ignoring empty handshake")
  547. elif handshake.startswith(s2b("<policy-file-request/>")):
  548. # Answer Flash policy request
  549. handshake = sock.recv(1024)
  550. sock.send(s2b(self.policy_response))
  551. raise self.EClose("Sending flash policy response")
  552. elif handshake[0] in ("\x16", "\x80", 22, 128):
  553. # SSL wrap the connection
  554. if not ssl:
  555. raise self.EClose("SSL connection but no 'ssl' module")
  556. if not os.path.exists(self.cert):
  557. raise self.EClose("SSL connection but '%s' not found"
  558. % self.cert)
  559. retsock = None
  560. try:
  561. retsock = ssl.wrap_socket(
  562. sock,
  563. server_side=True,
  564. certfile=self.cert,
  565. keyfile=self.key)
  566. except ssl.SSLError:
  567. _, x, _ = sys.exc_info()
  568. if x.args[0] == ssl.SSL_ERROR_EOF:
  569. if len(x.args) > 1:
  570. raise self.EClose(x.args[1])
  571. else:
  572. raise self.EClose("Got SSL_ERROR_EOF")
  573. else:
  574. raise
  575. self.scheme = "wss"
  576. stype = "SSL/TLS (wss://)"
  577. elif self.ssl_only:
  578. raise self.EClose("non-SSL connection received but disallowed")
  579. else:
  580. retsock = sock
  581. self.scheme = "ws"
  582. stype = "Plain non-SSL (ws://)"
  583. wsh = WSRequestHandler(retsock, address, not self.web)
  584. if wsh.last_code == 101:
  585. # Continue on to handle WebSocket upgrade
  586. pass
  587. elif wsh.last_code == 405:
  588. raise self.EClose("Normal web request received but disallowed")
  589. elif wsh.last_code < 200 or wsh.last_code >= 300:
  590. raise self.EClose(wsh.last_message)
  591. elif self.verbose:
  592. raise self.EClose(wsh.last_message)
  593. else:
  594. raise self.EClose("")
  595. response = self.do_websocket_handshake(wsh.headers, wsh.path)
  596. self.msg("%s: %s WebSocket connection" % (address[0], stype))
  597. self.msg("%s: Version %s, base64: '%s'" % (address[0],
  598. self.version, self.base64))
  599. if self.path != '/':
  600. self.msg("%s: Path: '%s'" % (address[0], self.path))
  601. # Send server WebSockets handshake response
  602. #self.msg("sending response [%s]" % response)
  603. retsock.send(s2b(response))
  604. # Return the WebSockets socket which may be SSL wrapped
  605. return retsock
  606. #
  607. # Events that can/should be overridden in sub-classes
  608. #
  609. def started(self):
  610. """ Called after WebSockets startup """
  611. self.vmsg("WebSockets server started")
  612. def poll(self):
  613. """ Run periodically while waiting for connections. """
  614. #self.vmsg("Running poll()")
  615. pass
  616. def fallback_SIGCHLD(self, sig, stack):
  617. # Reap zombies when using os.fork() (python 2.4)
  618. self.vmsg("Got SIGCHLD, reaping zombies")
  619. try:
  620. result = os.waitpid(-1, os.WNOHANG)
  621. while result[0]:
  622. self.vmsg("Reaped child process %s" % result[0])
  623. result = os.waitpid(-1, os.WNOHANG)
  624. except (OSError):
  625. pass
  626. def do_SIGINT(self, sig, stack):
  627. self.msg("Got SIGINT, exiting")
  628. sys.exit(0)
  629. def top_new_client(self, startsock, address):
  630. """ Do something with a WebSockets client connection. """
  631. # Initialize per client settings
  632. self.send_parts = []
  633. self.recv_part = None
  634. self.base64 = False
  635. self.rec = None
  636. self.start_time = int(time.time()*1000)
  637. # handler process
  638. try:
  639. try:
  640. self.client = self.do_handshake(startsock, address)
  641. if self.record:
  642. # Record raw frame data as JavaScript array
  643. fname = "%s.%s" % (self.record,
  644. self.handler_id)
  645. self.msg("opening record file: %s" % fname)
  646. self.rec = open(fname, 'w+')
  647. self.rec.write("var VNC_frame_data = [\n")
  648. self.ws_connection = True
  649. self.new_client()
  650. except self.CClose:
  651. # Close the client
  652. _, exc, _ = sys.exc_info()
  653. if self.client:
  654. self.send_close(exc.args[0], exc.args[1])
  655. except self.EClose:
  656. _, exc, _ = sys.exc_info()
  657. # Connection was not a WebSockets connection
  658. if exc.args[0]:
  659. self.msg("%s: %s" % (address[0], exc.args[0]))
  660. except Exception:
  661. _, exc, _ = sys.exc_info()
  662. self.msg("handler exception: %s" % str(exc))
  663. if self.verbose:
  664. self.msg(traceback.format_exc())
  665. finally:
  666. if self.rec:
  667. self.rec.write("'EOF']\n")
  668. self.rec.close()
  669. if self.client and self.client != startsock:
  670. # Close the SSL wrapped socket
  671. # Original socket closed by caller
  672. self.client.close()
  673. def new_client(self):
  674. """ Do something with a WebSockets client connection. """
  675. raise("WebSocketServer.new_client() must be overloaded")
  676. def start_server(self):
  677. """
  678. Daemonize if requested. Listen for for connections. Run
  679. do_handshake() method for each connection. If the connection
  680. is a WebSockets client then call new_client() method (which must
  681. be overridden) for each new client connection.
  682. """
  683. lsock = self.socket(self.listen_host, self.listen_port, False, self.prefer_ipv6)
  684. if self.daemon:
  685. self.daemonize(keepfd=lsock.fileno(), chdir=self.web)
  686. self.started() # Some things need to happen after daemonizing
  687. # Allow override of SIGINT
  688. signal.signal(signal.SIGINT, self.do_SIGINT)
  689. if not multiprocessing:
  690. # os.fork() (python 2.4) child reaper
  691. signal.signal(signal.SIGCHLD, self.fallback_SIGCHLD)
  692. last_active_time = self.launch_time
  693. while True:
  694. try:
  695. try:
  696. self.client = None
  697. startsock = None
  698. pid = err = 0
  699. child_count = 0
  700. if multiprocessing and self.idle_timeout:
  701. child_count = len(multiprocessing.active_children())
  702. time_elapsed = time.time() - self.launch_time
  703. if self.timeout and time_elapsed > self.timeout:
  704. self.msg('listener exit due to --timeout %s'
  705. % self.timeout)
  706. break
  707. if self.idle_timeout:
  708. idle_time = 0
  709. if child_count == 0:
  710. idle_time = time.time() - last_active_time
  711. else:
  712. idle_time = 0
  713. last_active_time = time.time()
  714. if idle_time > self.idle_timeout and child_count == 0:
  715. self.msg('listener exit due to --idle-timeout %s'
  716. % self.idle_timeout)
  717. break
  718. try:
  719. self.poll()
  720. ready = select.select([lsock], [], [], 1)[0]
  721. if lsock in ready:
  722. startsock, address = lsock.accept()
  723. else:
  724. continue
  725. except Exception:
  726. _, exc, _ = sys.exc_info()
  727. if hasattr(exc, 'errno'):
  728. err = exc.errno
  729. elif hasattr(exc, 'args'):
  730. err = exc.args[0]
  731. else:
  732. err = exc[0]
  733. if err == errno.EINTR:
  734. self.vmsg("Ignoring interrupted syscall")
  735. continue
  736. else:
  737. raise
  738. if self.run_once:
  739. # Run in same process if run_once
  740. self.top_new_client(startsock, address)
  741. if self.ws_connection :
  742. self.msg('%s: exiting due to --run-once'
  743. % address[0])
  744. break
  745. elif multiprocessing:
  746. self.vmsg('%s: new handler Process' % address[0])
  747. p = multiprocessing.Process(
  748. target=self.top_new_client,
  749. args=(startsock, address))
  750. p.start()
  751. # child will not return
  752. else:
  753. # python 2.4
  754. self.vmsg('%s: forking handler' % address[0])
  755. pid = os.fork()
  756. if pid == 0:
  757. # child handler process
  758. self.top_new_client(startsock, address)
  759. break # child process exits
  760. # parent process
  761. self.handler_id += 1
  762. except KeyboardInterrupt:
  763. _, exc, _ = sys.exc_info()
  764. print("In KeyboardInterrupt")
  765. pass
  766. except SystemExit:
  767. _, exc, _ = sys.exc_info()
  768. print("In SystemExit")
  769. break
  770. except Exception:
  771. _, exc, _ = sys.exc_info()
  772. self.msg("handler exception: %s" % str(exc))
  773. if self.verbose:
  774. self.msg(traceback.format_exc())
  775. finally:
  776. if startsock:
  777. startsock.close()
  778. # HTTP handler with WebSocket upgrade support
  779. class WSRequestHandler(SimpleHTTPRequestHandler):
  780. def __init__(self, req, addr, only_upgrade=False):
  781. self.only_upgrade = only_upgrade # only allow upgrades
  782. SimpleHTTPRequestHandler.__init__(self, req, addr, object())
  783. def do_GET(self):
  784. if (self.headers.get('upgrade') and
  785. self.headers.get('upgrade').lower() == 'websocket'):
  786. if (self.headers.get('sec-websocket-key1') or
  787. self.headers.get('websocket-key1')):
  788. # For Hixie-76 read out the key hash
  789. self.headers.__setitem__('key3', self.rfile.read(8))
  790. # Just indicate that an WebSocket upgrade is needed
  791. self.last_code = 101
  792. self.last_message = "101 Switching Protocols"
  793. elif self.only_upgrade:
  794. # Normal web request responses are disabled
  795. self.last_code = 405
  796. self.last_message = "405 Method Not Allowed"
  797. else:
  798. SimpleHTTPRequestHandler.do_GET(self)
  799. def send_response(self, code, message=None):
  800. # Save the status code
  801. self.last_code = code
  802. SimpleHTTPRequestHandler.send_response(self, code, message)
  803. def log_message(self, f, *args):
  804. # Save instead of printing
  805. self.last_message = f % args