diff --git a/monkey/infection_monkey/exploit/tools/helpers.py b/monkey/infection_monkey/exploit/tools/helpers.py index b5fce82f7..e26f6ff01 100644 --- a/monkey/infection_monkey/exploit/tools/helpers.py +++ b/monkey/infection_monkey/exploit/tools/helpers.py @@ -46,19 +46,19 @@ def build_monkey_commandline_explicitly(parent=None, tunnel=None, server=None, d cmdline = "" if parent is not None: - cmdline += " -p " + parent + cmdline += f" -p {parent}" if tunnel is not None: - cmdline += " -t " + tunnel + cmdline += f" -t {tunnel}" if server is not None: - cmdline += " -s " + server + cmdline += f" -s {server}" if depth is not None: - if depth < 0: + if int(depth) < 0: depth = 0 cmdline += f" -d {depth}" if location is not None: cmdline += f" -l {location}" if vulnerable_port is not None: - cmdline += f" -vp {str(vulnerable_port)}" + cmdline += f" -vp {vulnerable_port}" return cmdline diff --git a/monkey/infection_monkey/exploit/tools/test_helpers.py b/monkey/infection_monkey/exploit/tools/test_helpers.py new file mode 100644 index 000000000..5d7dd422d --- /dev/null +++ b/monkey/infection_monkey/exploit/tools/test_helpers.py @@ -0,0 +1,28 @@ +import unittest + +from infection_monkey.exploit.tools.helpers import build_monkey_commandline_explicitly + + +class TestHelpers(unittest.TestCase): + + def test_build_monkey_commandline_explicitly(self): + test1 = " -p 101010 -t 10.10.101.10 -s 127.127.127.127:5000 -d 0 -l C:\\windows\\abc -vp 80" + result1 = build_monkey_commandline_explicitly(101010, + "10.10.101.10", + "127.127.127.127:5000", + 0, + "C:\\windows\\abc", + 80) + + test2 = " -p parent -s 127.127.127.127:5000 -d 0 -vp 80" + result2 = build_monkey_commandline_explicitly(parent="parent", + server="127.127.127.127:5000", + depth="0", + vulnerable_port="80") + + self.assertEqual(test1, result1) + self.assertEqual(test2, result2) + + +if __name__ == '__main__': + unittest.main()