# -*- coding: utf-8 -*-

import zlib

import pytest

import wsproto.extensions as wpext
import wsproto.frame_protocol as fp


class TestPerMessageDeflate(object):
    parameter_sets = [
        {
            'client_no_context_takeover': False,
            'client_max_window_bits': 15,
            'server_no_context_takeover': False,
            'server_max_window_bits': 15,
        },
        {
            'client_no_context_takeover': True,
            'client_max_window_bits': 9,
            'server_no_context_takeover': False,
            'server_max_window_bits': 15,
        },
        {
            'client_no_context_takeover': False,
            'client_max_window_bits': 15,
            'server_no_context_takeover': True,
            'server_max_window_bits': 9,
        },
        {
            'client_no_context_takeover': True,
            'client_max_window_bits': 8,
            'server_no_context_takeover': True,
            'server_max_window_bits': 9,
        },
        {
            'client_no_context_takeover': True,
            'server_max_window_bits': 9,
        },
        {
            'server_no_context_takeover': True,
            'client_max_window_bits': 8,
        },
        {
            'client_max_window_bits': None,
            'server_max_window_bits': None,
        },
        {},
    ]

    def make_offer_string(self, params):
        offer = ['permessage-deflate']

        if 'client_max_window_bits' in params:
            if params['client_max_window_bits'] is None:
                offer.append('client_max_window_bits')
            else:
                offer.append('client_max_window_bits=%d' %
                             params['client_max_window_bits'])
        if 'server_max_window_bits' in params:
            if params['server_max_window_bits'] is None:
                offer.append('server_max_window_bits')
            else:
                offer.append('server_max_window_bits=%d' %
                             params['server_max_window_bits'])
        if params.get('client_no_context_takeover', False):
            offer.append('client_no_context_takeover')
        if params.get('server_no_context_takeover', False):
            offer.append('server_no_context_takeover')

        return '; '.join(offer)

    def compare_params_to_string(self, params, ext, param_string):
        if 'client_max_window_bits' in params:
            if params['client_max_window_bits'] is None:
                bits = ext.client_max_window_bits
            else:
                bits = params['client_max_window_bits']
            assert 'client_max_window_bits=%d' % bits in param_string
        if 'server_max_window_bits' in params:
            if params['server_max_window_bits'] is None:
                bits = ext.server_max_window_bits
            else:
                bits = params['server_max_window_bits']
            assert 'server_max_window_bits=%d' % bits in param_string
        if params.get('client_no_context_takeover', False):
            assert 'client_no_context_takeover' in param_string
        if params.get('server_no_context_takeover', False):
            assert 'server_no_context_takeover' in param_string

    @pytest.mark.parametrize('params', parameter_sets)
    def test_offer(self, params):
        ext = wpext.PerMessageDeflate(**params)
        offer = ext.offer(None)

        self.compare_params_to_string(params, ext, offer)

    @pytest.mark.parametrize('params', parameter_sets)
    def test_finalize(self, params):
        ext = wpext.PerMessageDeflate()
        assert not ext.enabled()

        params = dict(params)
        if 'client_max_window_bits' in params:
            if params['client_max_window_bits'] is None:
                del params['client_max_window_bits']
        if 'server_max_window_bits' in params:
            if params['server_max_window_bits'] is None:
                del params['server_max_window_bits']
        offer = self.make_offer_string(params)
        ext.finalize(None, offer)

        if params.get('client_max_window_bits', None):
            assert ext.client_max_window_bits == \
                params['client_max_window_bits']
        if params.get('server_max_window_bits', None):
            assert ext.server_max_window_bits == \
                params['server_max_window_bits']
        assert ext.client_no_context_takeover is \
            params.get('client_no_context_takeover', False)
        assert ext.server_no_context_takeover is \
            params.get('server_no_context_takeover', False)

        assert ext.enabled()

    def test_finalize_ignores_rubbish(self):
        ext = wpext.PerMessageDeflate()
        assert not ext.enabled()

        ext.finalize(None, 'i am the lizard queen; worship me')

        assert ext.enabled()

    @pytest.mark.parametrize('params', parameter_sets)
    def test_accept(self, params):
        ext = wpext.PerMessageDeflate()
        assert not ext.enabled()

        offer = self.make_offer_string(params)
        response = ext.accept(None, offer)

        if ext.client_no_context_takeover:
            assert 'client_no_context_takeover' in response
        if ext.server_no_context_takeover:
            assert 'server_no_context_takeover' in response
        if 'client_max_window_bits' in params:
            if params['client_max_window_bits'] is None:
                bits = ext.client_max_window_bits
            else:
                bits = params['client_max_window_bits']
            assert ext.client_max_window_bits == bits
            assert 'client_max_window_bits=%d' % bits in response
        if 'server_max_window_bits' in params:
            if params['server_max_window_bits'] is None:
                bits = ext.server_max_window_bits
            else:
                bits = params['server_max_window_bits']
            assert ext.server_max_window_bits == bits
            assert 'server_max_window_bits=%d' % bits in response

    def test_accept_ignores_rubbish(self):
        ext = wpext.PerMessageDeflate()
        assert not ext.enabled()

        ext.accept(None, 'i am the lizard queen; worship me')

        assert ext.enabled()

    def test_inbound_uncompressed_control_frame(self):
        payload = b'x' * 23

        ext = wpext.PerMessageDeflate()
        ext._enabled = True
        proto = fp.FrameProtocol(client=True, extensions=[ext])

        result = ext.frame_inbound_header(proto, fp.Opcode.PING,
                                          fp.RsvBits(False, False, False),
                                          len(payload))
        assert result.rsv1

        data = ext.frame_inbound_payload_data(proto, payload)
        assert data == payload

        assert ext.frame_inbound_complete(proto, True) is None

    def test_inbound_compressed_control_frame(self):
        payload = b'x' * 23

        ext = wpext.PerMessageDeflate()
        ext._enabled = True
        proto = fp.FrameProtocol(client=True, extensions=[ext])

        result = ext.frame_inbound_header(proto, fp.Opcode.PING,
                                          fp.RsvBits(True, False, False),
                                          len(payload))
        assert result == fp.CloseReason.PROTOCOL_ERROR

    def test_inbound_compressed_continuation_frame(self):
        payload = b'x' * 23

        ext = wpext.PerMessageDeflate()
        ext._enabled = True
        proto = fp.FrameProtocol(client=True, extensions=[ext])

        result = ext.frame_inbound_header(proto, fp.Opcode.CONTINUATION,
                                          fp.RsvBits(True, False, False),
                                          len(payload))
        assert result == fp.CloseReason.PROTOCOL_ERROR

    def test_inbound_uncompressed_data_frame(self):
        payload = b'x' * 23

        ext = wpext.PerMessageDeflate()
        ext._enabled = True
        proto = fp.FrameProtocol(client=True, extensions=[ext])

        result = ext.frame_inbound_header(proto, fp.Opcode.BINARY,
                                          fp.RsvBits(False, False, False),
                                          len(payload))
        assert result.rsv1

        data = ext.frame_inbound_payload_data(proto, payload)
        assert data == payload

        assert ext.frame_inbound_complete(proto, True) is None

    @pytest.mark.parametrize('client', [True, False])
    def test_client_inbound_compressed_single_data_frame(self, client):
        payload = b'x' * 23
        compressed_payload = b'\xaa\xa8\xc0\n\x00\x00'

        ext = wpext.PerMessageDeflate()
        ext._enabled = True
        proto = fp.FrameProtocol(client=client, extensions=[ext])

        result = ext.frame_inbound_header(proto, fp.Opcode.BINARY,
                                          fp.RsvBits(True, False, False),
                                          len(compressed_payload))
        assert result.rsv1

        data = ext.frame_inbound_payload_data(proto, compressed_payload)
        data += ext.frame_inbound_complete(proto, True)
        assert data == payload

    @pytest.mark.parametrize('client', [True, False])
    def test_client_inbound_compressed_multiple_data_frames(self, client):
        payload = b'x' * 23
        compressed_payload = b'\xaa\xa8\xc0\n\x00\x00'
        split = 3
        data = b''

        ext = wpext.PerMessageDeflate()
        ext._enabled = True
        proto = fp.FrameProtocol(client=client, extensions=[ext])

        result = ext.frame_inbound_header(proto, fp.Opcode.BINARY,
                                          fp.RsvBits(True, False, False),
                                          split)
        assert result.rsv1
        result = ext.frame_inbound_payload_data(proto,
                                                compressed_payload[:split])
        assert not isinstance(result, fp.CloseReason)
        data += result
        assert ext.frame_inbound_complete(proto, False) is None

        result = ext.frame_inbound_header(proto, fp.Opcode.CONTINUATION,
                                          fp.RsvBits(False, False, False),
                                          len(compressed_payload) - split)
        assert result.rsv1
        result = ext.frame_inbound_payload_data(proto,
                                                compressed_payload[split:])
        assert not isinstance(result, fp.CloseReason)
        data += result

        result = ext.frame_inbound_complete(proto, True)
        assert not isinstance(result, fp.CloseReason)
        data += result

        assert data == payload

    @pytest.mark.parametrize('client', [True, False])
    def test_client_decompress_after_uncompressible_frame(self, client):
        ext = wpext.PerMessageDeflate()
        ext._enabled = True
        proto = fp.FrameProtocol(client=client, extensions=[ext])

        # A PING frame
        result = ext.frame_inbound_header(proto, fp.Opcode.PING,
                                          fp.RsvBits(False, False, False),
                                          0)
        result = ext.frame_inbound_payload_data(proto, b'')
        assert not isinstance(result, fp.CloseReason)
        assert ext.frame_inbound_complete(proto, True) is None

        # A compressed TEXT frame
        payload = b'x' * 23
        compressed_payload = b'\xaa\xa8\xc0\n\x00\x00'

        result = ext.frame_inbound_header(proto, fp.Opcode.TEXT,
                                          fp.RsvBits(True, False, False),
                                          len(compressed_payload))
        assert result.rsv1
        result = ext.frame_inbound_payload_data(proto, compressed_payload)
        assert result == payload

        result = ext.frame_inbound_complete(proto, True)
        assert not isinstance(result, fp.CloseReason)

    def test_inbound_bad_zlib_payload(self):
        compressed_payload = b'x' * 23

        ext = wpext.PerMessageDeflate()
        ext._enabled = True
        proto = fp.FrameProtocol(client=True, extensions=[ext])

        result = ext.frame_inbound_header(proto, fp.Opcode.BINARY,
                                          fp.RsvBits(True, False, False),
                                          len(compressed_payload))
        assert result.rsv1
        result = ext.frame_inbound_payload_data(proto, compressed_payload)
        assert result is fp.CloseReason.INVALID_FRAME_PAYLOAD_DATA

    def test_inbound_bad_zlib_decoder_end_state(self, monkeypatch):
        compressed_payload = b'x' * 23

        ext = wpext.PerMessageDeflate()
        ext._enabled = True
        proto = fp.FrameProtocol(client=True, extensions=[ext])

        result = ext.frame_inbound_header(proto, fp.Opcode.BINARY,
                                          fp.RsvBits(True, False, False),
                                          len(compressed_payload))
        assert result.rsv1

        class FailDecompressor(object):
            def decompress(self, data):
                return b''

            def flush(self):
                raise zlib.error()

        monkeypatch.setattr(ext, '_decompressor', FailDecompressor())

        result = ext.frame_inbound_complete(proto, True)
        assert result is fp.CloseReason.INVALID_FRAME_PAYLOAD_DATA

    @pytest.mark.parametrize('client,no_context_takeover', [
        (True, True),
        (True, False),
        (False, True),
        (False, False),
    ])
    def test_decompressor_reset(self, client, no_context_takeover):
        if client:
            args = {'server_no_context_takeover': no_context_takeover}
        else:
            args = {'client_no_context_takeover': no_context_takeover}
        ext = wpext.PerMessageDeflate(**args)
        ext._enabled = True
        proto = fp.FrameProtocol(client=client, extensions=[ext])

        result = ext.frame_inbound_header(proto, fp.Opcode.BINARY,
                                          fp.RsvBits(True, False, False), 0)
        assert result.rsv1

        assert ext._decompressor is not None

        result = ext.frame_inbound_complete(proto, True)
        assert not isinstance(result, fp.CloseReason)

        if no_context_takeover:
            assert ext._decompressor is None
        else:
            assert ext._decompressor is not None

        result = ext.frame_inbound_header(proto, fp.Opcode.BINARY,
                                          fp.RsvBits(True, False, False), 0)
        assert result.rsv1

        assert ext._decompressor is not None

    def test_outbound_uncompressible_opcode(self):
        ext = wpext.PerMessageDeflate()
        ext._enabled = True
        proto = fp.FrameProtocol(client=True, extensions=[ext])

        rsv = fp.RsvBits(False, False, False)
        payload = b'x' * 23

        rsv, data = ext.frame_outbound(proto, fp.Opcode.PING, rsv, payload,
                                       True)

        assert rsv.rsv1 is False
        assert data == payload

    @pytest.mark.parametrize('client', [True, False])
    def test_outbound_compress_single_frame(self, client):
        ext = wpext.PerMessageDeflate()
        ext._enabled = True
        proto = fp.FrameProtocol(client=client, extensions=[ext])

        rsv = fp.RsvBits(False, False, False)
        payload = b'x' * 23
        compressed_payload = b'\xaa\xa8\xc0\n\x00\x00'

        rsv, data = ext.frame_outbound(proto, fp.Opcode.BINARY, rsv, payload,
                                       True)

        assert rsv.rsv1 is True
        assert data == compressed_payload

    @pytest.mark.parametrize('client', [True, False])
    def test_outbound_compress_multiple_frames(self, client):
        ext = wpext.PerMessageDeflate()
        ext._enabled = True
        proto = fp.FrameProtocol(client=client, extensions=[ext])

        rsv = fp.RsvBits(False, False, False)
        payload = b'x' * 23
        split = 12
        compressed_payload = b'\xaa\xa8\xc0\n\x00\x00'

        rsv, data = ext.frame_outbound(proto, fp.Opcode.BINARY, rsv,
                                       payload[:split], False)
        assert rsv.rsv1 is True

        rsv = fp.RsvBits(False, False, False)
        rsv, more_data = ext.frame_outbound(proto, fp.Opcode.CONTINUATION, rsv,
                                            payload[split:], True)
        assert rsv.rsv1 is False
        assert data + more_data == compressed_payload

    @pytest.mark.parametrize('client,no_context_takeover', [
        (True, True),
        (True, False),
        (False, True),
        (False, False),
    ])
    def test_compressor_reset(self, client, no_context_takeover):
        if client:
            args = {'client_no_context_takeover': no_context_takeover}
        else:
            args = {'server_no_context_takeover': no_context_takeover}
        ext = wpext.PerMessageDeflate(**args)
        ext._enabled = True
        proto = fp.FrameProtocol(client=client, extensions=[ext])
        rsv = fp.RsvBits(False, False, False)

        rsv, data = ext.frame_outbound(proto, fp.Opcode.BINARY, rsv, b'',
                                       False)
        assert rsv.rsv1 is True
        assert ext._compressor is not None

        rsv = fp.RsvBits(False, False, False)
        rsv, data = ext.frame_outbound(proto, fp.Opcode.CONTINUATION, rsv, b'',
                                       True)
        assert rsv.rsv1 is False
        if no_context_takeover:
            assert ext._compressor is None
        else:
            assert ext._compressor is not None

        rsv = fp.RsvBits(False, False, False)
        rsv, data = ext.frame_outbound(proto, fp.Opcode.BINARY, rsv, b'',
                                       False)
        assert rsv.rsv1 is True
        assert ext._compressor is not None

    @pytest.mark.parametrize('params', parameter_sets)
    def test_repr(self, params):
        ext = wpext.PerMessageDeflate(**params)
        self.compare_params_to_string(params, ext, repr(ext))
