Completed
Push — master ( 22ac2f...f571c7 )
by Juan José
17s queued 12s
created

ospd.scan.ScanCollection.get_exclude_hosts()   A

Complexity

Conditions 1

Size

Total Lines 4
Code Lines 2

Duplication

Lines 0
Ratio 0 %

Importance

Changes 0
Metric Value
cc 1
eloc 2
nop 2
dl 0
loc 4
rs 10
c 0
b 0
f 0
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
26
27
from ospd.network import target_str_to_list
28
29
LOGGER = logging.getLogger(__name__)
30
31
32
class ScanStatus(Enum):
33
    """Scan status. """
34
35
    INIT = 0
36
    RUNNING = 1
37
    STOPPED = 2
38
    FINISHED = 3
39
40
41
class ScanCollection:
42
43
    """ Scans collection, managing scans and results read and write, exposing
44
    only needed information.
45
46
    Each scan has meta-information such as scan ID, current progress (from 0 to
47
    100), start time, end time, scan target and options and a list of results.
48
49
    There are 4 types of results: Alarms, Logs, Errors and Host Details.
50
51
    Todo:
52
    - Better checking for Scan ID existence and handling otherwise.
53
    - More data validation.
54
    - Mutex access per table/scan_info.
55
56
    """
57
58
    def __init__(self) -> None:
59
        """ Initialize the Scan Collection. """
60
61
        self.data_manager = (
62
            None
63
        )  # type: Optional[multiprocessing.managers.SyncManager]
64
        self.scans_table = dict()  # type: Dict
65
66
    def add_result(
67
        self,
68
        scan_id: str,
69
        result_type: int,
70
        host: str = '',
71
        hostname: str = '',
72
        name: str = '',
73
        value: str = '',
74
        port: str = '',
75
        test_id: str = '',
76
        severity: str = '',
77
        qod: str = '',
78
    ) -> None:
79
        """ Add a result to a scan in the table. """
80
81
        assert scan_id
82
        assert len(name) or len(value)
83
84
        result = OrderedDict()  # type: Dict
85
        result['type'] = result_type
86
        result['name'] = name
87
        result['severity'] = severity
88
        result['test_id'] = test_id
89
        result['value'] = value
90
        result['host'] = host
91
        result['hostname'] = hostname
92
        result['port'] = port
93
        result['qod'] = qod
94
        results = self.scans_table[scan_id]['results']
95
        results.append(result)
96
97
        # Set scan_info's results to propagate results to parent process.
98
        self.scans_table[scan_id]['results'] = results
99
100
    def add_result_list(
101
        self, scan_id: str, result_list: Iterable[Dict[str, str]]
102
    ) -> None:
103
        """ Add a batch of results to the result's table for the corresponding scan_id """
104
        results = self.scans_table[scan_id]['results']
105
        results.extend(result_list)
106
107
        # Set scan_info's results to propagate results to parent process.
108
        self.scans_table[scan_id]['results'] = results
109
110
    def remove_hosts_from_target_progress(
111
        self, scan_id: str, hosts: List
112
    ) -> None:
113
        """Remove a list of hosts from the main scan progress table to avoid
114
        the hosts to be included in the calculation of the scan progress"""
115
        if not hosts:
116
            return
117
118
        target = self.scans_table[scan_id].get('target_progress')
119
        for host in hosts:
120
            if host in target:
121
                del target[host]
122
123
        # Set scan_info's target_progress to propagate progresses
124
        # to parent process.
125
        self.scans_table[scan_id]['target_progress'] = target
126
127
    def set_progress(self, scan_id: str, progress: int) -> None:
128
        """ Sets scan_id scan's progress. """
129
130
        if progress > 0 and progress <= 100:
131
            self.scans_table[scan_id]['progress'] = progress
132
133
        if progress == 100:
134
            self.scans_table[scan_id]['end_time'] = int(time.time())
135
136
    def set_host_progress(
137
        self, scan_id: str, host_progress_batch: Dict[str, int]
138
    ) -> None:
139
        """ Sets scan_id scan's progress. """
140
        host_progresses = self.scans_table[scan_id].get('target_progress')
141
        host_progresses.update(host_progress_batch)
142
143
        # Set scan_info's target_progress to propagate progresses
144
        # to parent process.
145
        self.scans_table[scan_id]['target_progress'] = host_progresses
146
147
    def set_host_finished(self, scan_id: str, hosts: List[str]) -> None:
148
        """ Add the host in a list of finished hosts """
149
        finished_hosts = self.scans_table[scan_id].get('finished_hosts')
150
        finished_hosts.extend(hosts)
151
        uniq_set_of_hosts = set(finished_hosts)
152
153
        self.scans_table[scan_id]['finished_hosts'] = list(uniq_set_of_hosts)
154
155
    def get_hosts_unfinished(self, scan_id: str) -> List[Any]:
156
        """ Get a list of unfinished hosts."""
157
158
        unfinished_hosts = target_str_to_list(self.get_host_list(scan_id))
