"""Unittests for eap_state_machine.FullEAPStateMachine"""
# pylint: disable=missing-docstring

import logging
import tempfile
import unittest
from queue import Queue

from chewie.eap import Eap
from chewie.event import EventMessageReceived, EventRadiusMessageReceived, EventPortStatusChange
from chewie.mac_address import MacAddress
from chewie.message_parser import EapolStartMessage, IdentityMessage, Md5ChallengeMessage, \
    SuccessMessage, LegacyNakMessage, TtlsMessage, FailureMessage, EapolLogoffMessage
from chewie.radius_attributes import State
from chewie.state_machines.eap_state_machine import FullEAPStateMachine

from helpers import FakeTimerScheduler


def check_counters(_func=None, *,
                   expected_auth_counter=0, expected_failure_counter=0, expected_logoff_counter=0):
    """Decorator to check the handlers have been called the
     correct number of times at the end of each test"""

    def decorator_check_counters(func):
        def wrapper(self):
            start_auth_counter = self.auth_counter
            start_failure_counter = self.failure_counter
            start_logoff_counter = self.logoff_counter
            ret = func(self)
            self.assertEqual(self.auth_counter,
                             start_auth_counter + expected_auth_counter)
            self.assertEqual(self.failure_counter,
                             start_failure_counter + expected_failure_counter)
            self.assertEqual(self.logoff_counter,
                             start_logoff_counter + expected_logoff_counter)
            return ret

        return wrapper

    if _func is None:
        return decorator_check_counters
    else:
        return decorator_check_counters(_func)


