ospd.scan.ScanCollection.get_ports() - Code Metrics - Inspection of "Improve communication with clients" - greenbone/ospd - Measure and Improve Code Quality continuously with Scrutinizer

Passed

Pull Request — master (#287)

by Juan José

created 2020-06-22 15:10 UTC

ospd.scan.ScanCollection.get_ports() A

↳ Parent: ospd.scan

Complexity

Conditions

Size

Total Lines	7
Code Lines	5

Duplication

Lines	0
Ratio	0 %

Importance

Changes

Metric	Value
cc	1
eloc	5
nop	2
dl	0
loc	7
rs	10
c	0
b	0
f	0

# Copyright (C) 2014-2020 Greenbone Networks GmbH
#
# SPDX-License-Identifier: AGPL-3.0-or-later
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.

import logging
import multiprocessing
import time
import uuid

from collections import OrderedDict
from enum import Enum
from typing import List, Any, Dict, Iterator, Optional, Iterable, Union

from ospd.network import target_str_to_list
from ospd.datapickler import DataPickler
from ospd.errors import OspdCommandError

LOGGER = logging.getLogger(__name__)


class ScanStatus(Enum):
    """Scan status. """

    QUEUED = 0
    INIT = 1
    RUNNING = 2
    STOPPED = 3
    FINISHED = 4


class ScanCollection:

    """ Scans collection, managing scans and results read and write, exposing
    only needed information.

    Each scan has meta-information such as scan ID, current progress (from 0 to
    100), start time, end time, scan target and options and a list of results.

    There are 4 types of results: Alarms, Logs, Errors and Host Details.

    Todo:
    - Better checking for Scan ID existence and handling otherwise.
    - More data validation.
    - Mutex access per table/scan_info.

    """

    def __init__(self, file_storage_dir: str) -> None:
        """ Initialize the Scan Collection. """

        self.data_manager = (
            None
        )  # type: Optional[multiprocessing.managers.SyncManager]
        self.scans_table = dict()  # type: Dict
        self.file_storage_dir = file_storage_dir

    def init(self):
        self.data_manager = multiprocessing.Manager()

    def add_result(
        self,
        scan_id: str,
        result_type: int,
        host: str = '',
        hostname: str = '',
        name: str = '',
        value: str = '',
        port: str = '',
        test_id: str = '',
        severity: str = '',
        qod: str = '',
        uri: str = '',
    ) -> None:
        """ Add a result to a scan in the table. """

        assert scan_id
        assert len(name) or len(value)

        result = OrderedDict()  # type: Dict
        result['type'] = result_type
        result['name'] = name
        result['severity'] = severity
        result['test_id'] = test_id
        result['value'] = value
        result['host'] = host
        result['hostname'] = hostname
        result['port'] = port
        result['qod'] = qod
        result['uri'] = uri
        results = self.scans_table[scan_id]['results']
        results.append(result)

        # Set scan_info's results to propagate results to parent process.
        self.scans_table[scan_id]['results'] = results

    def add_result_list(
        self, scan_id: str, result_list: Iterable[Dict[str, str]]
    ) -> None:
        """
        Add a batch of results to the result's table for the corresponding
        scan_id
        """
        results = self.scans_table[scan_id]['results']
        results.extend(result_list)

        # Set scan_info's results to propagate results to parent process.
        self.scans_table[scan_id]['results'] = results

    def remove_hosts_from_target_progress(
        self, scan_id: str, hosts: List
    ) -> None:
        """Remove a list of hosts from the main scan progress table to avoid
        the hosts to be included in the calculation of the scan progress"""
        if not hosts:
            return

        target = self.scans_table[scan_id].get('target_progress')
        for host in hosts:
            if host in target:
                del target[host]

        # Set scan_info's target_progress to propagate progresses
        # to parent process.
        self.scans_table[scan_id]['target_progress'] = target

    def set_progress(self, scan_id: str, progress: int) -> None:
        """ Sets scan_id scan's progress. """

        if progress > 0 and progress <= 100:
            self.scans_table[scan_id]['progress'] = progress

        if progress == 100:
            self.scans_table[scan_id]['end_time'] = int(time.time())

    def set_host_progress(
        self, scan_id: str, host_progress_batch: Dict[str, int]
    ) -> None:
        """ Sets scan_id scan's progress. """

        host_progresses = self.scans_table[scan_id].get('target_progress')
        host_progresses.update(host_progress_batch)

        # Set scan_info's target_progress to propagate progresses
        # to parent process.
        self.scans_table[scan_id]['target_progress'] = host_progresses

    def set_host_finished(self, scan_id: str, hosts: List[str]) -> None:
        """ Increase the amount of finished hosts which were alive."""

        total_finished = len(hosts)
        count_alive = (
            self.scans_table[scan_id].get('count_alive') + total_finished
        )
        self.scans_table[scan_id]['count_alive'] = count_alive

    def set_host_dead(self, scan_id: str, hosts: List[str]) -> None:
        """ Increase the amount of dead hosts. """

        total_dead = len(hosts)
        count_dead = self.scans_table[scan_id].get('count_dead') + total_dead
        self.scans_table[scan_id]['count_dead'] = count_dead

    def set_amount_dead_hosts(self, scan_id: str, total_dead: int) -> None:
        """ Increase the amount of dead hosts. """

        count_dead = self.scans_table[scan_id].get('count_dead') + total_dead
        self.scans_table[scan_id]['count_dead'] = count_dead

    def clean_temp_result_list(self, scan_id):
        """ Clean the results stored in the temporary list. """
        self.scans_table[scan_id]['temp_results'] = list()

    def restore_temp_result_list(self, scan_id):
        """ Add the results stored in the temporary list into the results
        list again. """
        result_aux = self.scans_table[scan_id].get('results', list())
        result_aux.extend(self.scans_table[scan_id].get('temp_results', list()))

        # Propagate results
        self.scans_table[scan_id]['results'] = result_aux
        self.clean_temp_result_list(scan_id)

    def results_iterator(
        self, scan_id: str, pop_res: bool = False, max_res: int = None
    ) -> Iterator[Any]:
        """ Returns an iterator over scan_id scan's results. If pop_res is True,
        it removed the fetched results from the list.

        If max_res is None, return all the results.
        Otherwise, if max_res = N > 0 return N as maximum number of results.

        max_res works only together with pop_results.
        """
        if pop_res and max_res:
            result_aux = self.scans_table[scan_id].get('results', list())
            self.scans_table[scan_id]['results'] = result_aux[max_res:]
            self.scans_table[scan_id]['temp_results'] = result_aux[:max_res]
            return iter(self.scans_table[scan_id]['temp_results'])
        elif pop_res:
            self.scans_table[scan_id]['temp_results'] = self.scans_table[
                scan_id
            ].get('results', list())
            self.scans_table[scan_id]['results'] = list()
            return iter(self.scans_table[scan_id]['temp_results'])

        return iter(self.scans_table[scan_id]['results'])

    def ids_iterator(self) -> Iterator[str]:
        """ Returns an iterator over the collection's scan IDS. """

        return iter(self.scans_table.keys())

    def clean_up_pickled_scan_info(self) -> None:
        """ Remove files of pickled scan info """
        for scan_id in self.ids_iterator():
            if self.get_status(scan_id) == ScanStatus.QUEUED:
                self.remove_file_pickled_scan_info(scan_id)

    def remove_file_pickled_scan_info(self, scan_id: str) -> None:
        pickler = DataPickler(self.file_storage_dir)
        pickler.remove_file(scan_id)

    def unpickle_scan_info(self, scan_id: str) -> None:
        """ Unpickle a stored scan_inf correspinding to the scan_id
        and store it in the scan_table """

        scan_info = self.scans_table.get(scan_id)
        scan_info_hash = scan_info.pop('scan_info_hash')

        pickler = DataPickler(self.file_storage_dir)
        unpickled_scan_info = pickler.load_data(scan_id, scan_info_hash)

        if not unpickled_scan_info:
            pickler.remove_file(scan_id)
            raise OspdCommandError(
                'Not possible to unpickle stored scan info for %s' % scan_id,
                'start_scan',
            )

        scan_info['results'] = list()
        scan_info['temp_results'] = list()
        scan_info['progress'] = 0
        scan_info['target_progress'] = dict()
        scan_info['count_alive'] = 0
        scan_info['count_dead'] = 0
        scan_info['target'] = unpickled_scan_info.pop('target')
        scan_info['vts'] = unpickled_scan_info.pop('vts')
        scan_info['options'] = unpickled_scan_info.pop('options')
        scan_info['start_time'] = int(time.time())
        scan_info['end_time'] = 0

        self.scans_table[scan_id] = scan_info

        pickler.remove_file(scan_id)

    def create_scan(
        self,
        scan_id: str = '',
        target: Dict = None,
        options: Optional[Dict] = None,
        vts: Dict = None,
    ) -> str:
        """ Creates a new scan with provided scan information. """

        if not options:
            options = dict()

        credentials = target.pop('credentials')

        scan_info = self.data_manager.dict()  # type: Dict
        scan_info['status'] = ScanStatus.QUEUED
        scan_info['credentials'] = credentials
        scan_info['start_time'] = int(time.time())

        scan_info_to_pickle = {
            'target': target,
            'options': options,
            'vts': vts,
        }

        if scan_id is None or scan_id == '':
            scan_id = str(uuid.uuid4())

        pickler = DataPickler(self.file_storage_dir)
        scan_info_hash = None
        try:
            scan_info_hash = pickler.store_data(scan_id, scan_info_to_pickle)
        except OspdCommandError as e:
            LOGGER.error(e)
            return

        scan_info['scan_id'] = scan_id
        scan_info['scan_info_hash'] = scan_info_hash

        self.scans_table[scan_id] = scan_info
        return scan_id

    def set_status(self, scan_id: str, status: ScanStatus) -> None:
        """ Sets scan_id scan's status. """
        self.scans_table[scan_id]['status'] = status
        if status == ScanStatus.STOPPED:
            self.scans_table[scan_id]['end_time'] = int(time.time())

    def get_status(self, scan_id: str) -> ScanStatus:
        """ Get scan_id scans's status."""

        return self.scans_table[scan_id].get('status')

    def get_options(self, scan_id: str) -> Dict:
        """ Get scan_id scan's options list. """

        return self.scans_table[scan_id].get('options')

    def set_option(self, scan_id, name: str, value: Any) -> None:
        """ Set a scan_id scan's name option to value. """

        self.scans_table[scan_id]['options'][name] = value

    def get_progress(self, scan_id: str) -> int:
        """ Get a scan's current progress value. """

        return self.scans_table[scan_id].get('progress', 0)

    def get_count_dead(self, scan_id: str) -> int:
        """ Get a scan's current dead host count. """

        return self.scans_table[scan_id]['count_dead']

    def get_count_alive(self, scan_id: str) -> int:
        """ Get a scan's current dead host count. """

        return self.scans_table[scan_id]['count_alive']

    def get_current_target_progress(self, scan_id: str) -> Dict[str, int]:
        """ Get a scan's current hosts progress """
        return self.scans_table[scan_id]['target_progress']

    def simplify_exclude_host_count(self, scan_id: str) -> int:
        """ Remove from exclude_hosts the received hosts in the finished_hosts
        list sent by the client.
        The finished hosts are sent also as exclude hosts for backward
        compatibility purposses.

        Return:
            Count of excluded host.
        """

        exc_hosts_list = target_str_to_list(self.get_exclude_hosts(scan_id))

        finished_hosts_list = target_str_to_list(
            self.get_finished_hosts(scan_id)
        )

        if finished_hosts_list and exc_hosts_list:
            for finished in finished_hosts_list:
                if finished in exc_hosts_list:
                    exc_hosts_list.remove(finished)

        return len(exc_hosts_list) if exc_hosts_list else 0

    def calculate_target_progress(self, scan_id: str) -> int:
        """ Get a target's current progress value.
        The value is calculated with the progress of each single host
        in the target."""

        total_hosts = self.get_host_count(scan_id)
        exc_hosts = self.simplify_exclude_host_count(scan_id)
        count_alive = self.get_count_alive(scan_id)
        count_dead = self.get_count_dead(scan_id)
        host_progresses = self.get_current_target_progress(scan_id)

        try:
            t_prog = int(
                (sum(host_progresses.values()) + 100 * count_alive)
                / (total_hosts - exc_hosts - count_dead)
            )
        except ZeroDivisionError:
            LOGGER.error(
                "Zero division error in %s",
                self.calculate_target_progress.__name__,
            )
            raise

        return t_prog

    def get_start_time(self, scan_id: str) -> str:
        """ Get a scan's start time. """

        return self.scans_table[scan_id]['start_time']

    def get_end_time(self, scan_id: str) -> str:
        """ Get a scan's end time. """

        return self.scans_table[scan_id]['end_time']

    def get_host_list(self, scan_id: str) -> Dict:
        """ Get a scan's host list. """

        return self.scans_table[scan_id]['target'].get('hosts')

    def get_host_count(self, scan_id: str) -> int:
        """ Get total host count in the target. """
        host = self.get_host_list(scan_id)
        total_hosts = len(target_str_to_list(host))

        return total_hosts

    def get_ports(self, scan_id: str) -> str:
        """ Get a scan's ports list.
        """
        target = self.scans_table[scan_id].get('target')
        ports = target.pop('ports')
        self.scans_table[scan_id]['target'] = target
        return ports

    def get_exclude_hosts(self, scan_id: str) -> str:
        """ Get an exclude host list for a given target.
        """
        return self.scans_table[scan_id]['target'].get('exclude_hosts')

    def get_finished_hosts(self, scan_id: str) -> str:
        """ Get the finished host list sent by the client for a given target.
        """
        return self.scans_table[scan_id]['target'].get('finished_hosts')

    def get_credentials(self, scan_id: str) -> Dict[str, Dict[str, str]]:
        """ Get a scan's credential list. It return dictionary with
        the corresponding credential for a given target.
        """
        return self.scans_table[scan_id].get('credentials')

    def get_target_options(self, scan_id: str) -> Dict[str, str]:
        """ Get a scan's target option dictionary.
        It return dictionary with the corresponding options for
        a given target.
        """
        return self.scans_table[scan_id]['target'].get('options')

    def get_vts(self, scan_id: str) -> Dict[str, Union[Dict[str, str], List]]:
        """ Get a scan's vts. """
        scan_info = self.scans_table[scan_id]
        vts = scan_info.pop('vts')
        self.scans_table[scan_id] = scan_info

        return vts

    def id_exists(self, scan_id: str) -> bool:
        """ Check whether a scan exists in the table. """

        return self.scans_table.get(scan_id) is not None

    def delete_scan(self, scan_id: str) -> bool:
        """ Delete a scan if fully finished. """

        if self.get_status(scan_id) == ScanStatus.RUNNING:
            return False

        scans_table = self.scans_table
        del scans_table[scan_id]
        self.scans_table = scans_table

        return True


1			# Copyright (C) 2014-2020 Greenbone Networks GmbH
2			#
3			# SPDX-License-Identifier: AGPL-3.0-or-later
4			#
5			# This program is free software: you can redistribute it and/or modify
6			# it under the terms of the GNU Affero General Public License as
7			# published by the Free Software Foundation, either version 3 of the
8			# License, or (at your option) any later version.
9			#
10			# This program is distributed in the hope that it will be useful,
11			# but WITHOUT ANY WARRANTY; without even the implied warranty of
12			# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
13			# GNU Affero General Public License for more details.
14			#
15			# You should have received a copy of the GNU Affero General Public License
16			# along with this program. If not, see <http://www.gnu.org/licenses/>.
17
18			import logging
19			import multiprocessing
20			import time
21			import uuid
22
23			from collections import OrderedDict
24			from enum import Enum
25			from typing import List, Any, Dict, Iterator, Optional, Iterable, Union
26
27			from ospd.network import target_str_to_list
28			from ospd.datapickler import DataPickler
29			from ospd.errors import OspdCommandError
30
31			LOGGER = logging.getLogger(__name__)
32
33
34			class ScanStatus(Enum):
35			"""Scan status. """
36
37			QUEUED = 0
38			INIT = 1
39			RUNNING = 2
40			STOPPED = 3
41			FINISHED = 4
42
43
44			class ScanCollection:
45
46			""" Scans collection, managing scans and results read and write, exposing
47			only needed information.
48
49			Each scan has meta-information such as scan ID, current progress (from 0 to
50			100), start time, end time, scan target and options and a list of results.
51
52			There are 4 types of results: Alarms, Logs, Errors and Host Details.
53
54			Todo:
55			- Better checking for Scan ID existence and handling otherwise.
56			- More data validation.
57			- Mutex access per table/scan_info.
58
59			"""
60
61			def __init__(self, file_storage_dir: str) -> None:
62			""" Initialize the Scan Collection. """
63
64			self.data_manager = (
65			None
66			) # type: Optional[multiprocessing.managers.SyncManager]
67			self.scans_table = dict() # type: Dict
68			self.file_storage_dir = file_storage_dir
69
70			def init(self):
71			self.data_manager = multiprocessing.Manager()
72
73			def add_result(
74			self,
75			scan_id: str,
76			result_type: int,
77			host: str = '',
78			hostname: str = '',
79			name: str = '',
80			value: str = '',
81			port: str = '',
82			test_id: str = '',
83			severity: str = '',
84			qod: str = '',
85			uri: str = '',
86			) -> None:
87			""" Add a result to a scan in the table. """
88
89			assert scan_id
90			assert len(name) or len(value)
91
92			result = OrderedDict() # type: Dict
93			result['type'] = result_type
94			result['name'] = name
95			result['severity'] = severity
96			result['test_id'] = test_id
97			result['value'] = value
98			result['host'] = host
99			result['hostname'] = hostname
100			result['port'] = port
101			result['qod'] = qod
102			result['uri'] = uri
103			results = self.scans_table[scan_id]['results']
104			results.append(result)
105
106			# Set scan_info's results to propagate results to parent process.
107			self.scans_table[scan_id]['results'] = results
108
109			def add_result_list(
110			self, scan_id: str, result_list: Iterable[Dict[str, str]]
111			) -> None:
112			"""
113			Add a batch of results to the result's table for the corresponding
114			scan_id
115			"""
116			results = self.scans_table[scan_id]['results']
117			results.extend(result_list)
118
119			# Set scan_info's results to propagate results to parent process.
120			self.scans_table[scan_id]['results'] = results
121
122			def remove_hosts_from_target_progress(
123			self, scan_id: str, hosts: List
124			) -> None:
125			"""Remove a list of hosts from the main scan progress table to avoid
126			the hosts to be included in the calculation of the scan progress"""
127			if not hosts:
128			return
129
130			target = self.scans_table[scan_id].get('target_progress')
131			for host in hosts:
132			if host in target:
133			del target[host]
134
135			# Set scan_info's target_progress to propagate progresses
136			# to parent process.
137			self.scans_table[scan_id]['target_progress'] = target
138
139			def set_progress(self, scan_id: str, progress: int) -> None:
140			""" Sets scan_id scan's progress. """
141
142			if progress > 0 and progress <= 100:
143			self.scans_table[scan_id]['progress'] = progress
144
145			if progress == 100:
146			self.scans_table[scan_id]['end_time'] = int(time.time())
147
148			def set_host_progress(
149			self, scan_id: str, host_progress_batch: Dict[str, int]
150			) -> None:
151			""" Sets scan_id scan's progress. """
152
153			host_progresses = self.scans_table[scan_id].get('target_progress')
154			host_progresses.update(host_progress_batch)
155
156			# Set scan_info's target_progress to propagate progresses
157			# to parent process.
158			self.scans_table[scan_id]['target_progress'] = host_progresses
159
160			def set_host_finished(self, scan_id: str, hosts: List[str]) -> None:
161			""" Increase the amount of finished hosts which were alive."""
162
163			total_finished = len(hosts)
164			count_alive = (
165			self.scans_table[scan_id].get('count_alive') + total_finished
166			)
167			self.scans_table[scan_id]['count_alive'] = count_alive
168
169			def set_host_dead(self, scan_id: str, hosts: List[str]) -> None:
170			""" Increase the amount of dead hosts. """
171
172			total_dead = len(hosts)
173			count_dead = self.scans_table[scan_id].get('count_dead') + total_dead
174			self.scans_table[scan_id]['count_dead'] = count_dead
175
176			def set_amount_dead_hosts(self, scan_id: str, total_dead: int) -> None:
177			""" Increase the amount of dead hosts. """
178
179			count_dead = self.scans_table[scan_id].get('count_dead') + total_dead
180			self.scans_table[scan_id]['count_dead'] = count_dead
181
182			def clean_temp_result_list(self, scan_id):
183			""" Clean the results stored in the temporary list. """
184			self.scans_table[scan_id]['temp_results'] = list()
185
186			def restore_temp_result_list(self, scan_id):
187			""" Add the results stored in the temporary list into the results
188			list again. """
189			result_aux = self.scans_table[scan_id].get('results', list())
190			result_aux.extend(self.scans_table[scan_id].get('temp_results', list()))
191
192			# Propagate results
193			self.scans_table[scan_id]['results'] = result_aux
194			self.clean_temp_result_list(scan_id)
195
196			def results_iterator(
197			self, scan_id: str, pop_res: bool = False, max_res: int = None
198			) -> Iterator[Any]:
199			""" Returns an iterator over scan_id scan's results. If pop_res is True,
200			it removed the fetched results from the list.
201
202			If max_res is None, return all the results.
203			Otherwise, if max_res = N > 0 return N as maximum number of results.
204
205			max_res works only together with pop_results.
206			"""
207			if pop_res and max_res:
208			result_aux = self.scans_table[scan_id].get('results', list())
209			self.scans_table[scan_id]['results'] = result_aux[max_res:]
210			self.scans_table[scan_id]['temp_results'] = result_aux[:max_res]
211			return iter(self.scans_table[scan_id]['temp_results'])
212			elif pop_res:
213			self.scans_table[scan_id]['temp_results'] = self.scans_table[
214			scan_id
215			].get('results', list())
216			self.scans_table[scan_id]['results'] = list()
217			return iter(self.scans_table[scan_id]['temp_results'])
218
219			return iter(self.scans_table[scan_id]['results'])
220
221			def ids_iterator(self) -> Iterator[str]:
222			""" Returns an iterator over the collection's scan IDS. """
223
224			return iter(self.scans_table.keys())
225
226			def clean_up_pickled_scan_info(self) -> None:
227			""" Remove files of pickled scan info """
228			for scan_id in self.ids_iterator():
229			if self.get_status(scan_id) == ScanStatus.QUEUED:
230			self.remove_file_pickled_scan_info(scan_id)
231
232			def remove_file_pickled_scan_info(self, scan_id: str) -> None:
233			pickler = DataPickler(self.file_storage_dir)
234			pickler.remove_file(scan_id)
235
236			def unpickle_scan_info(self, scan_id: str) -> None:
237			""" Unpickle a stored scan_inf correspinding to the scan_id
238			and store it in the scan_table """
239
240			scan_info = self.scans_table.get(scan_id)
241			scan_info_hash = scan_info.pop('scan_info_hash')
242
243			pickler = DataPickler(self.file_storage_dir)
244			unpickled_scan_info = pickler.load_data(scan_id, scan_info_hash)
245
246			if not unpickled_scan_info:
247			pickler.remove_file(scan_id)
248			raise OspdCommandError(
249			'Not possible to unpickle stored scan info for %s' % scan_id,
250			'start_scan',
251			)
252
253			scan_info['results'] = list()
254			scan_info['temp_results'] = list()
255			scan_info['progress'] = 0
256			scan_info['target_progress'] = dict()
257			scan_info['count_alive'] = 0
258			scan_info['count_dead'] = 0
259			scan_info['target'] = unpickled_scan_info.pop('target')
260			scan_info['vts'] = unpickled_scan_info.pop('vts')
261			scan_info['options'] = unpickled_scan_info.pop('options')
262			scan_info['start_time'] = int(time.time())
263			scan_info['end_time'] = 0
264
265			self.scans_table[scan_id] = scan_info
266
267			pickler.remove_file(scan_id)
268
269			def create_scan(
270			self,
271			scan_id: str = '',
272			target: Dict = None,
273			options: Optional[Dict] = None,
274			vts: Dict = None,
275			) -> str:
276			""" Creates a new scan with provided scan information. """
277
278			if not options:
279			options = dict()
280
281			credentials = target.pop('credentials')
282
283			scan_info = self.data_manager.dict() # type: Dict
284			scan_info['status'] = ScanStatus.QUEUED
285			scan_info['credentials'] = credentials
286			scan_info['start_time'] = int(time.time())
287
288			scan_info_to_pickle = {
289			'target': target,
290			'options': options,
291			'vts': vts,
292			}
293
294			if scan_id is None or scan_id == '':
295			scan_id = str(uuid.uuid4())
296
297			pickler = DataPickler(self.file_storage_dir)
298			scan_info_hash = None
299			try:
300			scan_info_hash = pickler.store_data(scan_id, scan_info_to_pickle)
301			except OspdCommandError as e:
302			LOGGER.error(e)
303			return
304
305			scan_info['scan_id'] = scan_id
306			scan_info['scan_info_hash'] = scan_info_hash
307
308			self.scans_table[scan_id] = scan_info
309			return scan_id
310
311			def set_status(self, scan_id: str, status: ScanStatus) -> None:
312			""" Sets scan_id scan's status. """
313			self.scans_table[scan_id]['status'] = status
314			if status == ScanStatus.STOPPED:
315			self.scans_table[scan_id]['end_time'] = int(time.time())
316
317			def get_status(self, scan_id: str) -> ScanStatus:
318			""" Get scan_id scans's status."""
319
320			return self.scans_table[scan_id].get('status')
321
322			def get_options(self, scan_id: str) -> Dict:
323			""" Get scan_id scan's options list. """
324
325			return self.scans_table[scan_id].get('options')
326
327			def set_option(self, scan_id, name: str, value: Any) -> None:
328			""" Set a scan_id scan's name option to value. """
329
330			self.scans_table[scan_id]['options'][name] = value
331
332			def get_progress(self, scan_id: str) -> int:
333			""" Get a scan's current progress value. """
334
335			return self.scans_table[scan_id].get('progress', 0)
336
337			def get_count_dead(self, scan_id: str) -> int:
338			""" Get a scan's current dead host count. """
339
340			return self.scans_table[scan_id]['count_dead']
341
342			def get_count_alive(self, scan_id: str) -> int:
343			""" Get a scan's current dead host count. """
344
345			return self.scans_table[scan_id]['count_alive']
346
347			def get_current_target_progress(self, scan_id: str) -> Dict[str, int]:
348			""" Get a scan's current hosts progress """
349			return self.scans_table[scan_id]['target_progress']
350
351			def simplify_exclude_host_count(self, scan_id: str) -> int:
352			""" Remove from exclude_hosts the received hosts in the finished_hosts
353			list sent by the client.
354			The finished hosts are sent also as exclude hosts for backward
355			compatibility purposses.
356
357			Return:
358			Count of excluded host.
359			"""
360
361			exc_hosts_list = target_str_to_list(self.get_exclude_hosts(scan_id))
362
363			finished_hosts_list = target_str_to_list(
364			self.get_finished_hosts(scan_id)
365			)
366
367			if finished_hosts_list and exc_hosts_list:
368			for finished in finished_hosts_list:
369			if finished in exc_hosts_list:
370			exc_hosts_list.remove(finished)
371
372			return len(exc_hosts_list) if exc_hosts_list else 0
373
374			def calculate_target_progress(self, scan_id: str) -> int:
375			""" Get a target's current progress value.
376			The value is calculated with the progress of each single host
377			in the target."""
378
379			total_hosts = self.get_host_count(scan_id)
380			exc_hosts = self.simplify_exclude_host_count(scan_id)
381			count_alive = self.get_count_alive(scan_id)
382			count_dead = self.get_count_dead(scan_id)
383			host_progresses = self.get_current_target_progress(scan_id)
384
385			try:
386			t_prog = int(
387			(sum(host_progresses.values()) + 100 * count_alive)
388			/ (total_hosts - exc_hosts - count_dead)
389			)
390			except ZeroDivisionError:
391			LOGGER.error(
392			"Zero division error in %s",
393			self.calculate_target_progress.__name__,
394			)
395			raise
396
397			return t_prog
398
399			def get_start_time(self, scan_id: str) -> str:
400			""" Get a scan's start time. """
401
402			return self.scans_table[scan_id]['start_time']
403
404			def get_end_time(self, scan_id: str) -> str:
405			""" Get a scan's end time. """
406
407			return self.scans_table[scan_id]['end_time']
408
409			def get_host_list(self, scan_id: str) -> Dict:
410			""" Get a scan's host list. """
411
412			return self.scans_table[scan_id]['target'].get('hosts')
413
414			def get_host_count(self, scan_id: str) -> int:
415			""" Get total host count in the target. """
416			host = self.get_host_list(scan_id)
417			total_hosts = len(target_str_to_list(host))
418
419			return total_hosts
420
421			def get_ports(self, scan_id: str) -> str:
422			""" Get a scan's ports list.
423			"""
424			target = self.scans_table[scan_id].get('target')
425			ports = target.pop('ports')
426			self.scans_table[scan_id]['target'] = target
427			return ports
428
429			def get_exclude_hosts(self, scan_id: str) -> str:
430			""" Get an exclude host list for a given target.
431			"""
432			return self.scans_table[scan_id]['target'].get('exclude_hosts')
433
434			def get_finished_hosts(self, scan_id: str) -> str:
435			""" Get the finished host list sent by the client for a given target.
436			"""
437			return self.scans_table[scan_id]['target'].get('finished_hosts')
438
439			def get_credentials(self, scan_id: str) -> Dict[str, Dict[str, str]]:
440			""" Get a scan's credential list. It return dictionary with
441			the corresponding credential for a given target.
442			"""
443			return self.scans_table[scan_id].get('credentials')
444
445			def get_target_options(self, scan_id: str) -> Dict[str, str]:
446			""" Get a scan's target option dictionary.
447			It return dictionary with the corresponding options for
448			a given target.
449			"""
450			return self.scans_table[scan_id]['target'].get('options')
451
452			def get_vts(self, scan_id: str) -> Dict[str, Union[Dict[str, str], List]]:
453			""" Get a scan's vts. """
454			scan_info = self.scans_table[scan_id]
455			vts = scan_info.pop('vts')
456			self.scans_table[scan_id] = scan_info
457
458			return vts
459
460			def id_exists(self, scan_id: str) -> bool:
461			""" Check whether a scan exists in the table. """
462
463			return self.scans_table.get(scan_id) is not None
464
465			def delete_scan(self, scan_id: str) -> bool:
466			""" Delete a scan if fully finished. """
467
468			if self.get_status(scan_id) == ScanStatus.RUNNING:
469			return False
470
471			scans_table = self.scans_table
472			del scans_table[scan_id]
473			self.scans_table = scans_table
474
475			return True
476

greenbone / ospd

Pull Request — master (#287)

ospd.scan.ScanCollection.get_ports() A

Complexity

Size

Duplication

Importance

Duplication Side-by-Side

Filter issues like