websocket.py 32 KB

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