###############################################################################
#
# The MIT License (MIT)
#
# Copyright (c) typedef int GmbH
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
#
###############################################################################

import unittest
from unittest.mock import Mock

from autobahn.testutil import FakeTransport
from autobahn.twisted.rawsocket import (
    WampRawSocketClientFactory,
    WampRawSocketClientProtocol,
    WampRawSocketServerFactory,
    WampRawSocketServerProtocol,
)


class RawSocketHandshakeTests(unittest.TestCase):
    def test_handshake_succeeds(self):
        """
        A client can connect to a server.
        """
        session_mock = Mock()
        t = FakeTransport()
        f = WampRawSocketClientFactory(lambda: session_mock)
        p = WampRawSocketClientProtocol()
        p.transport = t
        p.factory = f

        server_session_mock = Mock()
        st = FakeTransport()
        sf = WampRawSocketServerFactory(lambda: server_session_mock)
        sp = WampRawSocketServerProtocol()
        sp.transport = st
        sp.factory = sf

        sp.connectionMade()
        p.connectionMade()

        # Send the server the client handshake
        sp.dataReceived(t._written[0:1])
        sp.dataReceived(t._written[1:4])

        # Send the client the server handshake
        p.dataReceived(st._written)

        # The handshake succeeds, a session on each end is created
        # onOpen is called on the session
        session_mock.onOpen.assert_called_once_with(p)
        server_session_mock.onOpen.assert_called_once_with(sp)

    def test_receive_limit_advertised_and_enforced(self):
        """
        ``setProtocolOptions(maxMessagePayloadSize=...)`` configures the server's
        received-message size cap: the advertised handshake length exponent and
        the enforced ``MAX_LENGTH`` are the configured max rounded up to the next
        power of two. This mirrors the asyncio backend (see
        ``test_aio_rawsocket.RECV_LIMIT_CASES``) so both backends make identical
        accept/reject decisions for the same configuration (#1911).
        """
        # (maxMessagePayloadSize, advertised length exponent, enforced cap)
        cases = [
            (512, 0, 512),
            (1000, 1, 1024),  # rounded up to the next power of two
            (1024, 1, 1024),
            (4096, 3, 4096),
            (2**20, 11, 2**20),
            (2**24, 15, 2**24),
        ]
        for max_size, exp, cap in cases:
            with self.subTest(max_size=max_size):
                sf = WampRawSocketServerFactory(lambda: Mock())
                sf.setProtocolOptions(maxMessagePayloadSize=max_size)
                sp = sf.buildProtocol(None)
                sp.transport = FakeTransport()
                sp.connectionMade()
                ser_id = sorted(sf._serializers.keys())[0]
                # client opening handshake: magic, (length-exp 15 | serializer)
                sp.dataReceived(bytes([0x7F, 0xF0 | ser_id, 0, 0]))
                written = sp.transport._written
                self.assertEqual(written[1] >> 4, exp)
                self.assertEqual(sp.MAX_LENGTH, cap)

    def test_server_bad_magic_byte_aborts_cleanly(self):
        """
        A server receiving an invalid magic byte in the opening handshake
        aborts the transport cleanly instead of raising ``TransportLost`` out
        of ``dataReceived`` (regression test for #1850).
        """
        server_session_mock = Mock()
        st = FakeTransport()
        sf = WampRawSocketServerFactory(lambda: server_session_mock)
        sp = WampRawSocketServerProtocol()
        sp.transport = st
        sp.factory = sf

        sp.connectionMade()

        # any non-0x7f first octet is an invalid magic byte; this must not raise
        sp.dataReceived(b"GET ")

        # the transport was aborted and no WAMP session was started
        self.assertTrue(st.abort_called())
        server_session_mock.onOpen.assert_not_called()

    def test_client_bad_magic_byte_aborts_cleanly(self):
        """
        A client receiving an invalid magic byte in the opening handshake
        aborts the transport cleanly instead of raising ``TransportLost`` out
        of ``dataReceived`` (regression test for #1850).
        """
        session_mock = Mock()
        t = FakeTransport()
        f = WampRawSocketClientFactory(lambda: session_mock)
        p = WampRawSocketClientProtocol()
        p.transport = t
        p.factory = f

        p.connectionMade()

        # any non-0x7f first octet is an invalid magic byte; this must not raise
        p.dataReceived(b"GET ")

        # the transport was aborted and no WAMP session was started
        self.assertTrue(t.abort_called())
        session_mock.onOpen.assert_not_called()