class FullStateMachineStartTestCase(unittest.TestCase):
    # TODO tests could be more thorough, and test that
    # the correct packet (type/content) has been put in its respective queue.
    # Would also be nice to check that the states are correctly transitioned through,
    # and not just the final resting spot. Not sure how to do that - maybe parse the log??

    PORT_ID_MAC = MacAddress.from_string("00:00:00:00:00:01")

    def setUp(self):
        logger = logging.getLogger()
        logger.level = logging.DEBUG
        self.log_file = tempfile.NamedTemporaryFile()
        logger.addHandler(logging.FileHandler(self.log_file.name))

        self.eap_output_queue = Queue()
        self.radius_output_queue = Queue()
        self.timer_scheduler = FakeTimerScheduler()
        self.src_mac = MacAddress.from_string("00:12:34:56:78:90")
        log_prefix = "chewie.SM - port: %s, client: %s" % (self.src_mac, self.PORT_ID_MAC)

        self.sm = FullEAPStateMachine(self.eap_output_queue, self.radius_output_queue, self.src_mac,
                                      self.timer_scheduler,
                                      self.auth_handler, self.failure_handler, self.logoff_handler,
                                      log_prefix)
        # find ways to inject these - overriding consts isn't ideal
        self.MAX_RETRANSMITS = 3
        self.sm.MAX_RETRANS = self.MAX_RETRANSMITS
        self.sm.DEFAULT_TIMEOUT = 0.1
        self.sm.port_enabled = True
        self.sm.eap_restart = True

        self.auth_counter = 0
        self.failure_counter = 0
        self.logoff_counter = 0

    def tearDown(self):
        with open(self.log_file.name) as log:
            self.assertNotIn('aaaEapResp is true. but data is false. This should never happen',
                             log.read())

    def auth_handler(self, client_mac, port_id_mac, timer, vlan_name,
                     filter_id):  # pylint: disable=unused-argument
        self.auth_counter += 1
        print('Successful auth from MAC %s' % str(client_mac))

    def failure_handler(self, client_mac, port_id_mac):  # pylint: disable=unused-argument
        self.failure_counter += 1
        print('failure from MAC %s' % str(client_mac))

    def logoff_handler(self, client_mac, port_id_mac):  # pylint: disable=unused-argument
        self.logoff_counter += 1
        print('logoff from MAC %s' % str(client_mac))

    @check_counters
    def test_eap_start(self):
        # input EAPStart packet.
        # output EAPIdentityRequest on eap_output_q
        message = EapolStartMessage(self.src_mac)
        self.sm.event(EventMessageReceived(message, self.PORT_ID_MAC))
        self.assertEqual(self.sm.state, self.sm.IDLE)

        self.assertEqual(self.eap_output_queue.qsize(), 1)
        output = self.eap_output_queue.get_nowait()[0]
        self.assertIsInstance(output, IdentityMessage)
        self.assertEqual(self.radius_output_queue.qsize(), 0)

        return output  # Used by test_identity_response

    @check_counters
    def test_eap_restart(self):
        self.test_eap_start()
        message = EapolStartMessage(self.src_mac)
        self.sm.event(EventMessageReceived(message, self.PORT_ID_MAC))
        self.assertEqual(self.sm.state, self.sm.IDLE)

        self.assertEqual(self.eap_output_queue.qsize(), 1)
        output = self.eap_output_queue.get_nowait()[0]
        self.assertIsInstance(output, IdentityMessage)
        self.assertEqual(self.radius_output_queue.qsize(), 0)

    @check_counters
    def test_no_radius_state_attribute_after_restart(self):
        """The RADIUS State Attribute is returned by the RADIUS server,
        and MUST be sent back in the next RADIUS Message.
        Unless the authentication process has been restarted or its a new one.
        Here we restart part way through another."""
        self.test_md5_challenge_request()

        self.assertEqual(self.sm.radius_state_attribute.data_type.data(), b'random state')

        self.test_eap_restart()
        self.assertIsNone(self.sm.radius_state_attribute)

    @check_counters
    def test_timeout_failure_from_max_retransmits(self):
        """Go to timeout failure from exceeding max retransmits (also tests retransmitting)"""
        output0 = self.test_eap_start()

        old_radius_count = self.radius_output_queue.qsize()
        self.timer_scheduler.run_jobs()

        output1 = self.eap_output_queue.get_nowait()[0]
        output2 = self.eap_output_queue.get_nowait()[0]
        output3 = self.eap_output_queue.get_nowait()[0]
        self.assertIsInstance(output0, IdentityMessage)
        self.assertIsInstance(output1, IdentityMessage)
        self.assertEqual(output0, output1)
        self.assertEqual(output0, output2)
        self.assertEqual(output0, output3)
        self.assertEqual(self.radius_output_queue.qsize(), 0)

        self.assertEqual(self.sm.state, self.sm.TIMEOUT_FAILURE)
        self.assertEqual(self.eap_output_queue.qsize(), 0)
        self.assertEqual(old_radius_count, self.radius_output_queue.qsize())

    @check_counters(expected_auth_counter=1)
    def test_auth_success_after_timeout_failure2_from_max_retransmits(self):
        self.test_timeout_failure2_from_max_retransmits()
        self.eap_output_queue.queue.clear()
        self.test_success2()

    @check_counters(expected_auth_counter=1)
    def test_auth_success_after_timeout_failure2_from_aaa_timeout(self):
        self.test_timeout_failure2_from_aaa_timeout()
        self.test_success2()

    @check_counters(expected_auth_counter=1)
    def test_auth_success_after_timeout_failure_after_max_retransmits(self):
        self.test_timeout_failure_from_max_retransmits()
        self.test_success2()

    @check_counters()
    def test_leave_timeout_failure2_with_identiy_response(self):
        self.test_timeout_failure2_from_max_retransmits()
        start_eap_q_size = self.eap_output_queue.qsize()
        start_radius_q_size = self.radius_output_queue.qsize()
        message = IdentityMessage(self.src_mac, 25, Eap.RESPONSE, "host1user")
        self.sm.event(EventMessageReceived(message, self.PORT_ID_MAC))
        self.assertEqual(self.sm.state, self.sm.AAA_IDLE)

        self.assertEqual(self.eap_output_queue.qsize(), start_eap_q_size + 0)
        self.assertEqual(self.radius_output_queue.qsize(), start_radius_q_size + 1)
        self.assertIsInstance(self.radius_output_queue.get_nowait()[0], IdentityMessage)

    @check_counters
    def test_timeout_failure2_from_aaa_timeout(self):
        """no response from AAA server equals timeout_failure2"""
        self.test_md5_challenge_response()
        old_eap_count = self.eap_output_queue.qsize()
        old_radius_count = self.radius_output_queue.qsize()

        self.timer_scheduler.run_jobs()

        self.assertEqual(self.sm.state, self.sm.TIMEOUT_FAILURE2)
        self.assertEqual(old_eap_count, self.eap_output_queue.qsize())
        self.assertEqual(old_radius_count, self.radius_output_queue.qsize())

    @check_counters
    def test_timeout_failure2_from_max_retransmits(self):
        """If client does not respond when in passthrough mode,
         send again and again until max retransmit counter is reached."""
        self.test_md5_challenge_request()
        self.assertEqual(self.eap_output_queue.qsize(), 0)
        old_radius_count = self.radius_output_queue.qsize()

        self.timer_scheduler.run_jobs()

        self.assertEqual(self.sm.state, self.sm.TIMEOUT_FAILURE2)
        self.assertEqual(self.MAX_RETRANSMITS, self.eap_output_queue.qsize())
        self.assertEqual(old_radius_count, self.radius_output_queue.qsize())

    @check_counters(expected_auth_counter=1)
    def test_disabled_state(self):
        """move to disabled and then from disabled"""
        self.test_success2()
        # move to disabled. e.g. link down.
        self.sm.event(EventPortStatusChange(False))
        self.assertEqual(self.sm.state, self.sm.DISABLED)

        self.assertEqual(self.eap_output_queue.qsize(), 0)
        self.assertEqual(self.radius_output_queue.qsize(), 0)
        self.assertEqual(self.auth_counter, 1)
        # don't transition to initialize (still not enabled)
        message = EapolStartMessage(self.src_mac)
        self.sm.event(EventMessageReceived(message, self.PORT_ID_MAC))

        self.assertEqual(self.sm.state, self.sm.DISABLED)
        self.assertEqual(self.eap_output_queue.qsize(), 0)
        self.assertEqual(self.radius_output_queue.qsize(), 0)

        # port is enabled again
        self.sm.event(EventPortStatusChange(True))

        self.assertEqual(self.sm.state, self.sm.IDLE)

        self.assertEqual(self.eap_output_queue.qsize(), 1)
        self.assertEqual(self.radius_output_queue.qsize(), 0)

    @check_counters
    def test_identity_response(self):
        _id = self.test_eap_start().message_id
        # input EapIdentityResponse
        # output EapIdentityResponse on radius_output_q
        message = IdentityMessage(self.src_mac, _id, Eap.RESPONSE, "host1user")
        self.sm.event(EventMessageReceived(message, self.PORT_ID_MAC))

        self.assertEqual(self.sm.state, self.sm.AAA_IDLE)

        self.assertEqual(self.eap_output_queue.qsize(), 0)
        self.assertEqual(self.radius_output_queue.qsize(), 1)
        self.assertIsInstance(self.radius_output_queue.get_nowait()[0], IdentityMessage)

    @check_counters
    def test_md5_challenge_request(self):
        self.test_identity_response()

        eap_message = Md5ChallengeMessage(self.src_mac, 2, Eap.REQUEST,
                                          bytes.fromhex("74d3db089b727d9cc5774599e4a32a29"),
                                          b"host1user")
        self.sm.event(EventRadiusMessageReceived(eap_message, State.create(b"random state")))

        self.assertEqual(self.sm.state, self.sm.IDLE2)

        self.assertEqual(self.eap_output_queue.qsize(), 1)
        output = self.eap_output_queue.get_nowait()[0]
        self.assertIsInstance(output, Md5ChallengeMessage)

        self.assertEqual(self.radius_output_queue.qsize(), 0)
        return output

    @check_counters
    def test_md5_challenge_response(self):
        self.test_md5_challenge_request()

        message = Md5ChallengeMessage(self.src_mac, 2, Eap.RESPONSE,
                                      bytes.fromhex("3a535f0ee8c6b34fe714aa7dad9a0e15"),
                                      b"host1user")
        self.sm.event(EventMessageReceived(message, self.PORT_ID_MAC))

        self.assertEqual(self.sm.state, self.sm.AAA_IDLE)
        self.assertEqual(self.eap_output_queue.qsize(), 0)
        self.assertEqual(self.radius_output_queue.qsize(), 1)
        self.assertIsInstance(self.radius_output_queue.get_nowait()[0], Md5ChallengeMessage)

    @check_counters(expected_auth_counter=1)
    def test_success2(self):
        self.test_md5_challenge_response()

        message = SuccessMessage(self.src_mac, 3)
        self.sm.event(EventRadiusMessageReceived(message, None))

        self.assertEqual(self.sm.state, self.sm.SUCCESS2)
        self.assertEqual(self.eap_output_queue.qsize(), 1)
        self.assertIsInstance(self.eap_output_queue.get_nowait()[0], SuccessMessage)
        self.assertEqual(self.radius_output_queue.qsize(), 0)

    @check_counters(expected_auth_counter=2)
    def test_two_success2(self):
        self.test_success2()

        expiry_job = self.sm.session_timeout_job
        self.assertFalse(expiry_job.cancelled())

        self.test_success2()
        self.assertTrue(expiry_job.cancelled())

        expiry_job = self.sm.session_timeout_job
        self.assertFalse(expiry_job.cancelled())

    @check_counters(expected_auth_counter=2, expected_logoff_counter=1)
    def test_logoff2(self):
        """Test logoff from success2 state."""
        self.test_success2()

        # test the second logon expires the firsts session timeout event.
        expiry_job = self.sm.session_timeout_job
        self.assertFalse(expiry_job.cancelled())

        message = EapolLogoffMessage(self.src_mac)
        self.sm.event(EventRadiusMessageReceived(message, self.PORT_ID_MAC))

        self.assertTrue(expiry_job.cancelled())

        self.assertEqual(self.sm.state, self.sm.LOGOFF2)
        self.assertEqual(self.eap_output_queue.qsize(), 0)
        self.assertEqual(self.radius_output_queue.qsize(), 0)

        self.test_success2()

    @check_counters
    def test_logoff_from_idle2(self):
        """Test logoff from middle of authentication. should be ignored"""
        self.test_md5_challenge_request()

        message = EapolLogoffMessage(self.src_mac)
        self.sm.event(EventRadiusMessageReceived(message, self.PORT_ID_MAC))

        # should be in same state as when test_md5_challenge_request returned.

        self.assertEqual(self.sm.state, self.sm.IDLE2)
        self.assertEqual(self.eap_output_queue.qsize(), 0)
        self.assertEqual(self.radius_output_queue.qsize(), 0)

    @check_counters(expected_failure_counter=1)
    def test_failure2(self):
        self.test_md5_challenge_response()
        message = FailureMessage(self.src_mac, 3)
        self.sm.event(EventRadiusMessageReceived(message, None))

        self.assertEqual(self.sm.state, self.sm.FAILURE2)
        self.assertEqual(self.eap_output_queue.qsize(), 1)
        self.assertIsInstance(self.eap_output_queue.get_nowait()[0], FailureMessage)
        self.assertEqual(self.radius_output_queue.qsize(), 0)

    @check_counters
    def test_discard2(self):
        request = self.test_md5_challenge_request()

        message = Md5ChallengeMessage(self.src_mac, request.message_id + 10, Eap.RESPONSE,
                                      bytes.fromhex("3a535f0ee8c6b34fe714aa7dad9a0e15"),
                                      b"host1user")
        self.sm.event(EventMessageReceived(message, self.PORT_ID_MAC))
        self.assertEqual(self.sm.state, self.sm.IDLE2)
        self.assertEqual(self.eap_output_queue.qsize(), 0)
        self.assertEqual(self.radius_output_queue.qsize(), 0)

    @check_counters
    def test_discard(self):
        message = self.test_eap_start()
        # Make a message that will be discarded (id here is not sequential)
        message = IdentityMessage(self.src_mac, message.message_id + 10, Eap.RESPONSE, "host1user")
        self.sm.event(EventMessageReceived(message, self.PORT_ID_MAC))
        self.assertEqual(self.sm.state, self.sm.IDLE)

        self.assertEqual(self.eap_output_queue.qsize(), 0)
        self.assertEqual(self.radius_output_queue.qsize(), 0)

    @check_counters
    def test_ttls_request(self):
        self.test_md5_challenge_request()
        self.assertEqual(self.eap_output_queue.qsize(), 0)
        self.assertEqual(self.radius_output_queue.qsize(), 0)

        message = LegacyNakMessage(self.src_mac, 2, Eap.RESPONSE, 21)
        self.sm.event(EventMessageReceived(message, self.PORT_ID_MAC))
        self.assertEqual(self.sm.state, self.sm.AAA_IDLE)
        self.assertEqual(self.eap_output_queue.qsize(), 0)
        self.assertEqual(self.radius_output_queue.qsize(), 1)

        message = TtlsMessage(self.src_mac, 3, Eap.REQUEST, 0x20, b'')
        self.sm.event(EventRadiusMessageReceived(message, State.create(b"more random state")))
        self.assertEqual(self.sm.state, self.sm.IDLE2)
        self.assertEqual(self.eap_output_queue.qsize(), 1)
        self.assertEqual(self.radius_output_queue.qsize(), 1)

        message = TtlsMessage(self.src_mac, 3, Eap.RESPONSE, 0x00,
                              bytes.fromhex(
                                  '16030101280100012403032c36dbf8ee16b94b28efdb8c5603e07823f9b716557b5ef2624b026daea115760000aac030c02cc028c024c014c00a00a500a300a1009f006b006a0069006800390038003700360088008700860085c032c02ec02ac026c00fc005009d003d00350084c02fc02bc027c023c013c00900a400a200a0009e00670040003f003e0033003200310030009a0099009800970045004400430042c031c02dc029c025c00ec004009c003c002f00960041c011c007c00cc00200050004c012c008001600130010000dc00dc003000a00ff01000051000b000403000102000a001c001a00170019001c001b0018001a0016000e000d000b000c0009000a000d0020001e060106020603050105020503040104020403030103020303020102020203000f000101'))
        self.sm.event(EventMessageReceived(message, self.PORT_ID_MAC))
        self.assertEqual(self.sm.state, self.sm.AAA_IDLE)
        self.assertEqual(self.eap_output_queue.qsize(), 1)
        self.assertEqual(self.radius_output_queue.qsize(), 2)

        message = TtlsMessage(self.src_mac, 4, Eap.REQUEST, 0x00, bytes.fromhex(
            ' '))
        self.sm.event(EventRadiusMessageReceived(message, State.create(b"some more random state")))
        self.assertEqual(self.sm.state, self.sm.IDLE2)
        self.assertEqual(self.eap_output_queue.qsize(), 2)
        self.assertEqual(self.radius_output_queue.qsize(), 2)

    @check_counters(expected_auth_counter=1, expected_logoff_counter=1)
    def test_deauth_timer(self):
        self.sm.DEFAULT_SESSION_TIMEOUT = 2
        self.test_success2()

        self.timer_scheduler.run_jobs()
        self.assertEqual(self.sm.state, self.sm.LOGOFF2)

    @check_counters(expected_auth_counter=2, expected_logoff_counter=0)
    def test_port_flap(self):
        """Test logoff from success2 state."""
        self.test_success2()
        # Put port down
        self.sm.event(EventPortStatusChange(False))
        self.sm.event(EventPortStatusChange(True))

        self.assertEqual(self.sm.state, self.sm.NO_STATE)
        self.assertFalse(self.sm.aaa_success)
        self.assertFalse(self.sm.eap_success)

        self.assertEqual(self.eap_output_queue.qsize(), 0)
        self.assertEqual(self.radius_output_queue.qsize(), 0)

        self.test_success2()
