#! /usr/bin/env python3
import exploit
import argparse
import unittest

RELEASE_PORT = 8080
DEBUG_PORT = 8081


def __get_activation_key__() -> bytes:
    global activation_key_file
    assert (activation_key_file is not None)
    with open(activation_key_file, "rb") as f:
        data = f.read()
    return data


class ExploitTest(unittest.TestCase):

    def __check_extract_activation_key__(self, is_debug: bool):
        port = DEBUG_PORT if is_debug else RELEASE_PORT
        exploit.PORT = port
        key = exploit.extract_premium_key(is_debug)
        self.assertEqual(key, __get_activation_key__())

    def test_extract_activation_key_debug(self):
        self.__check_extract_activation_key__(is_debug=True)

    def test_extract_activation_key_release(self):
        self.__check_extract_activation_key__(is_debug=False)

    def __check_get_flag__(self, is_debug: bool):
        port = DEBUG_PORT if is_debug else RELEASE_PORT
        exploit.PORT = port
        p = exploit.connect(__get_activation_key__(), True)
        flag = exploit.get_flag(p, is_debug)
        self.assertRegex(flag, "flag_[0-9a-f]{32}")

    def test_get_flag_debug(self):
        self.__check_get_flag__(True)

    def test_get_flag_release(self):
        self.__check_get_flag__(False)

    def __check_combined__(self, is_debug: bool):
        port = DEBUG_PORT if is_debug else RELEASE_PORT
        exploit.PORT = port
        activation_key = exploit.extract_premium_key(is_debug)
        p = exploit.connect(activation_key, True)
        flag = exploit.get_flag(p, is_debug)
        self.assertRegex(flag, "flag_[0-9a-f]{32}")

    def test_combined_debug(self):
        self.__check_combined__(True)

    def test_combined_release(self):
        self.__check_combined__(False)


if __name__ == "__main__":
    global activation_key_file
    arg_parser = argparse.ArgumentParser(prog="Exploit CI tester")
    arg_parser.add_argument("-f", "--file", type=str, required=True, help="File containing the activation key.")
    args = arg_parser.parse_args()
    activation_key_file = args.file

    unittest.main(argv=[""])