Browse Source

Modified check_payload_len() to raise an exception instead of calling exit()

Stefan Schmidt 6 years ago
parent
commit
92dcc5524c
2 changed files with 5 additions and 8 deletions
  1. 3 3
      code/ID2TLib/Utility.py
  2. 2 5
      code/Test/test_Utility.py

+ 3 - 3
code/ID2TLib/Utility.py

@@ -247,7 +247,7 @@ def get_rnd_bytes(count=1, ignore=None):
     return result
 
 
-def check_payload_len(payload_len: int, limit: int):
+def check_payload_len(payload_len: int, limit: int) -> None:
     """
     Checks if the len of the payload exceeds a given limit
 
@@ -256,8 +256,8 @@ def check_payload_len(payload_len: int, limit: int):
     """
 
     if payload_len > limit:
-        print("\nCustom payload too long: ", payload_len, " bytes. Should be a maximum of ", limit, " bytes.")
-        exit(1)
+        raise ValueError("Custom payload too long: " + str(payload_len) +
+                         " bytes. Should be a maximum of " + str(limit) + " bytes.")
 
 
 def get_bytes_from_file(filepath):

+ 2 - 5
code/Test/test_Utility.py

@@ -201,14 +201,11 @@ class TestUtility(unittest.TestCase):
         self.assertEqual(Utility.handle_most_used_outputs(test_input), 0)
 
     def test_check_payload_len_exceeded(self):
-        with self.assertRaises(SystemExit):
+        with self.assertRaises(ValueError):
             Utility.check_payload_len(10, 5)
 
     def test_check_payload_len_valid(self):
-        try:
-            Utility.check_payload_len(5, 10)
-        except SystemExit:
-            self.fail()
+        Utility.check_payload_len(5, 10)
 
     def test_remove_generic_ending_attack(self):
         self.assertEqual(Utility.remove_generic_ending("someattack"), "some")