159
160
        finished_hosts = self.get_hosts_finished(scan_id)
161
162
        for host in finished_hosts:
163
            unfinished_hosts.remove(host)
164
165
        return unfinished_hosts
166
167
    def get_hosts_finished(self, scan_id: str) -> List:
168
        """ Get a list of finished hosts."""
169
170
        return self.scans_table[scan_id].get('finished_hosts')
171
172
    def results_iterator(
173
        self, scan_id: str, pop_res: bool = False, max_res: int = None
174
    ) -> Iterator[Any]:
175
        """ Returns an iterator over scan_id scan's results. If pop_res is True,
176
        it removed the fetched results from the list.
177
178
        If max_res is None, return all the results.
179
        Otherwise, if max_res = N > 0 return N as maximum number of results.
180
181
        max_res works only together with pop_results.
182
        """
183
        if pop_res and max_res:
184
            result_aux = self.scans_table[scan_id]['results']
185
            self.scans_table[scan_id]['results'] = result_aux[max_res:]
186
            return iter(result_aux[:max_res])
187
        elif pop_res:
188
            result_aux = self.scans_table[scan_id]['results']
189
            self.scans_table[scan_id]['results'] = list()
190
            return iter(result_aux)
191
192
        return iter(self.scans_table[scan_id]['results'])
193
194
    def ids_iterator(self) -> Iterator[str]:
195
        """ Returns an iterator over the collection's scan IDS. """
196
197
        return iter(self.scans_table.keys())
198
199
    def remove_single_result(
200
        self, scan_id: str, result: Dict[str, str]
201
    ) -> None:
202
        """Removes a single result from the result list in scan_table.
203
204
        Parameters:
205
            scan_id (uuid): Scan ID to identify the scan process to be resumed.
206
            result (dict): The result to be removed from the results list.
207
        """
208
        results = self.scans_table[scan_id]['results']
209
        results.remove(result)
210
        self.scans_table[scan_id]['results'] = results
211
212
    def del_results_for_stopped_hosts(self, scan_id: str) -> None:
213
        """ Remove results from the result table for those host
214
        """
215
        unfinished_hosts = self.get_hosts_unfinished(scan_id)
216
        for result in self.results_iterator(
217
            scan_id, pop_res=False, max_res=None
218
        ):
219
            if result['host'] in unfinished_hosts:
220
                self.remove_single_result(scan_id, result)
221
222
    def resume_scan(self, scan_id: str, options: Optional[Dict]) -> str:
223
        """ Reset the scan status in the scan_table to INIT.
224
        Also, overwrite the options, because a resume task cmd
225
        can add some new option. E.g. exclude hosts list.
226
        Parameters:
227
            scan_id (uuid): Scan ID to identify the scan process to be resumed.
228
            options (dict): Options for the scan to be resumed. This options
229
                            are not added to the already existent ones.
230
                            The old ones are removed
231
232
        Return:
233
            Scan ID which identifies the current scan.
234
        """
235
        self.scans_table[scan_id]['status'] = ScanStatus.INIT
236
        if options:
237
            self.scans_table[scan_id]['options'] = options
238
239
        self.del_results_for_stopped_hosts(scan_id)
240
241
        return scan_id
242
243
    def create_scan(
244
        self,
245
        scan_id: str = '',
246
        target: Dict = None,
247
        options: Optional[Dict] = None,
248
        vts: Dict = None,
249
    ) -> str:
250
        """ Creates a new scan with provided scan information. """
251
252
        if not target:
253
            target = {}
254
255
        if self.data_manager is None:
256
            self.data_manager = multiprocessing.Manager()
257
258
        # Check if it is possible to resume task. To avoid to resume, the
259
        # scan must be deleted from the scans_table.
260
        if (
261
            scan_id
262
            and self.id_exists(scan_id)
263
            and (self.get_status(scan_id) == ScanStatus.STOPPED)
264
        ):
265
            self.scans_table[scan_id]['end_time'] = 0
266
267
            return self.resume_scan(scan_id, options)
268
269
        if not options:
270
            options = dict()
271
272
        scan_info = self.data_manager.dict()  # type: Dict
273
        scan_info['results'] = list()
274
        scan_info['finished_hosts'] = list()
275
        scan_info['progress'] = 0
276
        scan_info['target_progress'] = dict()
277
        scan_info['target'] = target
278
        scan_info['vts'] = vts
279
        scan_info['options'] = options
280
        scan_info['start_time'] = int(time.time())
281
        scan_info['end_time'] = 0
282
        scan_info['status'] = ScanStatus.INIT
283
284
        if scan_id is None or scan_id == '':
285
            scan_id = str(uuid.uuid4())
286
287
        scan_info['scan_id'] = scan_id
288
289
        self.scans_table[scan_id] = scan_info
290
        return scan_id
291
292
    def set_status(self, scan_id: str, status: ScanStatus) -> None:
293
        """ Sets scan_id scan's status. """
294
        self.scans_table[scan_id]['status'] = status
295
        if status == ScanStatus.STOPPED:
