import sched
import time
from queue import Queue
import unittest

from netils import build_byte_string

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


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??

    def setUp(self):
        self.eap_output_queue = Queue()
        self.radius_output_queue = Queue()
        self.timer_scheduler = sched.scheduler(time.time, time.sleep)
        self.src_mac = MacAddress.from_string("00:12:34:56:78:90")
        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)
        self.MAX_RETRANSMITS = 3
        self.sm.MAX_RETRANS = self.MAX_RETRANSMITS
        self.sm.DEFAULT_TIMEOUT = 0.1
        self.sm.portEnabled = True
        self.sm.eapRestart = True

    def auth_handler(self, client_mac, port_id_mac):
        print('Successful auth from MAC %s' % str(client_mac))

    def failure_handler(self, client_mac, port_id_mac):
        print('failure from MAC %s' % str(client_mac))

    def logoff_handler(self, client_mac, port_id_mac):
        print('logoff from MAC %s' % str(client_mac))

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

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

    def test_eap_restart(self):
        self.test_eap_start()
        message = EapolStartMessage(self.src_mac)
        self.sm.event(EventMessageReceived(message, None))
        self.assertEqual(self.sm.currentState, self.sm.IDLE)

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

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

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

        output0 = self.eap_output_queue.queue[0][0]
        output1 = self.eap_output_queue.queue[1][0]
        output2 = self.eap_output_queue.queue[2][0]
        self.assertIsInstance(output0, IdentityMessage)
        self.assertIsInstance(output1, IdentityMessage)
        self.assertEqual(output0, output1)
        self.assertEqual(output0, output2)
        self.assertEqual(self.radius_output_queue.qsize(), 0)

        self.assertEqual(self.sm.currentState, self.sm.TIMEOUT_FAILURE)
        self.assertEqual(self.MAX_RETRANSMITS, self.eap_output_queue.qsize())
        self.assertEqual(old_radius_count, self.radius_output_queue.qsize())

    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()

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

    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.eap_output_queue.queue.clear()
        old_radius_count = self.radius_output_queue.qsize()

        self.timer_scheduler.run()

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

    def test_disabled_state(self):
        """move to disabled and then from disabled"""
        # move to disabled. e.g. link down.
        self.test_eap_start()
        self.sm.event(EventPortStatusChange(False))
        self.assertEqual(self.sm.currentState, self.sm.DISABLED)

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

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

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

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

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

    def test_identity_response(self):
        self.test_eap_start()
        # input EapIdentityResponse
        # output EapIdentityResponse on radius_output_q
        _id = self.eap_output_queue.queue[0][0].message_id
        message = IdentityMessage(self.src_mac, _id, Eap.RESPONSE, "host1user")
        self.sm.event(EventMessageReceived(message, None))
        self.assertEqual(self.sm.currentState, self.sm.AAA_IDLE)

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

    def test_md5_challenge_request(self):
        self.test_identity_response()

        eap_message = Md5ChallengeMessage(self.src_mac, 2, Eap.REQUEST,
                                          build_byte_string("74d3db089b727d9cc5774599e4a32a29"), b"host1user")
        self.sm.event(EventRadiusMessageReceived(eap_message, None))
        self.assertEqual(self.sm.currentState, self.sm.IDLE2)

        self.assertEqual(self.eap_output_queue.qsize(), 2)
        self.assertIsInstance(self.eap_output_queue.queue[1][0], Md5ChallengeMessage)

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

    def test_md5_challenge_response(self):
        self.test_md5_challenge_request()

        message = Md5ChallengeMessage(self.src_mac, 2, Eap.RESPONSE,
                                      build_byte_string("3a535f0ee8c6b34fe714aa7dad9a0e15"), b"host1user")
        self.sm.event(EventMessageReceived(message, None))
        self.assertEqual(self.sm.currentState, self.sm.AAA_IDLE)
        self.assertEqual(self.eap_output_queue.qsize(), 2)
        self.assertEqual(self.radius_output_queue.qsize(), 2)
        self.assertIsInstance(self.radius_output_queue.queue[1][0], Md5ChallengeMessage)

    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.currentState, self.sm.SUCCESS2)
        self.assertEqual(self.eap_output_queue.qsize(), 3)
        self.assertIsInstance(self.eap_output_queue.queue[2][0], SuccessMessage)
        self.assertEqual(self.radius_output_queue.qsize(), 2)

    def test_logoff2(self):
        """Test logoff from success2 state."""
        self.test_success2()

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

        self.assertEqual(self.sm.currentState, self.sm.LOGOFF2)
        self.assertEqual(self.eap_output_queue.qsize(), 3)
        self.assertEqual(self.radius_output_queue.qsize(), 2)

        self.eap_output_queue.queue.clear()
        self.radius_output_queue.queue.clear()
        self.test_success2()

    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, None))

        # should be in same state as when test_md5_challenge_request returned.
        self.assertEqual(self.sm.currentState, self.sm.IDLE2)
        self.assertEqual(self.eap_output_queue.qsize(), 2)
        self.assertIsInstance(self.eap_output_queue.queue[1][0], Md5ChallengeMessage)
        self.assertEqual(self.radius_output_queue.qsize(), 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.currentState, self.sm.FAILURE2)
        self.assertEqual(self.eap_output_queue.qsize(), 3)
        self.assertIsInstance(self.eap_output_queue.queue[2][0], FailureMessage)
        self.assertEqual(self.radius_output_queue.qsize(), 2)

    def test_discard2(self):
        self.test_md5_challenge_request()

        message = Md5ChallengeMessage(self.src_mac, 222, Eap.RESPONSE,
                                      build_byte_string("3a535f0ee8c6b34fe714aa7dad9a0e15"), b"host1user")
        self.sm.event(EventMessageReceived(message, None))
        self.assertEqual(self.sm.currentState, self.sm.IDLE2)
        self.assertEqual(self.eap_output_queue.qsize(), 2)
        self.assertEqual(self.radius_output_queue.qsize(), 1)

    def test_discard(self):
        self.test_eap_start()

        message = IdentityMessage(self.src_mac, 40, Eap.RESPONSE, "host1user")
        self.sm.event(EventMessageReceived(message, None))
        self.assertEqual(self.sm.currentState, self.sm.IDLE)

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

    def test_ttls_request(self):
        self.test_md5_challenge_request()

        message = LegacyNakMessage(self.src_mac, 2, Eap.RESPONSE, 21)
        self.sm.event(EventMessageReceived(message, None))
        self.assertEqual(self.sm.currentState, self.sm.AAA_IDLE)
        self.assertEqual(self.eap_output_queue.qsize(), 2)
        self.assertEqual(self.radius_output_queue.qsize(), 2)

        message = TtlsMessage(self.src_mac, 3, Eap.REQUEST, 0x20, b'')
        self.sm.event(EventRadiusMessageReceived(message, None))
        self.assertEqual(self.sm.currentState, self.sm.IDLE2)
        self.assertEqual(self.eap_output_queue.qsize(), 3)
        self.assertEqual(self.radius_output_queue.qsize(), 2)

        message = TtlsMessage(self.src_mac, 3, Eap.RESPONSE, 0x00,
                              build_byte_string('16030101280100012403032c36dbf8ee16b94b28efdb8c5603e07823f9b716557b5ef2624b026daea115760000aac030c02cc028c024c014c00a00a500a300a1009f006b006a0069006800390038003700360088008700860085c032c02ec02ac026c00fc005009d003d00350084c02fc02bc027c023c013c00900a400a200a0009e00670040003f003e0033003200310030009a0099009800970045004400430042c031c02dc029c025c00ec004009c003c002f00960041c011c007c00cc00200050004c012c008001600130010000dc00dc003000a00ff01000051000b000403000102000a001c001a00170019001c001b0018001a0016000e000d000b000c0009000a000d0020001e060106020603050105020503040104020403030103020303020102020203000f000101'))
        self.sm.event(EventMessageReceived(message, None))
        self.assertEqual(self.sm.currentState, self.sm.AAA_IDLE)
        self.assertEqual(self.eap_output_queue.qsize(), 3)
        self.assertEqual(self.radius_output_queue.qsize(), 3)

        message = TtlsMessage(self.src_mac, 4, Eap.REQUEST, 0x00, build_byte_string(
            ' '))
        self.sm.event(EventRadiusMessageReceived(message, None))
        self.assertEqual(self.sm.currentState, self.sm.IDLE2)
        self.assertEqual(self.eap_output_queue.qsize(), 4)
        self.assertEqual(self.radius_output_queue.qsize(), 3)
