Gh0st
/
Nettacker
mirror of https://github.com/OWASP/Nettacker.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Refactor code for future `numpy` removal

This commit is contained in:
Arkadii Yakovets 2024-10-21 14:46:26 -07:00
parent 637aa76508
commit 8921009c14
No known key found for this signature in database
GPG Key ID: 7350E7F17DFE6846
7 changed files with 136 additions and 76 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
nettacker/core/app.py
View File

@ -6,7 +6,6 @@ import sys
from threading import Thread
import multiprocess
import numpy
from nettacker import logger
from nettacker.config import Config, version_info
@ -26,7 +25,7 @@ from nettacker.core.ip import (
from nettacker.core.messages import messages as _
from nettacker.core.module import Module
from nettacker.core.socks_proxy import set_socks_proxy
from nettacker.core.utils import common as utils
from nettacker.core.utils import common as common_utils
from nettacker.core.utils.common import wait_for_threads_to_finish
from nettacker.database.db import find_events, remove_old_logs
from nettacker.database.mysql import mysql_create_database, mysql_create_tables
@ -192,7 +191,7 @@ class Nettacker(ArgParser):
Returns:
True when it ends
"""
scan_id = utils.generate_random_token(32)
scan_id = common_utils.generate_random_token(32)
log.info("ScanID: {0}".format(scan_id))
log.info(_("regrouping_targets"))
# find total number of targets + types + expand (subdomain, IPRanges, etc)
@ -210,18 +209,9 @@ class Nettacker(ArgParser):
return exit_code
def start_scan(self, scan_id):
number_of_total_targets = len(self.arguments.targets)
target_groups = [
targets.tolist()
for targets in numpy.array_split(
self.arguments.targets,
(
self.arguments.set_hardware_usage
if self.arguments.set_hardware_usage <= len(self.arguments.targets)
else number_of_total_targets
),
target_groups = common_utils.generate_target_groups(
self.arguments.targets, self.arguments.set_hardware_usage
)
]
log.info(_("removing_old_db_records"))
for target_group in target_groups:
@ -239,7 +229,7 @@ class Nettacker(ArgParser):
for _i in range(target_groups.count([])):
target_groups.remove([])
log.info(_("start_multi_process").format(number_of_total_targets, len(target_groups)))
log.info(_("start_multi_process").format(len(self.arguments.targets), len(target_groups)))
active_processes = []
for t_id, target_groups in enumerate(target_groups):
process = multiprocess.Process(

12
nettacker/core/arg_parser.py
View File

@ -17,7 +17,7 @@ from nettacker.core.ip import (
)
from nettacker.core.messages import messages as _
from nettacker.core.template import TemplateLoader
from nettacker.core.utils import common as utils
from nettacker.core.utils import common as common_utils
from nettacker.logger import TerminalCodes, get_logger
log = get_logger()
@ -92,7 +92,7 @@ class ArgParser(ArgumentParser):
if len(module_names) == limit:
module_names["..."] = {}
break
module_names = utils.sort_dictionary(module_names)
module_names = common_utils.sort_dictionary(module_names)
module_names["all"] = {}
return module_names
@ -118,11 +118,11 @@ class ArgParser(ArgumentParser):
else:
profiles[tag].append(key)
if len(profiles) == limit:
profiles = utils.sort_dictionary(profiles)
profiles = common_utils.sort_dictionary(profiles)
profiles["..."] = []
profiles["all"] = []
return profiles
profiles = utils.sort_dictionary(profiles)
profiles = common_utils.sort_dictionary(profiles)
profiles["all"] = []
return profiles
@ -612,7 +612,9 @@ class ArgParser(ArgumentParser):
# threading & processing
if options.set_hardware_usage not in {"low", "normal", "high", "maximum"}:
die_failure(_("wrong_hardware_usage"))
options.set_hardware_usage = utils.select_maximum_cpu_core(options.set_hardware_usage)
options.set_hardware_usage = common_utils.select_maximum_cpu_core(
options.set_hardware_usage
)
options.thread_per_host = int(options.thread_per_host)
if options.thread_per_host < 1:

18
nettacker/core/utils/common.py
View File

@ -11,6 +11,8 @@ import string
import sys
import time
import numpy
from nettacker import logger
log = logger.get_logger()
@ -181,9 +183,21 @@ def generate_and_replace_md5(content):
)
def arrays_to_matrix(arrays):
import numpy
def generate_target_groups(targets, set_hardware_usage):
if not targets:
return targets
targets_total = len(targets)
return [
targets.tolist()
for targets in numpy.array_split(
targets,
(set_hardware_usage if set_hardware_usage <= targets_total else targets_total),
)
]
def arrays_to_matrix(arrays):
return (
numpy.array(numpy.meshgrid(*[arrays[array_name] for array_name in arrays]))
.T.reshape(-1, len(arrays.keys()))

0
tests/core/test_socket.py → tests/core/lib/test_socket.py
View File

0
tests/core/test_ssl.py → tests/core/lib/test_ssl.py
View File

53
tests/core/test_utils.py
View File

@ -1,53 +0,0 @@
from unittest.mock import patch
from nettacker.core.utils import common as utils
from tests.common import TestCase
class TestUtils(TestCase):
def test_sort_dictionary(self):
input_dict = {
"a": 1,
"c": 3,
"d": 23,
"b": 2,
}
expected_dict = {
"a": 1,
"b": 2,
"c": 3,
"d": 23,
}
input_dict_keys = tuple(input_dict.keys())
expected_dict_keys = tuple(expected_dict.keys())
self.assertNotEqual(input_dict_keys, expected_dict_keys)
sorted_dict_keys = tuple(utils.sort_dictionary(input_dict).keys())
self.assertEqual(sorted_dict_keys, expected_dict_keys)
@patch("multiprocessing.cpu_count")
def test_select_maximum_cpu_core(self, cpu_count_mock):
cores_mapping = {
1: {"low": 1, "normal": 1, "high": 1, "maximum": 1},
2: {"low": 1, "normal": 1, "high": 1, "maximum": 1},
4: {"low": 1, "normal": 1, "high": 2, "maximum": 3},
6: {"low": 1, "normal": 1, "high": 3, "maximum": 5},
8: {"low": 1, "normal": 2, "high": 4, "maximum": 7},
10: {"low": 1, "normal": 2, "high": 5, "maximum": 9},
12: {"low": 1, "normal": 3, "high": 6, "maximum": 11},
16: {"low": 2, "normal": 4, "high": 8, "maximum": 15},
32: {"low": 4, "normal": 8, "high": 16, "maximum": 31},
48: {"low": 6, "normal": 12, "high": 24, "maximum": 47},
64: {"low": 8, "normal": 16, "high": 32, "maximum": 63},
}
for num_cores, levels in cores_mapping.items():
cpu_count_mock.return_value = num_cores
for level in ("low", "normal", "high", "maximum"):
self.assertEqual(
utils.select_maximum_cpu_core(level),
levels[level],
f"It should be {utils.select_maximum_cpu_core(level)} "
"of {num_cores} cores for '{level}' mode",
)
self.assertEqual(utils.select_maximum_cpu_core("invalid"), 1)

107
tests/core/utils/test_common.py Normal file
View File

@ -0,0 +1,107 @@
from unittest.mock import patch
from nettacker.core.utils import common as common_utils
from tests.common import TestCase
class TestCommon(TestCase):
def test_arrays_to_matrix(self):
(
self.assertEqual(
sorted(
common_utils.arrays_to_matrix(
{"ports": [1, 2, 3, 4, 5]},
)
),
[[1], [2], [3], [4], [5]],
),
)
self.assertEqual(
sorted(
common_utils.arrays_to_matrix(
{"x": [1, 2], "y": [3, 4], "z": [5, 6]},
)
),
[
[1, 3, 5],
[1, 3, 6],
[1, 4, 5],
[1, 4, 6],
[2, 3, 5],
[2, 3, 6],
[2, 4, 5],
[2, 4, 6],
],
)
def test_generate_target_groups_empty_list(self):
targets = []
set_hardware_usage = 3
result = common_utils.generate_target_groups(targets, set_hardware_usage)
assert result == []
def test_generate_target_groups_set_hardware_less_than_targets_total(self):
targets = [1, 2, 3, 4, 5]
set_hardware_usage = 2
result = common_utils.generate_target_groups(targets, set_hardware_usage)
assert result == [[1, 2, 3], [4, 5]]
def test_generate_target_groups_set_hardware_equal_to_targets_total(self):
targets = [1, 2, 3, 4, 5]
set_hardware_usage = 5
result = common_utils.generate_target_groups(targets, set_hardware_usage)
assert result == [[1], [2], [3], [4], [5]]
def test_generate_target_groups_set_hardware_greater_than_targets_total(self):
targets = [1, 2, 3]
set_hardware_usage = 5
result = common_utils.generate_target_groups(targets, set_hardware_usage)
assert result == [[1], [2], [3]]
def test_sort_dictionary(self):
input_dict = {
"a": 1,
"c": 3,
"d": 23,
"b": 2,
}
expected_dict = {
"a": 1,
"b": 2,
"c": 3,
"d": 23,
}
input_dict_keys = tuple(input_dict.keys())
expected_dict_keys = tuple(expected_dict.keys())
self.assertNotEqual(input_dict_keys, expected_dict_keys)
sorted_dict_keys = tuple(common_utils.sort_dictionary(input_dict).keys())
self.assertEqual(sorted_dict_keys, expected_dict_keys)
@patch("multiprocessing.cpu_count")
def test_select_maximum_cpu_core(self, cpu_count_mock):
cores_mapping = {
1: {"low": 1, "normal": 1, "high": 1, "maximum": 1},
2: {"low": 1, "normal": 1, "high": 1, "maximum": 1},
4: {"low": 1, "normal": 1, "high": 2, "maximum": 3},
6: {"low": 1, "normal": 1, "high": 3, "maximum": 5},
8: {"low": 1, "normal": 2, "high": 4, "maximum": 7},
10: {"low": 1, "normal": 2, "high": 5, "maximum": 9},
12: {"low": 1, "normal": 3, "high": 6, "maximum": 11},
16: {"low": 2, "normal": 4, "high": 8, "maximum": 15},
32: {"low": 4, "normal": 8, "high": 16, "maximum": 31},
48: {"low": 6, "normal": 12, "high": 24, "maximum": 47},
64: {"low": 8, "normal": 16, "high": 32, "maximum": 63},
}
for num_cores, levels in cores_mapping.items():
cpu_count_mock.return_value = num_cores
for level in ("low", "normal", "high", "maximum"):
self.assertEqual(
common_utils.select_maximum_cpu_core(level),
levels[level],
f"It should be {common_utils.select_maximum_cpu_core(level)} "
"of {num_cores} cores for '{level}' mode",
)
self.assertEqual(common_utils.select_maximum_cpu_core("invalid"), 1)