websocket.py 28 KB

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