update SMBScanAttack parameters

code/Attack/SMBScanAttack.py

@@ -61,14 +61,12 @@ class SMBScanAttack(BaseAttack.BaseAttack):
         # (values are overwritten if user specifies them)
         most_used_ip_address = self.statistics.get_most_used_ip_address()
 
+        self.add_param_value(atkParam.Parameter.TARGET_COUNT, 200)
+
         self.add_param_value(atkParam.Parameter.IP_SOURCE, most_used_ip_address)
         self.add_param_value(atkParam.Parameter.IP_SOURCE_RANDOMIZE, 'False')
         self.add_param_value(atkParam.Parameter.MAC_SOURCE, ft.partial(self.statistics.get_mac_address,
                                                                        most_used_ip_address))
-
-        self.add_param_value(atkParam.Parameter.TARGET_COUNT, 200)
-        self.add_param_value(atkParam.Parameter.IP_DESTINATION, "1.1.1.1")
-
         self.add_param_value(atkParam.Parameter.PORT_SOURCE, ft.partial(rnd.randint, 1024, 65535))
         self.add_param_value(atkParam.Parameter.PORT_SOURCE_RANDOMIZE, 'True')
         self.add_param_value(atkParam.Parameter.PACKET_LIMIT_PER_SECOND,
@@ -87,12 +85,21 @@ class SMBScanAttack(BaseAttack.BaseAttack):
         self.add_param_value(atkParam.Parameter.INJECT_PPS, 0)
 
         self.add_param_value(atkParam.Parameter.HOSTING_PERCENTAGE, 0.5)
-        self.add_param_value(atkParam.Parameter.HOSTING_IP, "1.1.1.1")
         self.add_param_value(atkParam.Parameter.HOSTING_VERSION, ft.partial(SMBLib.get_smb_version,
                                                                             platform=self.host_os))
         self.add_param_value(atkParam.Parameter.SOURCE_PLATFORM, Util.get_rnd_os)
         self.add_param_value(atkParam.Parameter.PROTOCOL_VERSION, "1")
 
+        # Check smb version
+        smb_version = self.get_param_value(atkParam.Parameter.PROTOCOL_VERSION)
+        print(smb_version)
+        if not SMBLib.is_valid_version(smb_version):
+            exit(1)
+        hosting_version = self.get_param_value(atkParam.Parameter.HOSTING_VERSION)
+        print(hosting_version)
+        if not SMBLib.is_valid_version(hosting_version):
+            exit(1)
+
         ip_source = self.get_param_value(atkParam.Parameter.IP_SOURCE)
 
         dest_ip_count = self.get_param_value(atkParam.Parameter.TARGET_COUNT)
@@ -104,7 +111,7 @@ class SMBScanAttack(BaseAttack.BaseAttack):
         ip_destinations = self.get_param_value(atkParam.Parameter.IP_DESTINATION)
         if isinstance(ip_destinations, list):
             dest_ip_count = dest_ip_count - len(ip_destinations)
-        elif ip_destinations is not "1.1.1.1":
+        elif ip_destinations is not None:
             dest_ip_count = dest_ip_count - 1
             ip_destinations = [ip_destinations]
         else:
@@ -129,12 +136,13 @@ class SMBScanAttack(BaseAttack.BaseAttack):
         hosting_ip = self.get_param_value(atkParam.Parameter.HOSTING_IP)
         if isinstance(hosting_ip, list):
             rnd_ip_count = rnd_ip_count - len(hosting_ip)
-        elif hosting_ip is not "1.1.1.1":
+        elif hosting_ip is not None:
             rnd_ip_count = rnd_ip_count - 1
             hosting_ip = [hosting_ip]
         else:
             hosting_ip = []
 
+        hosting_ip = hosting_ip + ip_destinations[:int(rnd_ip_count)]
         self.add_param_value(atkParam.Parameter.HOSTING_IP, hosting_ip)
 
     def generate_attack_packets(self):
@@ -161,6 +169,10 @@ class SMBScanAttack(BaseAttack.BaseAttack):
 
         hosting_ip = self.get_param_value(atkParam.Parameter.HOSTING_IP)
 
+        # Get smb version
+        smb_version = self.get_param_value(atkParam.Parameter.PROTOCOL_VERSION)
+        hosting_version = self.get_param_value(atkParam.Parameter.HOSTING_VERSION)
+
         # Shuffle targets
         rnd.shuffle(ip_destinations)
 
@@ -168,13 +180,6 @@ class SMBScanAttack(BaseAttack.BaseAttack):
         mac_source = self.get_param_value(atkParam.Parameter.MAC_SOURCE)
         mac_dest = self.get_param_value(atkParam.Parameter.MAC_DESTINATION)
 
-        # Check smb version
-        smb_version = self.get_param_value(atkParam.Parameter.PROTOCOL_VERSION)
-        if smb_version not in SMBLib.smb_versions:
-            SMBLib.invalid_smb_version(smb_version)
-        hosting_version = self.get_param_value(atkParam.Parameter.HOSTING_VERSION)
-        if hosting_version not in SMBLib.smb_versions:
-            SMBLib.invalid_smb_version(hosting_version)
         # Check source platform
         src_platform = self.get_param_value(atkParam.Parameter.SOURCE_PLATFORM).lower()
         self.packets = []

code/Core/AttackController.py

@@ -216,4 +216,4 @@ class AttackController:
         :return: None
         """
         for param_key, param_value in params.items():
-            self.current_attack.add_param_value(param_key, param_value)
+            self.current_attack.add_param_value(param_key, param_value, user_specified=True)

code/ID2TLib/SMBLib.py

@@ -97,12 +97,13 @@ def get_smb_platform_data(platform: str, timestamp: float):
     return server_guid, security_blob, capabilities, data_size, server_start_time
 
 
-def invalid_smb_version(version: str):
+def is_valid_version(version: str) -> bool:
     """
-    Prints an error and exits
-
-    :param version: the invalid SMB
+    :param version: the SMB version
+    :return: True if version is valid, False if otherwise.
     """
-    print("\nInvalid smb version: " + version +
-          "\nPlease select one of the following versions: ", smb_versions)
-    exit(1)
+    if version not in smb_versions:
+        print("\nInvalid smb version: " + version +
+              "\nPlease select one of the following versions: ", smb_versions)
+        return False
+    return True

code/Test/test_SMBLib.py

@@ -49,5 +49,4 @@ class TestSMBLib(unittest.TestCase):
         self.assertTrue(time <= Utility.get_filetime_format(100))
 
     def test_invalid_smb_version(self):
-        with self.assertRaises(SystemExit):
-            SMBLib.invalid_smb_version("abc")
+        self.assertEqual(SMBLib.is_valid_version("abc"), False)

code/Test/test_SMBScanAttack.py

@@ -2,10 +2,10 @@ import unittest.mock as mock
 
 import Test.ID2TAttackTest as Test
 
-sha_default = '53252b4235c6fa0cdfd8966a3057c1e4d9f5042c17595951f3019f41711c7a9f'
-sha_dest_mac_only = 'b4fe3af17b7c364dd7cf5e2c8b02b07fde91659884177d1046574f696a94204e'
+sha_default = '68eeaed9682bdd63581cfe55c314ca3c024b487ebcbc2c539453b0ec85bcb21e'
+sha_dest_mac_only = '3b79e327a96dd6c8b49193cee794e813cfd6be50dd50c66e9082d984e8e91fc4'
 sha_multiple_victims_macos = '35c94ff8efd0958c39807e57477e967fdc100f9e12a68dcb6b368bb7de76056b'
-sha_one_victim_linux = 'da1ca7d73efc34abe4e6935c6e5b8e9fc7a3c9cc495545ed9c94522822238d5a'
+sha_one_victim_linux = 'aeb46ccd80e3862856931dc1d430e5a5e126deef5b291f9581735952f61facad'
 sha_port_shuffle = '9d124e51966324528e335e43a2a7bb635b20427a76337a5e1c91d56cc152acab'
 sha_smb2 = 'a3b3eebb41da0c8a8b6fd55855bdec44e30d2152947652113d496ec8f2aa60d9'
 sha_ip_src_shuffle = '79d86343ec9334cd750a8383f233e619142863c7010e48c28c4bc10a574c67e7'