| # -*- coding: utf-8 -*- |
| from struct import pack, unpack |
| |
| from ambari_ws4py.exc import FrameTooLargeException, ProtocolException |
| from ambari_ws4py.compat import py3k, ord, range |
| |
| # Frame opcodes defined in the spec. |
| OPCODE_CONTINUATION = 0x0 |
| OPCODE_TEXT = 0x1 |
| OPCODE_BINARY = 0x2 |
| OPCODE_CLOSE = 0x8 |
| OPCODE_PING = 0x9 |
| OPCODE_PONG = 0xa |
| |
| __all__ = ['Frame'] |
| |
| class Frame(object): |
| def __init__(self, opcode=None, body=b'', masking_key=None, fin=0, rsv1=0, rsv2=0, rsv3=0): |
| """ |
| Implements the framing protocol as defined by RFC 6455. |
| |
| .. code-block:: python |
| :linenos: |
| |
| >>> test_mask = 'XXXXXX' # perhaps from os.urandom(4) |
| >>> f = Frame(OPCODE_TEXT, 'hello world', masking_key=test_mask, fin=1) |
| >>> bytes = f.build() |
| >>> bytes.encode('hex') |
| '818bbe04e66ad6618a06d1249105cc6882' |
| >>> f = Frame() |
| >>> f.parser.send(bytes[0]) |
| 1 |
| >>> f.parser.send(bytes[1]) |
| 4 |
| |
| .. seealso:: Data Framing http://tools.ietf.org/html/rfc6455#section-5.2 |
| """ |
| if not isinstance(body, bytes): |
| raise TypeError("The body must be properly encoded") |
| |
| self.opcode = opcode |
| self.body = body |
| self.masking_key = masking_key |
| self.fin = fin |
| self.rsv1 = rsv1 |
| self.rsv2 = rsv2 |
| self.rsv3 = rsv3 |
| self.payload_length = len(body) |
| |
| self._parser = None |
| |
| @property |
| def parser(self): |
| if self._parser is None: |
| self._parser = self._parsing() |
| # Python generators must be initialized once. |
| next(self.parser) |
| return self._parser |
| |
| def _cleanup(self): |
| if self._parser: |
| self._parser.close() |
| self._parser = None |
| |
| def build(self): |
| """ |
| Builds a frame from the instance's attributes and returns |
| its bytes representation. |
| """ |
| header = b'' |
| |
| if self.fin > 0x1: |
| raise ValueError('FIN bit parameter must be 0 or 1') |
| |
| if 0x3 <= self.opcode <= 0x7 or 0xB <= self.opcode: |
| raise ValueError('Opcode cannot be a reserved opcode') |
| |
| ## +-+-+-+-+-------+ |
| ## |F|R|R|R| opcode| |
| ## |I|S|S|S| (4) | |
| ## |N|V|V|V| | |
| ## | |1|2|3| | |
| ## +-+-+-+-+-------+ |
| header = pack('!B', ((self.fin << 7) |
| | (self.rsv1 << 6) |
| | (self.rsv2 << 5) |
| | (self.rsv3 << 4) |
| | self.opcode)) |
| |
| ## +-+-------------+-------------------------------+ |
| ## |M| Payload len | Extended payload length | |
| ## |A| (7) | (16/63) | |
| ## |S| | (if payload len==126/127) | |
| ## |K| | | |
| ## +-+-+-+-+-------+-+-------------+ - - - - - - - - - - - - - - - + |
| ## | Extended payload length continued, if payload len == 127 | |
| ## + - - - - - - - - - - - - - - - +-------------------------------+ |
| if self.masking_key: mask_bit = 1 << 7 |
| else: mask_bit = 0 |
| |
| length = self.payload_length |
| if length < 126: |
| header += pack('!B', (mask_bit | length)) |
| elif length < (1 << 16): |
| header += pack('!B', (mask_bit | 126)) + pack('!H', length) |
| elif length < (1 << 63): |
| header += pack('!B', (mask_bit | 127)) + pack('!Q', length) |
| else: |
| raise FrameTooLargeException() |
| |
| ## + - - - - - - - - - - - - - - - +-------------------------------+ |
| ## | |Masking-key, if MASK set to 1 | |
| ## +-------------------------------+-------------------------------+ |
| ## | Masking-key (continued) | Payload Data | |
| ## +-------------------------------- - - - - - - - - - - - - - - - + |
| ## : Payload Data continued ... : |
| ## + - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - + |
| ## | Payload Data continued ... | |
| ## +---------------------------------------------------------------+ |
| body = self.body |
| if not self.masking_key: |
| return bytes(header + body) |
| |
| return bytes(header + self.masking_key + self.mask(body)) |
| |
| def _parsing(self): |
| """ |
| Generator to parse bytes into a frame. Yields until |
| enough bytes have been read or an error is met. |
| """ |
| buf = b'' |
| some_bytes = b'' |
| |
| # yield until we get the first header's byte |
| while not some_bytes: |
| some_bytes = (yield 1) |
| |
| first_byte = some_bytes[0] if isinstance(some_bytes, bytearray) else ord(some_bytes[0]) |
| # frame-fin = %x0 ; more frames of this message follow |
| # / %x1 ; final frame of this message |
| self.fin = (first_byte >> 7) & 1 |
| self.rsv1 = (first_byte >> 6) & 1 |
| self.rsv2 = (first_byte >> 5) & 1 |
| self.rsv3 = (first_byte >> 4) & 1 |
| self.opcode = first_byte & 0xf |
| |
| # frame-rsv1 = %x0 ; 1 bit, MUST be 0 unless negotiated otherwise |
| # frame-rsv2 = %x0 ; 1 bit, MUST be 0 unless negotiated otherwise |
| # frame-rsv3 = %x0 ; 1 bit, MUST be 0 unless negotiated otherwise |
| if self.rsv1 or self.rsv2 or self.rsv3: |
| raise ProtocolException() |
| |
| # control frames between 3 and 7 as well as above 0xA are currently reserved |
| if 2 < self.opcode < 8 or self.opcode > 0xA: |
| raise ProtocolException() |
| |
| # control frames cannot be fragmented |
| if self.opcode > 0x7 and self.fin == 0: |
| raise ProtocolException() |
| |
| # do we already have enough some_bytes to continue? |
| some_bytes = some_bytes[1:] if some_bytes and len(some_bytes) > 1 else b'' |
| |
| # Yield until we get the second header's byte |
| while not some_bytes: |
| some_bytes = (yield 1) |
| |
| second_byte = some_bytes[0] if isinstance(some_bytes, bytearray) else ord(some_bytes[0]) |
| mask = (second_byte >> 7) & 1 |
| self.payload_length = second_byte & 0x7f |
| |
| # All control frames MUST have a payload length of 125 some_bytes or less |
| if self.opcode > 0x7 and self.payload_length > 125: |
| raise FrameTooLargeException() |
| |
| if some_bytes and len(some_bytes) > 1: |
| buf = some_bytes[1:] |
| some_bytes = buf |
| else: |
| buf = b'' |
| some_bytes = b'' |
| |
| if self.payload_length == 127: |
| # This will compute the actual application data size |
| if len(buf) < 8: |
| nxt_buf_size = 8 - len(buf) |
| some_bytes = (yield nxt_buf_size) |
| some_bytes = buf + (some_bytes or b'') |
| while len(some_bytes) < 8: |
| b = (yield 8 - len(some_bytes)) |
| if b is not None: |
| some_bytes = some_bytes + b |
| if len(some_bytes) > 8: |
| buf = some_bytes[8:] |
| some_bytes = some_bytes[:8] |
| else: |
| some_bytes = buf[:8] |
| buf = buf[8:] |
| extended_payload_length = some_bytes |
| self.payload_length = unpack( |
| '!Q', extended_payload_length)[0] |
| if self.payload_length > 0x7FFFFFFFFFFFFFFF: |
| raise FrameTooLargeException() |
| elif self.payload_length == 126: |
| if len(buf) < 2: |
| nxt_buf_size = 2 - len(buf) |
| some_bytes = (yield nxt_buf_size) |
| some_bytes = buf + (some_bytes or b'') |
| while len(some_bytes) < 2: |
| b = (yield 2 - len(some_bytes)) |
| if b is not None: |
| some_bytes = some_bytes + b |
| if len(some_bytes) > 2: |
| buf = some_bytes[2:] |
| some_bytes = some_bytes[:2] |
| else: |
| some_bytes = buf[:2] |
| buf = buf[2:] |
| extended_payload_length = some_bytes |
| self.payload_length = unpack( |
| '!H', extended_payload_length)[0] |
| |
| if mask: |
| if len(buf) < 4: |
| nxt_buf_size = 4 - len(buf) |
| some_bytes = (yield nxt_buf_size) |
| some_bytes = buf + (some_bytes or b'') |
| while not some_bytes or len(some_bytes) < 4: |
| b = (yield 4 - len(some_bytes)) |
| if b is not None: |
| some_bytes = some_bytes + b |
| if len(some_bytes) > 4: |
| buf = some_bytes[4:] |
| else: |
| some_bytes = buf[:4] |
| buf = buf[4:] |
| self.masking_key = some_bytes |
| |
| if len(buf) < self.payload_length: |
| nxt_buf_size = self.payload_length - len(buf) |
| some_bytes = (yield nxt_buf_size) |
| some_bytes = buf + (some_bytes or b'') |
| while len(some_bytes) < self.payload_length: |
| l = self.payload_length - len(some_bytes) |
| b = (yield l) |
| if b is not None: |
| some_bytes = some_bytes + b |
| else: |
| if self.payload_length == len(buf): |
| some_bytes = buf |
| else: |
| some_bytes = buf[:self.payload_length] |
| |
| self.body = some_bytes |
| yield |
| |
| def mask(self, data): |
| """ |
| Performs the masking or unmasking operation on data |
| using the simple masking algorithm: |
| |
| .. |
| j = i MOD 4 |
| transformed-octet-i = original-octet-i XOR masking-key-octet-j |
| |
| """ |
| masked = bytearray(data) |
| if py3k: key = self.masking_key |
| else: key = map(ord, self.masking_key) |
| for i in range(len(data)): |
| masked[i] = masked[i] ^ key[i%4] |
| return masked |
| unmask = mask |
