_abnf.py 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447
  1. """
  2. websocket - WebSocket client library for Python
  3. Copyright (C) 2010 Hiroki Ohtani(liris)
  4. This library is free software; you can redistribute it and/or
  5. modify it under the terms of the GNU Lesser General Public
  6. License as published by the Free Software Foundation; either
  7. version 2.1 of the License, or (at your option) any later version.
  8. This library is distributed in the hope that it will be useful,
  9. but WITHOUT ANY WARRANTY; without even the implied warranty of
  10. MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
  11. Lesser General Public License for more details.
  12. You should have received a copy of the GNU Lesser General Public
  13. License along with this library; if not, write to the Free Software
  14. Foundation, Inc., 51 Franklin Street, Fifth Floor,
  15. Boston, MA 02110-1335 USA
  16. """
  17. import array
  18. import os
  19. import struct
  20. import six
  21. from ._exceptions import *
  22. from ._utils import validate_utf8
  23. from threading import Lock
  24. try:
  25. if six.PY3:
  26. import numpy
  27. else:
  28. numpy = None
  29. except ImportError:
  30. numpy = None
  31. try:
  32. # If wsaccel is available we use compiled routines to mask data.
  33. if not numpy:
  34. from wsaccel.xormask import XorMaskerSimple
  35. def _mask(_m, _d):
  36. return XorMaskerSimple(_m).process(_d)
  37. except ImportError:
  38. # wsaccel is not available, we rely on python implementations.
  39. def _mask(_m, _d):
  40. for i in range(len(_d)):
  41. _d[i] ^= _m[i % 4]
  42. if six.PY3:
  43. return _d.tobytes()
  44. else:
  45. return _d.tostring()
  46. __all__ = [
  47. 'ABNF', 'continuous_frame', 'frame_buffer',
  48. 'STATUS_NORMAL',
  49. 'STATUS_GOING_AWAY',
  50. 'STATUS_PROTOCOL_ERROR',
  51. 'STATUS_UNSUPPORTED_DATA_TYPE',
  52. 'STATUS_STATUS_NOT_AVAILABLE',
  53. 'STATUS_ABNORMAL_CLOSED',
  54. 'STATUS_INVALID_PAYLOAD',
  55. 'STATUS_POLICY_VIOLATION',
  56. 'STATUS_MESSAGE_TOO_BIG',
  57. 'STATUS_INVALID_EXTENSION',
  58. 'STATUS_UNEXPECTED_CONDITION',
  59. 'STATUS_BAD_GATEWAY',
  60. 'STATUS_TLS_HANDSHAKE_ERROR',
  61. ]
  62. # closing frame status codes.
  63. STATUS_NORMAL = 1000
  64. STATUS_GOING_AWAY = 1001
  65. STATUS_PROTOCOL_ERROR = 1002
  66. STATUS_UNSUPPORTED_DATA_TYPE = 1003
  67. STATUS_STATUS_NOT_AVAILABLE = 1005
  68. STATUS_ABNORMAL_CLOSED = 1006
  69. STATUS_INVALID_PAYLOAD = 1007
  70. STATUS_POLICY_VIOLATION = 1008
  71. STATUS_MESSAGE_TOO_BIG = 1009
  72. STATUS_INVALID_EXTENSION = 1010
  73. STATUS_UNEXPECTED_CONDITION = 1011
  74. STATUS_BAD_GATEWAY = 1014
  75. STATUS_TLS_HANDSHAKE_ERROR = 1015
  76. VALID_CLOSE_STATUS = (
  77. STATUS_NORMAL,
  78. STATUS_GOING_AWAY,
  79. STATUS_PROTOCOL_ERROR,
  80. STATUS_UNSUPPORTED_DATA_TYPE,
  81. STATUS_INVALID_PAYLOAD,
  82. STATUS_POLICY_VIOLATION,
  83. STATUS_MESSAGE_TOO_BIG,
  84. STATUS_INVALID_EXTENSION,
  85. STATUS_UNEXPECTED_CONDITION,
  86. STATUS_BAD_GATEWAY,
  87. )
  88. class ABNF(object):
  89. """
  90. ABNF frame class.
  91. see http://tools.ietf.org/html/rfc5234
  92. and http://tools.ietf.org/html/rfc6455#section-5.2
  93. """
  94. # operation code values.
  95. OPCODE_CONT = 0x0
  96. OPCODE_TEXT = 0x1
  97. OPCODE_BINARY = 0x2
  98. OPCODE_CLOSE = 0x8
  99. OPCODE_PING = 0x9
  100. OPCODE_PONG = 0xa
  101. # available operation code value tuple
  102. OPCODES = (OPCODE_CONT, OPCODE_TEXT, OPCODE_BINARY, OPCODE_CLOSE,
  103. OPCODE_PING, OPCODE_PONG)
  104. # opcode human readable string
  105. OPCODE_MAP = {
  106. OPCODE_CONT: "cont",
  107. OPCODE_TEXT: "text",
  108. OPCODE_BINARY: "binary",
  109. OPCODE_CLOSE: "close",
  110. OPCODE_PING: "ping",
  111. OPCODE_PONG: "pong"
  112. }
  113. # data length threshold.
  114. LENGTH_7 = 0x7e
  115. LENGTH_16 = 1 << 16
  116. LENGTH_63 = 1 << 63
  117. def __init__(self, fin=0, rsv1=0, rsv2=0, rsv3=0,
  118. opcode=OPCODE_TEXT, mask=1, data=""):
  119. """
  120. Constructor for ABNF.
  121. please check RFC for arguments.
  122. """
  123. self.fin = fin
  124. self.rsv1 = rsv1
  125. self.rsv2 = rsv2
  126. self.rsv3 = rsv3
  127. self.opcode = opcode
  128. self.mask = mask
  129. if data is None:
  130. data = ""
  131. self.data = data
  132. self.get_mask_key = os.urandom
  133. def validate(self, skip_utf8_validation=False):
  134. """
  135. validate the ABNF frame.
  136. skip_utf8_validation: skip utf8 validation.
  137. """
  138. if self.rsv1 or self.rsv2 or self.rsv3:
  139. raise WebSocketProtocolException("rsv is not implemented, yet")
  140. if self.opcode not in ABNF.OPCODES:
  141. raise WebSocketProtocolException("Invalid opcode %r", self.opcode)
  142. if self.opcode == ABNF.OPCODE_PING and not self.fin:
  143. raise WebSocketProtocolException("Invalid ping frame.")
  144. if self.opcode == ABNF.OPCODE_CLOSE:
  145. l = len(self.data)
  146. if not l:
  147. return
  148. if l == 1 or l >= 126:
  149. raise WebSocketProtocolException("Invalid close frame.")
  150. if l > 2 and not skip_utf8_validation and not validate_utf8(self.data[2:]):
  151. raise WebSocketProtocolException("Invalid close frame.")
  152. code = 256 * \
  153. six.byte2int(self.data[0:1]) + six.byte2int(self.data[1:2])
  154. if not self._is_valid_close_status(code):
  155. raise WebSocketProtocolException("Invalid close opcode.")
  156. @staticmethod
  157. def _is_valid_close_status(code):
  158. return code in VALID_CLOSE_STATUS or (3000 <= code < 5000)
  159. def __str__(self):
  160. return "fin=" + str(self.fin) \
  161. + " opcode=" + str(self.opcode) \
  162. + " data=" + str(self.data)
  163. @staticmethod
  164. def create_frame(data, opcode, fin=1):
  165. """
  166. create frame to send text, binary and other data.
  167. data: data to send. This is string value(byte array).
  168. if opcode is OPCODE_TEXT and this value is unicode,
  169. data value is converted into unicode string, automatically.
  170. opcode: operation code. please see OPCODE_XXX.
  171. fin: fin flag. if set to 0, create continue fragmentation.
  172. """
  173. if opcode == ABNF.OPCODE_TEXT and isinstance(data, six.text_type):
  174. data = data.encode("utf-8")
  175. # mask must be set if send data from client
  176. return ABNF(fin, 0, 0, 0, opcode, 1, data)
  177. def format(self):
  178. """
  179. format this object to string(byte array) to send data to server.
  180. """
  181. if any(x not in (0, 1) for x in [self.fin, self.rsv1, self.rsv2, self.rsv3]):
  182. raise ValueError("not 0 or 1")
  183. if self.opcode not in ABNF.OPCODES:
  184. raise ValueError("Invalid OPCODE")
  185. length = len(self.data)
  186. if length >= ABNF.LENGTH_63:
  187. raise ValueError("data is too long")
  188. frame_header = chr(self.fin << 7
  189. | self.rsv1 << 6 | self.rsv2 << 5 | self.rsv3 << 4
  190. | self.opcode)
  191. if length < ABNF.LENGTH_7:
  192. frame_header += chr(self.mask << 7 | length)
  193. frame_header = six.b(frame_header)
  194. elif length < ABNF.LENGTH_16:
  195. frame_header += chr(self.mask << 7 | 0x7e)
  196. frame_header = six.b(frame_header)
  197. frame_header += struct.pack("!H", length)
  198. else:
  199. frame_header += chr(self.mask << 7 | 0x7f)
  200. frame_header = six.b(frame_header)
  201. frame_header += struct.pack("!Q", length)
  202. if not self.mask:
  203. return frame_header + self.data
  204. else:
  205. mask_key = self.get_mask_key(4)
  206. return frame_header + self._get_masked(mask_key)
  207. def _get_masked(self, mask_key):
  208. s = ABNF.mask(mask_key, self.data)
  209. if isinstance(mask_key, six.text_type):
  210. mask_key = mask_key.encode('utf-8')
  211. return mask_key + s
  212. @staticmethod
  213. def mask(mask_key, data):
  214. """
  215. mask or unmask data. Just do xor for each byte
  216. mask_key: 4 byte string(byte).
  217. data: data to mask/unmask.
  218. """
  219. if data is None:
  220. data = ""
  221. if isinstance(mask_key, six.text_type):
  222. mask_key = six.b(mask_key)
  223. if isinstance(data, six.text_type):
  224. data = six.b(data)
  225. if numpy:
  226. origlen = len(data)
  227. _mask_key = mask_key[3] << 24 | mask_key[2] << 16 | mask_key[1] << 8 | mask_key[0]
  228. # We need data to be a multiple of four...
  229. data += bytes(" " * (4 - (len(data) % 4)), "us-ascii")
  230. a = numpy.frombuffer(data, dtype="uint32")
  231. masked = numpy.bitwise_xor(a, [_mask_key]).astype("uint32")
  232. if len(data) > origlen:
  233. return masked.tobytes()[:origlen]
  234. return masked.tobytes()
  235. else:
  236. _m = array.array("B", mask_key)
  237. _d = array.array("B", data)
  238. return _mask(_m, _d)
  239. class frame_buffer(object):
  240. _HEADER_MASK_INDEX = 5
  241. _HEADER_LENGTH_INDEX = 6
  242. def __init__(self, recv_fn, skip_utf8_validation):
  243. self.recv = recv_fn
  244. self.skip_utf8_validation = skip_utf8_validation
  245. # Buffers over the packets from the layer beneath until desired amount
  246. # bytes of bytes are received.
  247. self.recv_buffer = []
  248. self.clear()
  249. self.lock = Lock()
  250. def clear(self):
  251. self.header = None
  252. self.length = None
  253. self.mask = None
  254. def has_received_header(self):
  255. return self.header is None
  256. def recv_header(self):
  257. header = self.recv_strict(2)
  258. b1 = header[0]
  259. if six.PY2:
  260. b1 = ord(b1)
  261. fin = b1 >> 7 & 1
  262. rsv1 = b1 >> 6 & 1
  263. rsv2 = b1 >> 5 & 1
  264. rsv3 = b1 >> 4 & 1
  265. opcode = b1 & 0xf
  266. b2 = header[1]
  267. if six.PY2:
  268. b2 = ord(b2)
  269. has_mask = b2 >> 7 & 1
  270. length_bits = b2 & 0x7f
  271. self.header = (fin, rsv1, rsv2, rsv3, opcode, has_mask, length_bits)
  272. def has_mask(self):
  273. if not self.header:
  274. return False
  275. return self.header[frame_buffer._HEADER_MASK_INDEX]
  276. def has_received_length(self):
  277. return self.length is None
  278. def recv_length(self):
  279. bits = self.header[frame_buffer._HEADER_LENGTH_INDEX]
  280. length_bits = bits & 0x7f
  281. if length_bits == 0x7e:
  282. v = self.recv_strict(2)
  283. self.length = struct.unpack("!H", v)[0]
  284. elif length_bits == 0x7f:
  285. v = self.recv_strict(8)
  286. self.length = struct.unpack("!Q", v)[0]
  287. else:
  288. self.length = length_bits
  289. def has_received_mask(self):
  290. return self.mask is None
  291. def recv_mask(self):
  292. self.mask = self.recv_strict(4) if self.has_mask() else ""
  293. def recv_frame(self):
  294. with self.lock:
  295. # Header
  296. if self.has_received_header():
  297. self.recv_header()
  298. (fin, rsv1, rsv2, rsv3, opcode, has_mask, _) = self.header
  299. # Frame length
  300. if self.has_received_length():
  301. self.recv_length()
  302. length = self.length
  303. # Mask
  304. if self.has_received_mask():
  305. self.recv_mask()
  306. mask = self.mask
  307. # Payload
  308. payload = self.recv_strict(length)
  309. if has_mask:
  310. payload = ABNF.mask(mask, payload)
  311. # Reset for next frame
  312. self.clear()
  313. frame = ABNF(fin, rsv1, rsv2, rsv3, opcode, has_mask, payload)
  314. frame.validate(self.skip_utf8_validation)
  315. return frame
  316. def recv_strict(self, bufsize):
  317. shortage = bufsize - sum(len(x) for x in self.recv_buffer)
  318. while shortage > 0:
  319. # Limit buffer size that we pass to socket.recv() to avoid
  320. # fragmenting the heap -- the number of bytes recv() actually
  321. # reads is limited by socket buffer and is relatively small,
  322. # yet passing large numbers repeatedly causes lots of large
  323. # buffers allocated and then shrunk, which results in
  324. # fragmentation.
  325. bytes_ = self.recv(min(16384, shortage))
  326. self.recv_buffer.append(bytes_)
  327. shortage -= len(bytes_)
  328. unified = six.b("").join(self.recv_buffer)
  329. if shortage == 0:
  330. self.recv_buffer = []
  331. return unified
  332. else:
  333. self.recv_buffer = [unified[bufsize:]]
  334. return unified[:bufsize]
  335. class continuous_frame(object):
  336. def __init__(self, fire_cont_frame, skip_utf8_validation):
  337. self.fire_cont_frame = fire_cont_frame
  338. self.skip_utf8_validation = skip_utf8_validation
  339. self.cont_data = None
  340. self.recving_frames = None
  341. def validate(self, frame):
  342. if not self.recving_frames and frame.opcode == ABNF.OPCODE_CONT:
  343. raise WebSocketProtocolException("Illegal frame")
  344. if self.recving_frames and \
  345. frame.opcode in (ABNF.OPCODE_TEXT, ABNF.OPCODE_BINARY):
  346. raise WebSocketProtocolException("Illegal frame")
  347. def add(self, frame):
  348. if self.cont_data:
  349. self.cont_data[1] += frame.data
  350. else:
  351. if frame.opcode in (ABNF.OPCODE_TEXT, ABNF.OPCODE_BINARY):
  352. self.recving_frames = frame.opcode
  353. self.cont_data = [frame.opcode, frame.data]
  354. if frame.fin:
  355. self.recving_frames = None
  356. def is_fire(self, frame):
  357. return frame.fin or self.fire_cont_frame
  358. def extract(self, frame):
  359. data = self.cont_data
  360. self.cont_data = None
  361. frame.data = data[1]
  362. if not self.fire_cont_frame and data[0] == ABNF.OPCODE_TEXT and not self.skip_utf8_validation and not validate_utf8(frame.data):
  363. raise WebSocketPayloadException(
  364. "cannot decode: " + repr(frame.data))
  365. return [data[0], frame]