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
|
|
|
) -> None: |
86
|
|
|
""" Add a result to a scan in the table. """ |
87
|
|
|
|
88
|
|
|
assert scan_id |
89
|
|
|
assert len(name) or len(value) |
90
|
|
|
|
91
|
|
|
result = OrderedDict() # type: Dict |
92
|
|
|
result['type'] = result_type |
93
|
|
|
result['name'] = name |
94
|
|
|
result['severity'] = severity |
95
|
|
|
result['test_id'] = test_id |
96
|
|
|
result['value'] = value |
97
|
|
|
result['host'] = host |
98
|
|
|
result['hostname'] = hostname |
99
|
|
|
result['port'] = port |
100
|
|
|
result['qod'] = qod |
101
|
|
|
results = self.scans_table[scan_id]['results'] |
102
|
|
|
results.append(result) |
103
|
|
|
|
104
|
|
|
# Set scan_info's results to propagate results to parent process. |
105
|
|
|
self.scans_table[scan_id]['results'] = results |
106
|
|
|
|
107
|
|
|
def add_result_list( |
108
|
|
|
self, scan_id: str, result_list: Iterable[Dict[str, str]] |
109
|
|
|
) -> None: |
110
|
|
|
""" |
111
|
|
|
Add a batch of results to the result's table for the corresponding |
112
|
|
|
scan_id |
113
|
|
|
""" |
114
|
|
|
results = self.scans_table[scan_id]['results'] |
115
|
|
|
results.extend(result_list) |
116
|
|
|
|
117
|
|
|
# Set scan_info's results to propagate results to parent process. |
118
|
|
|
self.scans_table[scan_id]['results'] = results |
119
|
|
|
|
120
|
|
|
def remove_hosts_from_target_progress( |
121
|
|
|
self, scan_id: str, hosts: List |
122
|
|
|
) -> None: |
123
|
|
|
"""Remove a list of hosts from the main scan progress table to avoid |
124
|
|
|
the hosts to be included in the calculation of the scan progress""" |
125
|
|
|
if not hosts: |
126
|
|
|
return |
127
|
|
|
|
128
|
|
|
target = self.scans_table[scan_id].get('target_progress') |
129
|
|
|
for host in hosts: |
130
|
|
|
if host in target: |
131
|
|
|
del target[host] |
132
|
|
|
|
133
|
|
|
# Set scan_info's target_progress to propagate progresses |
134
|
|
|
# to parent process. |
135
|
|
|
self.scans_table[scan_id]['target_progress'] = target |
136
|
|
|
|
137
|
|
|
def set_progress(self, scan_id: str, progress: int) -> None: |
138
|
|
|
""" Sets scan_id scan's progress. """ |
139
|
|
|
|
140
|
|
|
if progress > 0 and progress <= 100: |
141
|
|
|
self.scans_table[scan_id]['progress'] = progress |
142
|
|
|
|
143
|
|
|
if progress == 100: |
144
|
|
|
self.scans_table[scan_id]['end_time'] = int(time.time()) |
145
|
|
|
|
146
|
|
|
def set_host_progress( |
147
|
|
|
self, scan_id: str, host_progress_batch: Dict[str, int] |
148
|
|
|
) -> None: |
149
|
|
|
""" Sets scan_id scan's progress. """ |
150
|
|
|
|
151
|
|
|
host_progresses = self.scans_table[scan_id].get('target_progress') |
152
|
|
|
host_progresses.update(host_progress_batch) |
153
|
|
|
|
154
|
|
|
# Set scan_info's target_progress to propagate progresses |
155
|
|
|
# to parent process. |
156
|
|
|
self.scans_table[scan_id]['target_progress'] = host_progresses |
157
|
|
|
|
158
|
|
|
def set_host_finished(self, scan_id: str, hosts: List[str]) -> None: |
159
|
|
|
""" Increase the amount of finished hosts which were alive.""" |
160
|
|
|
|
161
|
|
|
total_finished = len(hosts) |
162
|
|
|
count_alive = ( |
163
|
|
|
self.scans_table[scan_id].get('count_alive') + total_finished |
164
|
|
|
) |
165
|
|
|
self.scans_table[scan_id]['count_alive'] = count_alive |
166
|
|
|
|
167
|
|
|
def set_host_dead(self, scan_id: str, hosts: List[str]) -> None: |
168
|
|
|
""" Increase the amount of dead hosts. """ |
169
|
|
|
|
170
|
|
|
total_dead = len(hosts) |
171
|
|
|
count_dead = self.scans_table[scan_id].get('count_dead') + total_dead |
172
|
|
|
self.scans_table[scan_id]['count_dead'] = count_dead |
173
|
|
|
|
174
|
|
|
def set_amount_dead_hosts(self, scan_id: str, total_dead: int) -> None: |
175
|
|
|
""" Increase the amount of dead hosts. """ |
176
|
|
|
|
177
|
|
|
count_dead = self.scans_table[scan_id].get('count_dead') + total_dead |
178
|
|
|
self.scans_table[scan_id]['count_dead'] = count_dead |
179
|
|
|
|
180
|
|
|
def results_iterator( |
181
|
|
|
self, scan_id: str, pop_res: bool = False, max_res: int = None |
182
|
|
|
) -> Iterator[Any]: |
183
|
|
|
""" Returns an iterator over scan_id scan's results. If pop_res is True, |
184
|
|
|
it removed the fetched results from the list. |
185
|
|
|
|
186
|
|
|
If max_res is None, return all the results. |
187
|
|
|
Otherwise, if max_res = N > 0 return N as maximum number of results. |
188
|
|
|
|
189
|
|
|
max_res works only together with pop_results. |
190
|
|
|
""" |
191
|
|
|
if pop_res and max_res: |
192
|
|
|
result_aux = self.scans_table[scan_id]['results'] |
193
|
|
|
self.scans_table[scan_id]['results'] = result_aux[max_res:] |
194
|
|
|
return iter(result_aux[:max_res]) |
195
|
|
|
elif pop_res: |
196
|
|
|
result_aux = self.scans_table[scan_id]['results'] |
197
|
|
|
self.scans_table[scan_id]['results'] = list() |
198
|
|
|
return iter(result_aux) |
199
|
|
|
|
200
|
|
|
return iter(self.scans_table[scan_id]['results']) |
201
|
|
|
|
202
|
|
|
def ids_iterator(self) -> Iterator[str]: |
203
|
|
|
""" Returns an iterator over the collection's scan IDS. """ |
204
|
|
|
|
205
|
|
|
return iter(self.scans_table.keys()) |
206
|
|
|
|
207
|
|
|
def clean_up_pickled_scan_info(self) -> None: |
208
|
|
|
""" Remove files of pickled scan info """ |
209
|
|
|
for scan_id in self.ids_iterator(): |
210
|
|
|
if self.get_status(scan_id) == ScanStatus.QUEUED: |
211
|
|
|
self.remove_file_pickled_scan_info(scan_id) |
212
|
|
|
|
213
|
|
|
def remove_file_pickled_scan_info(self, scan_id: str) -> None: |
214
|
|
|
pickler = DataPickler(self.file_storage_dir) |
215
|
|
|
pickler.remove_file(scan_id) |
216
|
|
|
|
217
|
|
|
def unpickle_scan_info(self, scan_id: str) -> None: |
218
|
|
|
""" Unpickle a stored scan_inf correspinding to the scan_id |
219
|
|
|
and store it in the scan_table """ |
220
|
|
|
|
221
|
|
|
scan_info = self.scans_table.get(scan_id) |
222
|
|
|
scan_info_hash = scan_info.pop('scan_info_hash') |
223
|
|
|
|
224
|
|
|
pickler = DataPickler(self.file_storage_dir) |
225
|
|
|
unpickled_scan_info = pickler.load_data(scan_id, scan_info_hash) |
226
|
|
|
|
227
|
|
|
if not unpickled_scan_info: |
228
|
|
|
pickler.remove_file(scan_id) |
229
|
|
|
raise OspdCommandError( |
230
|
|
|
'Not possible to unpickle stored scan info for %s' % scan_id, |
231
|
|
|
'start_scan', |
232
|
|
|
) |
233
|
|
|
|
234
|
|
|
scan_info['results'] = list() |
235
|
|
|
scan_info['progress'] = 0 |
236
|
|
|
scan_info['target_progress'] = dict() |
237
|
|
|
scan_info['count_alive'] = 0 |
238
|
|
|
scan_info['count_dead'] = 0 |
239
|
|
|
scan_info['target'] = unpickled_scan_info.pop('target') |
240
|
|
|
scan_info['vts'] = unpickled_scan_info.pop('vts') |
241
|
|
|
scan_info['options'] = unpickled_scan_info.pop('options') |
242
|
|
|
scan_info['start_time'] = int(time.time()) |
243
|
|
|
scan_info['end_time'] = 0 |
244
|
|
|
|
245
|
|
|
self.scans_table[scan_id] = scan_info |
246
|
|
|
|
247
|
|
|
pickler.remove_file(scan_id) |
248
|
|
|
|
249
|
|
|
def create_scan( |
250
|
|
|
self, |
251
|
|
|
scan_id: str = '', |
252
|
|
|
target: Dict = None, |
253
|
|
|
options: Optional[Dict] = None, |
254
|
|
|
vts: Dict = None, |
255
|
|
|
) -> str: |
256
|
|
|
""" Creates a new scan with provided scan information. """ |
257
|
|
|
|
258
|
|
|
if not options: |
259
|
|
|
options = dict() |
260
|
|
|
|
261
|
|
|
credentials = target.pop('credentials') |
262
|
|
|
|
263
|
|
|
scan_info = self.data_manager.dict() # type: Dict |
264
|
|
|
scan_info['status'] = ScanStatus.QUEUED |
265
|
|
|
scan_info['credentials'] = credentials |
266
|
|
|
scan_info['start_time'] = int(time.time()) |
267
|
|
|
|
268
|
|
|
scan_info_to_pickle = { |
269
|
|
|
'target': target, |
270
|
|
|
'options': options, |
271
|
|
|
'vts': vts, |
272
|
|
|
} |
273
|
|
|
|
274
|
|
|
if scan_id is None or scan_id == '': |
275
|
|
|
scan_id = str(uuid.uuid4()) |
276
|
|
|
|
277
|
|
|
pickler = DataPickler(self.file_storage_dir) |
278
|
|
|
scan_info_hash = None |
279
|
|
|
try: |
280
|
|
|
scan_info_hash = pickler.store_data(scan_id, scan_info_to_pickle) |
281
|
|
|
except OspdCommandError as e: |
282
|
|
|
LOGGER.error(e) |
283
|
|
|
return |
284
|
|
|
|
285
|
|
|
scan_info['scan_id'] = scan_id |
286
|
|
|
scan_info['scan_info_hash'] = scan_info_hash |
287
|
|
|
|
288
|
|
|
self.scans_table[scan_id] = scan_info |
289
|
|
|
return scan_id |
290
|
|
|
|
291
|
|
|
def set_status(self, scan_id: str, status: ScanStatus) -> None: |
292
|
|
|
""" Sets scan_id scan's status. """ |
293
|
|
|
self.scans_table[scan_id]['status'] = status |
294
|
|
|
if status == ScanStatus.STOPPED: |
295
|
|
|
self.scans_table[scan_id]['end_time'] = int(time.time()) |
296
|
|
|
|
297
|
|
|
def get_status(self, scan_id: str) -> ScanStatus: |
298
|
|
|
""" Get scan_id scans's status.""" |
299
|
|
|
|
300
|
|
|
return self.scans_table[scan_id].get('status') |
301
|
|
|
|
302
|
|
|
def get_options(self, scan_id: str) -> Dict: |
303
|
|
|
""" Get scan_id scan's options list. """ |
304
|
|
|
|
305
|
|
|
return self.scans_table[scan_id].get('options') |
306
|
|
|
|
307
|
|
|
def set_option(self, scan_id, name: str, value: Any) -> None: |
308
|
|
|
""" Set a scan_id scan's name option to value. """ |
309
|
|
|
|
310
|
|
|
self.scans_table[scan_id]['options'][name] = value |
311
|
|
|
|
312
|
|
|
def get_progress(self, scan_id: str) -> int: |
313
|
|
|
""" Get a scan's current progress value. """ |
314
|
|
|
|
315
|
|
|
return self.scans_table[scan_id].get('progress', 0) |
316
|
|
|
|
317
|
|
|
def get_count_dead(self, scan_id: str) -> int: |
318
|
|
|
""" Get a scan's current dead host count. """ |
319
|
|
|
|
320
|
|
|
return self.scans_table[scan_id]['count_dead'] |
321
|
|
|
|
322
|
|
|
def get_count_alive(self, scan_id: str) -> int: |
323
|
|
|
""" Get a scan's current dead host count. """ |
324
|
|
|
|
325
|
|
|
return self.scans_table[scan_id]['count_alive'] |
326
|
|
|
|
327
|
|
|
def get_current_target_progress(self, scan_id: str) -> Dict[str, int]: |
328
|
|
|
""" Get a scan's current hosts progress """ |
329
|
|
|
return self.scans_table[scan_id]['target_progress'] |
330
|
|
|
|
331
|
|
|
def simplify_exclude_host_count(self, scan_id: str) -> int: |
332
|
|
|
""" Remove from exclude_hosts the received hosts in the finished_hosts |
333
|
|
|
list sent by the client. |
334
|
|
|
The finished hosts are sent also as exclude hosts for backward |
335
|
|
|
compatibility purposses. |
336
|
|
|
|
337
|
|
|
Return: |
338
|
|
|
Count of excluded host. |
339
|
|
|
""" |
340
|
|
|
|
341
|
|
|
exc_hosts_list = target_str_to_list(self.get_exclude_hosts(scan_id)) |
342
|
|
|
|
343
|
|
|
finished_hosts_list = target_str_to_list( |
344
|
|
|
self.get_finished_hosts(scan_id) |
345
|
|
|
) |
346
|
|
|
|
347
|
|
|
if finished_hosts_list and exc_hosts_list: |
348
|
|
|
for finished in finished_hosts_list: |
349
|
|
|
if finished in exc_hosts_list: |
350
|
|
|
exc_hosts_list.remove(finished) |
351
|
|
|
|
352
|
|
|
return len(exc_hosts_list) if exc_hosts_list else 0 |
353
|
|
|
|
354
|
|
|
def calculate_target_progress(self, scan_id: str) -> int: |
355
|
|
|
""" Get a target's current progress value. |
356
|
|
|
The value is calculated with the progress of each single host |
357
|
|
|
in the target.""" |
358
|
|
|
|
359
|
|
|
total_hosts = self.get_host_count(scan_id) |
360
|
|
|
exc_hosts = self.simplify_exclude_host_count(scan_id) |
361
|
|
|
count_alive = self.get_count_alive(scan_id) |
362
|
|
|
count_dead = self.get_count_dead(scan_id) |
363
|
|
|
host_progresses = self.get_current_target_progress(scan_id) |
364
|
|
|
|
365
|
|
|
try: |
366
|
|
|
t_prog = int( |
367
|
|
|
(sum(host_progresses.values()) + 100 * count_alive) |
368
|
|
|
/ (total_hosts - exc_hosts - count_dead) |
369
|
|
|
) |
370
|
|
|
except ZeroDivisionError: |
371
|
|
|
LOGGER.error( |
372
|
|
|
"Zero division error in %s", |
373
|
|
|
self.calculate_target_progress.__name__, |
374
|
|
|
) |
375
|
|
|
raise |
376
|
|
|
|
377
|
|
|
return t_prog |
378
|
|
|
|
379
|
|
|
def get_start_time(self, scan_id: str) -> str: |
380
|
|
|
""" Get a scan's start time. """ |
381
|
|
|
|
382
|
|
|
return self.scans_table[scan_id]['start_time'] |
383
|
|
|
|
384
|
|
|
def get_end_time(self, scan_id: str) -> str: |
385
|
|
|
""" Get a scan's end time. """ |
386
|
|
|
|
387
|
|
|
return self.scans_table[scan_id]['end_time'] |
388
|
|
|
|
389
|
|
|
def get_host_list(self, scan_id: str) -> Dict: |
390
|
|
|
""" Get a scan's host list. """ |
391
|
|
|
|
392
|
|
|
return self.scans_table[scan_id]['target'].get('hosts') |
393
|
|
|
|
394
|
|
|
def get_host_count(self, scan_id: str) -> int: |
395
|
|
|
""" Get total host count in the target. """ |
396
|
|
|
host = self.get_host_list(scan_id) |
397
|
|
|
total_hosts = len(target_str_to_list(host)) |
398
|
|
|
|
399
|
|
|
return total_hosts |
400
|
|
|
|
401
|
|
|
def get_ports(self, scan_id: str) -> str: |
402
|
|
|
""" Get a scan's ports list. |
403
|
|
|
""" |
404
|
|
|
target = self.scans_table[scan_id].get('target') |
405
|
|
|
ports = target.pop('ports') |
406
|
|
|
self.scans_table[scan_id]['target'] = target |
407
|
|
|
return ports |
408
|
|
|
|
409
|
|
|
def get_exclude_hosts(self, scan_id: str) -> str: |
410
|
|
|
""" Get an exclude host list for a given target. |
411
|
|
|
""" |
412
|
|
|
return self.scans_table[scan_id]['target'].get('exclude_hosts') |
413
|
|
|
|
414
|
|
|
def get_finished_hosts(self, scan_id: str) -> str: |
415
|
|
|
""" Get the finished host list sent by the client for a given target. |
416
|
|
|
""" |
417
|
|
|
return self.scans_table[scan_id]['target'].get('finished_hosts') |
418
|
|
|
|
419
|
|
|
def get_credentials(self, scan_id: str) -> Dict[str, Dict[str, str]]: |
420
|
|
|
""" Get a scan's credential list. It return dictionary with |
421
|
|
|
the corresponding credential for a given target. |
422
|
|
|
""" |
423
|
|
|
return self.scans_table[scan_id].get('credentials') |
424
|
|
|
|
425
|
|
|
def get_target_options(self, scan_id: str) -> Dict[str, str]: |
426
|
|
|
""" Get a scan's target option dictionary. |
427
|
|
|
It return dictionary with the corresponding options for |
428
|
|
|
a given target. |
429
|
|
|
""" |
430
|
|
|
return self.scans_table[scan_id]['target'].get('options') |
431
|
|
|
|
432
|
|
|
def get_vts(self, scan_id: str) -> Dict[str, Union[Dict[str, str], List]]: |
433
|
|
|
""" Get a scan's vts. """ |
434
|
|
|
scan_info = self.scans_table[scan_id] |
435
|
|
|
vts = scan_info.pop('vts') |
436
|
|
|
self.scans_table[scan_id] = scan_info |
437
|
|
|
|
438
|
|
|
return vts |
439
|
|
|
|
440
|
|
|
def id_exists(self, scan_id: str) -> bool: |
441
|
|
|
""" Check whether a scan exists in the table. """ |
442
|
|
|
|
443
|
|
|
return self.scans_table.get(scan_id) is not None |
444
|
|
|
|
445
|
|
|
def delete_scan(self, scan_id: str) -> bool: |
446
|
|
|
""" Delete a scan if fully finished. """ |
447
|
|
|
|
448
|
|
|
if self.get_status(scan_id) == ScanStatus.RUNNING: |
449
|
|
|
return False |
450
|
|
|
|
451
|
|
|
scans_table = self.scans_table |
452
|
|
|
del scans_table[scan_id] |
453
|
|
|
self.scans_table = scans_table |
454
|
|
|
|
455
|
|
|
return True |
456
|
|
|
|