TestLibrary.py 3.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110
  1. import os
  2. import hashlib
  3. from definitions import ROOT_DIR
  4. test_resource_dir = ROOT_DIR + "/../resources/test"
  5. test_pcap = ROOT_DIR + "/../resources/test/reference_1998.pcap"
  6. test_pcap_ips = ["10.0.2.15", "52.85.173.182"]
  7. test_pcap_empty = []
  8. """
  9. helper functions for generic_test
  10. """
  11. def get_sha256(file):
  12. """
  13. Generates a sha256 checksum from file
  14. :param file: absolute path to file
  15. :return: sha256 checksum
  16. """
  17. sha = hashlib.sha256()
  18. with open(file, 'rb') as f:
  19. while True:
  20. data = f.read(0x100000)
  21. if not data:
  22. break
  23. sha.update(data)
  24. f.close()
  25. return sha.hexdigest()
  26. def clean_up(controller):
  27. """
  28. Removes the output files from a given controller
  29. :param controller: controller which created output files
  30. """
  31. os.remove(controller.pcap_dest_path)
  32. os.remove(controller.label_manager.label_file_path)
  33. """
  34. function patches for unittests
  35. """
  36. def get_bytes(count, ignore):
  37. """
  38. unittest patch for get_rnd_bytes (ID2TLib.Utility.py)
  39. :param count: count of requested bytes
  40. :param ignore: <not used>
  41. :return: a count of As
  42. """
  43. return b'A' * count
  44. def get_x86_nop(count, side_effect_free, char_filter):
  45. """
  46. unittest patch for get_rnd_x86_nop (ID2TLib.Utility.py)
  47. :param count: count of requested nops
  48. :param side_effect_free: <not used>
  49. :param char_filter: <not used>
  50. :return: a count of \x90
  51. """
  52. return b'\x90' * count
  53. def rename_test_result_files(controller, caller_function: str, attack_sub_dir=False, test_sub_dir=False):
  54. """
  55. :param controller: controller which created output files
  56. :param caller_function: the name of the function which called the generic test
  57. :param attack_sub_dir: create sub-directory for each attack-class if True
  58. :param test_sub_dir: create sub-directory for each test-function/case if True
  59. """
  60. tmp_path_tuple = controller.pcap_dest_path.rpartition("_")
  61. result_pcap_path = tmp_path_tuple[0] + tmp_path_tuple[1] + caller_function + "_" + tmp_path_tuple[2]
  62. tmp_label_path_tuple = controller.label_manager.label_file_path.rpartition("_")
  63. tmp_path_tuple = tmp_label_path_tuple[0].rpartition("_")
  64. result_labels_path = tmp_path_tuple[0] + tmp_path_tuple[1] + caller_function + "_" + tmp_path_tuple[2]
  65. result_labels_path = result_labels_path + tmp_label_path_tuple[1] + tmp_label_path_tuple[2]
  66. if attack_sub_dir:
  67. caller_attack = caller_function.replace("test_", "").partition("_")[0]
  68. tmp_dir_tuple = result_pcap_path.rpartition("/")
  69. result_dir = tmp_dir_tuple[0] + tmp_dir_tuple[1] + caller_attack + "/"
  70. result_pcap_path = result_dir + tmp_dir_tuple[2]
  71. os.makedirs(result_dir, exist_ok=True)
  72. tmp_dir_tuple = result_labels_path.rpartition("/")
  73. result_labels_path = result_dir + tmp_dir_tuple[2]
  74. if test_sub_dir:
  75. tmp_dir_tuple = result_pcap_path.rpartition("/")
  76. result_dir = tmp_dir_tuple[0] + tmp_dir_tuple[1] + (caller_function.replace("test_", "")) + "/"
  77. result_pcap_path = result_dir + tmp_dir_tuple[2]
  78. os.makedirs(result_dir, exist_ok=True)
  79. tmp_dir_tuple = result_labels_path.rpartition("/")
  80. result_labels_path = result_dir + tmp_dir_tuple[2]
  81. os.rename(controller.pcap_dest_path, result_pcap_path)
  82. controller.pcap_dest_path = result_pcap_path
  83. os.rename(controller.label_manager.label_file_path, result_labels_path)
  84. controller.label_manager.label_file_path = result_labels_path