diff --git a/monkey/infection_monkey/exploit/tools/payload_parsing.py b/monkey/infection_monkey/exploit/tools/payload_parsing.py deleted file mode 100644 index ef2e11426..000000000 --- a/monkey/infection_monkey/exploit/tools/payload_parsing.py +++ /dev/null @@ -1,64 +0,0 @@ -import textwrap - - -class Payload(object): - """ - Class for defining and parsing a payload (commands with prefixes/suffixes) - """ - - def __init__(self, command, prefix="", suffix=""): - self.command = command - self.prefix = prefix - self.suffix = suffix - - def get_payload(self, command=""): - """ - Returns prefixed and suffixed command (payload) - :param command: Command to suffix/prefix. If no command is passed than objects' property - is used - :return: prefixed and suffixed command (full payload) - """ - if not command: - command = self.command - return "{}{}{}".format(self.prefix, command, self.suffix) - - -class LimitedSizePayload(Payload): - """ - Class for defining and parsing commands/payloads - """ - - def __init__(self, command, max_length, prefix="", suffix=""): - """ - :param command: command - :param max_length: max length that payload(prefix + command + suffix) can have - :param prefix: commands prefix - :param suffix: commands suffix - """ - super(LimitedSizePayload, self).__init__(command, prefix, suffix) - self.max_length = max_length - - def is_suffix_and_prefix_too_long(self): - return self.payload_is_too_long(self.suffix + self.prefix) - - def split_into_array_of_smaller_payloads(self): - if self.is_suffix_and_prefix_too_long(): - raise Exception( - "Can't split command into smaller sub-commands because commands' prefix and " - "suffix already " - "exceeds required length of command." - ) - - elif self.command == "": - return [self.prefix + self.suffix] - wrapper = textwrap.TextWrapper( - drop_whitespace=False, width=self.get_max_sub_payload_length() - ) - commands = [self.get_payload(part) for part in wrapper.wrap(self.command)] - return commands - - def get_max_sub_payload_length(self): - return self.max_length - len(self.prefix) - len(self.suffix) - - def payload_is_too_long(self, command): - return len(command) >= self.max_length diff --git a/monkey/tests/unit_tests/infection_monkey/exploit/tools/test_payload.py b/monkey/tests/unit_tests/infection_monkey/exploit/tools/test_payload.py deleted file mode 100644 index 2656a7ada..000000000 --- a/monkey/tests/unit_tests/infection_monkey/exploit/tools/test_payload.py +++ /dev/null @@ -1,38 +0,0 @@ -from unittest import TestCase - -from infection_monkey.exploit.tools.payload_parsing import LimitedSizePayload, Payload - - -class TestPayload(TestCase): - def test_get_payload(self): - test_str1 = "abc" - test_str2 = "atc" - payload = Payload(command="b", prefix="a", suffix="c") - assert payload.get_payload() == test_str1 and payload.get_payload("t") == test_str2 - - def test_is_suffix_and_prefix_too_long(self): - pld_fail = LimitedSizePayload("b", 2, "a", "c") - pld_success = LimitedSizePayload("b", 3, "a", "c") - assert ( - pld_fail.is_suffix_and_prefix_too_long() - and not pld_success.is_suffix_and_prefix_too_long() - ) - - def test_split_into_array_of_smaller_payloads(self): - test_str1 = "123456789" - pld1 = LimitedSizePayload(test_str1, max_length=16, prefix="prefix", suffix="suffix") - array1 = pld1.split_into_array_of_smaller_payloads() - test1 = bool( - array1[0] == "prefix1234suffix" - and array1[1] == "prefix5678suffix" - and array1[2] == "prefix9suffix" - ) - - test_str2 = "12345678" - pld2 = LimitedSizePayload(test_str2, max_length=16, prefix="prefix", suffix="suffix") - array2 = pld2.split_into_array_of_smaller_payloads() - test2 = bool( - array2[0] == "prefix1234suffix" and array2[1] == "prefix5678suffix" and len(array2) == 2 - ) - - assert test1 and test2