3
0

TestLibrary.py 4.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141
  1. import hashlib
  2. import os
  3. import random as rnd
  4. import ID2TLib.Utility as Util
  5. # Directory of test resource files
  6. test_resource_dir = Util.TEST_DIR
  7. # Path to reference pcap
  8. test_pcap = Util.TEST_DIR + "reference_1998.pcap"
  9. # Several ips in the reference pcap
  10. test_pcap_ips = ["10.0.2.15", "52.85.173.182"]
  11. # Empty array for testing purposes
  12. test_pcap_empty = []
  13. """
  14. helper functions for ID2TAttackTest
  15. """
  16. def get_sha256(file):
  17. """
  18. Generates a sha256 checksum from file
  19. :param file: absolute path to file
  20. :return: sha256 checksum
  21. """
  22. sha = hashlib.sha256()
  23. with open(file, 'rb') as f:
  24. while True:
  25. data = f.read(0x100000)
  26. if not data:
  27. break
  28. sha.update(data)
  29. f.close()
  30. return sha.hexdigest()
  31. def clean_up(controller):
  32. """
  33. Removes the output files from a given controller
  34. :param controller: controller which created output files
  35. """
  36. os.remove(controller.pcap_dest_path)
  37. os.remove(controller.label_manager.label_file_path)
  38. def rename_test_result_files(controller, caller_function: str, attack_sub_dir=False, test_sub_dir=False):
  39. """
  40. :param controller: controller which created output files
  41. :param caller_function: the name of the function which called the generic test
  42. :param attack_sub_dir: create sub-directory for each attack-class if True
  43. :param test_sub_dir: create sub-directory for each test-function/case if True
  44. """
  45. tmp_path_tuple = controller.pcap_dest_path.rpartition("_")
  46. result_pcap_path = tmp_path_tuple[0] + tmp_path_tuple[1] + caller_function + "_" + tmp_path_tuple[2]
  47. tmp_label_path_tuple = controller.label_manager.label_file_path.rpartition("_")
  48. tmp_path_tuple = tmp_label_path_tuple[0].rpartition("_")
  49. result_labels_path = tmp_path_tuple[0] + tmp_path_tuple[1] + caller_function + "_" + tmp_path_tuple[2]
  50. result_labels_path = result_labels_path + tmp_label_path_tuple[1] + tmp_label_path_tuple[2]
  51. if attack_sub_dir:
  52. caller_attack = caller_function.replace("test_", "").partition("_")[0]
  53. tmp_dir_tuple = result_pcap_path.rpartition("/")
  54. result_dir = tmp_dir_tuple[0] + tmp_dir_tuple[1] + caller_attack + "/"
  55. result_pcap_path = result_dir + tmp_dir_tuple[2]
  56. os.makedirs(result_dir, exist_ok=True)
  57. tmp_dir_tuple = result_labels_path.rpartition("/")
  58. result_labels_path = result_dir + tmp_dir_tuple[2]
  59. if test_sub_dir:
  60. tmp_dir_tuple = result_pcap_path.rpartition("/")
  61. result_dir = tmp_dir_tuple[0] + tmp_dir_tuple[1] + (caller_function.replace("test_", "")) + "/"
  62. result_pcap_path = result_dir + tmp_dir_tuple[2]
  63. os.makedirs(result_dir, exist_ok=True)
  64. tmp_dir_tuple = result_labels_path.rpartition("/")
  65. result_labels_path = result_dir + tmp_dir_tuple[2]
  66. os.rename(controller.pcap_dest_path, result_pcap_path)
  67. controller.pcap_dest_path = result_pcap_path
  68. os.rename(controller.label_manager.label_file_path, result_labels_path)
  69. controller.label_manager.label_file_path = result_labels_path
  70. """
  71. function patches for unittests
  72. FYI: the parameters below, which are not used are needed to mock the mentioned function/method correctly
  73. """
  74. def get_bytes(count, ignore):
  75. """
  76. unittest patch for get_rnd_bytes (ID2TLib.Utility.py)
  77. :param count: count of requested bytes
  78. :param ignore: <not used>
  79. :return: a count of As
  80. """
  81. return b'A' * count
  82. def get_x86_nop(count, side_effect_free, char_filter):
  83. """
  84. unittest patch for get_rnd_x86_nop (ID2TLib.Utility.py)
  85. :param count: count of requested nops
  86. :param side_effect_free: <not used>
  87. :param char_filter: <not used>
  88. :return: a count of \x90
  89. """
  90. return b'\x90' * count
  91. def get_attacker_config(ip_source_list, ipAddress: str):
  92. """
  93. unittest patch for get_attacker_config (ID2TLib.Utility.py)
  94. :param ip_source_list: List of source IPs
  95. :param ipAddress: The IP address of the attacker
  96. :return: A tuple consisting of (port, ttlValue)
  97. """
  98. next_port = rnd.randint(0, 2 ** 16 - 1)
  99. ttl = rnd.randint(1, 255)
  100. return next_port, ttl
  101. def write_attack_pcap(self, packets: list, append_flag: bool = False, destination_path: str = None):
  102. """
  103. temporal efficiency test patch for write_attack_pcap (Attack.BaseAttack.py)
  104. :return: The path to a dummy pcap file.
  105. """
  106. # TODO: find another solution - copying influences efficiency tests
  107. os.system("cp " + test_pcap + " " + test_resource_dir + "dummy.pcap")
  108. return test_resource_dir + 'dummy.pcap'