forked from p15670423/monkey
Agent: Remove disused "Payload" classes
This commit is contained in:
parent
a54eca96ba
commit
819262ef73
|
@ -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
|
|
@ -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
|
Loading…
Reference in New Issue