websocket.py 31 KB

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