296
            self.scans_table[scan_id]['end_time'] = int(time.time())
297
298
    def get_status(self, scan_id: str) -> ScanStatus:
299
        """ Get scan_id scans's status."""
300
301
        return self.scans_table[scan_id]['status']
302
303
    def get_options(self, scan_id: str) -> Dict:
304
        """ Get scan_id scan's options list. """
305
306
        return self.scans_table[scan_id]['options']
307
308
    def set_option(self, scan_id, name: str, value: Any) -> None:
309
        """ Set a scan_id scan's name option to value. """
310
311
        self.scans_table[scan_id]['options'][name] = value
312
313
    def get_progress(self, scan_id: str) -> int:
314
        """ Get a scan's current progress value. """
315
316
        return self.scans_table[scan_id]['progress']
317
318
    def simplify_exclude_host_list(self, scan_id: str) -> List[Any]:
319
        """ Remove from exclude_hosts the received hosts in the finished_hosts
320
        list sent by the client.
321
        The finished hosts are sent also as exclude hosts for backward
322
        compatibility purposses.
323
        """
324
325
        exc_hosts_list = target_str_to_list(self.get_exclude_hosts(scan_id))
326
327
        finished_hosts_list = target_str_to_list(
328
            self.get_finished_hosts(scan_id)
329
        )
330
331
        if finished_hosts_list and exc_hosts_list:
332
            for finished in finished_hosts_list:
333
                if finished in exc_hosts_list:
334
                    exc_hosts_list.remove(finished)
335
336
        return exc_hosts_list
337
338
    def calculate_target_progress(self, scan_id: str) -> float:
339
        """ Get a target's current progress value.
340
        The value is calculated with the progress of each single host
341
        in the target."""
342
343
        host = self.get_host_list(scan_id)
344
        total_hosts = len(target_str_to_list(host))
345
        exc_hosts_list = self.simplify_exclude_host_list(scan_id)
346
        exc_hosts = len(exc_hosts_list) if exc_hosts_list else 0
347
        host_progresses = self.scans_table[scan_id].get('target_progress')
348
349
        try:
350
            t_prog = sum(host_progresses.values()) / (
351
                total_hosts - exc_hosts
352
            )  # type: float
353
        except ZeroDivisionError:
354
            LOGGER.error(
355
                "Zero division error in %s",
356
                self.calculate_target_progress.__name__,
357
            )
358
            raise
359
360
        return t_prog
361
362
    def get_start_time(self, scan_id: str) -> str:
363
        """ Get a scan's start time. """
364
365
        return self.scans_table[scan_id]['start_time']
366
367
    def get_end_time(self, scan_id: str) -> str:
368
        """ Get a scan's end time. """
369
370
        return self.scans_table[scan_id]['end_time']
371
372
    def get_host_list(self, scan_id: str) -> Dict:
373
        """ Get a scan's host list. """
374
375
        return self.scans_table[scan_id]['target'].get('hosts')
376
377
    def get_ports(self, scan_id: str):
378
        """ Get a scan's ports list.
379
        """
380
        return self.scans_table[scan_id]['target'].get('ports')
381
382
    def get_exclude_hosts(self, scan_id: str):
383
        """ Get an exclude host list for a given target.
384
        """
385
        return self.scans_table[scan_id]['target'].get('exclude_hosts')
386
387
    def get_finished_hosts(self, scan_id: str):
388
        """ Get the finished host list sent by the client for a given target.
389
        """
390
        return self.scans_table[scan_id]['target'].get('finished_hosts')
391
392
    def get_credentials(self, scan_id: str):
393
        """ Get a scan's credential list. It return dictionary with
394
        the corresponding credential for a given target.
395
        """
396
        return self.scans_table[scan_id]['target'].get('credentials')
397
398
    def get_target_options(self, scan_id: str):
399
        """ Get a scan's target option dictionary.
400
        It return dictionary with the corresponding options for
401
        a given target.
402
        """
403
        return self.scans_table[scan_id]['target'].get('options')
404
405
    def get_vts(self, scan_id: str) -> Dict:
406
        """ Get a scan's vts. """
407
408
        return self.scans_table[scan_id]['vts']
409
410
    def release_vts_list(self, scan_id: str) -> None:
411
        """ Release the memory used for the vts list. """
412
413
        scan_data = self.scans_table.get(scan_id)
414
        if scan_data and 'vts' in scan_data:
415
            del scan_data['vts']
416
417
    def id_exists(self, scan_id: str) -> bool:
418
        """ Check whether a scan exists in the table. """
419
420
        return self.scans_table.get(scan_id) is not None
421
422
    def delete_scan(self, scan_id: str) -> bool:
423
        """ Delete a scan if fully finished. """
424
425
        if self.get_status(scan_id) == ScanStatus.RUNNING:
426
            return False
427
428
        self.scans_table.pop(scan_id)
429
430
        if len(self.scans_table) == 0:
431
            del self.data_manager
432
            self.data_manager = None
433
434
        return True
435