Passed
Pull Request — master (#126)
by Juan José
01:32
created

ospd.misc.create_pid()   A

Complexity

Conditions 4

Size

Total Lines 19
Code Lines 13

Duplication

Lines 0
Ratio 0 %

Importance

Changes 0
Metric Value
cc 4
eloc 13
nop 1
dl 0
loc 19
rs 9.75
c 0
b 0
f 0
1
# Copyright (C) 2014-2018 Greenbone Networks GmbH
2
#
3
# SPDX-License-Identifier: GPL-2.0-or-later
4
#
5
# This program is free software; you can redistribute it and/or
6
# modify it under the terms of the GNU General Public License
7
# as published by the Free Software Foundation; either version 2
8
# of the 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 General Public License for more details.
14
#
15
# You should have received a copy of the GNU General Public License
16
# along with this program; if not, write to the Free Software
17
# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
18
19
# pylint: disable=too-many-lines
20
21
""" Miscellaneous classes and functions related to OSPD.
22
"""
23
24
import logging
25
import os
26
import sys
27
import time
28
import uuid
29
import multiprocessing
30
31
from enum import Enum
32
from collections import OrderedDict
33
from pathlib import Path
34
35
from ospd.network import target_str_to_list
36
37
LOGGER = logging.getLogger(__name__)
38
39
40
class ScanStatus(Enum):
41
    """Scan status. """
42
43
    INIT = 0
44
    RUNNING = 1
45
    STOPPED = 2
46
    FINISHED = 3
47
48
49
class ScanCollection(object):
50
51
    """ Scans collection, managing scans and results read and write, exposing
52
    only needed information.
53
54
    Each scan has meta-information such as scan ID, current progress (from 0 to
55
    100), start time, end time, scan target and options and a list of results.
56
57
    There are 4 types of results: Alarms, Logs, Errors and Host Details.
58
59
    Todo:
60
    - Better checking for Scan ID existence and handling otherwise.
61
    - More data validation.
62
    - Mutex access per table/scan_info.
63
64
    """
65
66
    def __init__(self):
67
        """ Initialize the Scan Collection. """
68
69
        self.data_manager = None
70
        self.scans_table = dict()
71
72
    def add_result(
73
        self,
74
        scan_id,
75
        result_type,
76
        host='',
77
        hostname='',
78
        name='',
79
        value='',
80
        port='',
81
        test_id='',
82
        severity='',
83
        qod='',
84
    ):
85
        """ Add a result to a scan in the table. """
86
87
        assert scan_id
88
        assert len(name) or len(value)
89
        result = OrderedDict()
90
        result['type'] = result_type
91
        result['name'] = name
92
        result['severity'] = severity
93
        result['test_id'] = test_id
94
        result['value'] = value
95
        result['host'] = host
96
        result['hostname'] = hostname
97
        result['port'] = port
98
        result['qod'] = qod
99
        results = self.scans_table[scan_id]['results']
100
        results.append(result)
101
        # Set scan_info's results to propagate results to parent process.
102
        self.scans_table[scan_id]['results'] = results
103
104
    def set_progress(self, scan_id, progress):
105
        """ Sets scan_id scan's progress. """
106
107
        if progress > 0 and progress <= 100:
108
            self.scans_table[scan_id]['progress'] = progress
109
        if progress == 100:
110
            self.scans_table[scan_id]['end_time'] = int(time.time())
111
112
    def set_target_progress(self, scan_id, target, host, progress):
113
        """ Sets scan_id scan's progress. """
114
        if progress > 0 and progress <= 100:
115
            targets = self.scans_table[scan_id]['target_progress']
116
            targets[target][host] = progress
117
            # Set scan_info's target_progress to propagate progresses
118
            # to parent process.
119
            self.scans_table[scan_id]['target_progress'] = targets
120
121
    def set_host_finished(self, scan_id, target, host):
122
        """ Add the host in a list of finished hosts """
123
        finished_hosts = self.scans_table[scan_id]['finished_hosts']
124
        finished_hosts[target].append(host)
125
        self.scans_table[scan_id]['finished_hosts'] = finished_hosts
126
127
    def get_hosts_unfinished(self, scan_id):
128
        """ Get a list of unfinished hosts."""
129
130
        unfinished_hosts = list()
131
        for target in self.scans_table[scan_id]['finished_hosts']:
132
            unfinished_hosts.extend(target_str_to_list(target))
133
        for target in self.scans_table[scan_id]['finished_hosts']:
134
            for host in self.scans_table[scan_id]['finished_hosts'][target]:
135
                unfinished_hosts.remove(host)
136
137
        return unfinished_hosts
138
139
    def get_hosts_finished(self, scan_id):
140
        """ Get a list of finished hosts."""
141
142
        finished_hosts = list()
143
        for target in self.scans_table[scan_id]['finished_hosts']:
144
            finished_hosts.extend(
145
                self.scans_table[scan_id]['finished_hosts'].get(target)
146
            )
147
148
        return finished_hosts
149
150
    def results_iterator(self, scan_id, pop_res):
151
        """ Returns an iterator over scan_id scan's results. If pop_res is True,
152
        it removed the fetched results from the list.
153
        """
154
        if pop_res:
155
            result_aux = self.scans_table[scan_id]['results']
156
            self.scans_table[scan_id]['results'] = list()
157
            return iter(result_aux)
158
159
        return iter(self.scans_table[scan_id]['results'])
160
161
    def ids_iterator(self):
162
        """ Returns an iterator over the collection's scan IDS. """
163
164
        return iter(self.scans_table.keys())
165
166
    def remove_single_result(self, scan_id, result):
167
        """Removes a single result from the result list in scan_table.
168
169
        Parameters:
170
            scan_id (uuid): Scan ID to identify the scan process to be resumed.
171
            result (dict): The result to be removed from the results list.
172
        """
173
        results = self.scans_table[scan_id]['results']
174
        results.remove(result)
175
        self.scans_table[scan_id]['results'] = results
176
177
    def del_results_for_stopped_hosts(self, scan_id):
178
        """ Remove results from the result table for those host
179
        """
180
        unfinished_hosts = self.get_hosts_unfinished(scan_id)
181
        for result in self.results_iterator(scan_id, False):
182
            if result['host'] in unfinished_hosts:
183
                self.remove_single_result(scan_id, result)
184
185
    def resume_scan(self, scan_id, options):
186
        """ Reset the scan status in the scan_table to INIT.
187
        Also, overwrite the options, because a resume task cmd
188
        can add some new option. E.g. exclude hosts list.
189
        Parameters:
190
            scan_id (uuid): Scan ID to identify the scan process to be resumed.
191
            options (dict): Options for the scan to be resumed. This options
192
                            are not added to the already existent ones.
193
                            The old ones are removed
194
195
        Return:
196
            Scan ID which identifies the current scan.
197
        """
198
        self.scans_table[scan_id]['status'] = ScanStatus.INIT
199
        if options:
200
            self.scans_table[scan_id]['options'] = options
201
202
        self.del_results_for_stopped_hosts(scan_id)
203
204
        return scan_id
205
206
    def create_scan(self, scan_id='', targets='', options=None, vts=''):
207
        """ Creates a new scan with provided scan information. """
208
209
        if self.data_manager is None:
210
            self.data_manager = multiprocessing.Manager()
211
212
        # Check if it is possible to resume task. To avoid to resume, the
213
        # scan must be deleted from the scans_table.
214
        if (
215
            scan_id
216
            and self.id_exists(scan_id)
217
            and (self.get_status(scan_id) == ScanStatus.STOPPED)
218
        ):
219
            return self.resume_scan(scan_id, options)
220
221
        if not options:
222
            options = dict()
223
        scan_info = self.data_manager.dict()
224
        scan_info['results'] = list()
225
        scan_info['finished_hosts'] = dict(
226
            [[target, []] for target, _, _, _ in targets]
227
        )
228
        scan_info['progress'] = 0
229
        scan_info['target_progress'] = dict(
230
            [[target, {}] for target, _, _, _ in targets]
231
        )
232
        scan_info['targets'] = targets
233
        scan_info['vts'] = vts
234
        scan_info['options'] = options
235
        scan_info['start_time'] = int(time.time())
236
        scan_info['end_time'] = "0"
237
        scan_info['status'] = ScanStatus.INIT
238
        if scan_id is None or scan_id == '':
239
            scan_id = str(uuid.uuid4())
240
        scan_info['scan_id'] = scan_id
241
        self.scans_table[scan_id] = scan_info
242
        return scan_id
243
244
    def set_status(self, scan_id, status):
245
        """ Sets scan_id scan's status. """
246
        self.scans_table[scan_id]['status'] = status
247
248
    def get_status(self, scan_id):
249
        """ Get scan_id scans's status."""
250
251
        return self.scans_table[scan_id]['status']
252
253
    def get_options(self, scan_id):
254
        """ Get scan_id scan's options list. """
255
256
        return self.scans_table[scan_id]['options']
257
258
    def set_option(self, scan_id, name, value):
259
        """ Set a scan_id scan's name option to value. """
260
261
        self.scans_table[scan_id]['options'][name] = value
262
263
    def get_progress(self, scan_id):
264
        """ Get a scan's current progress value. """
265
266
        return self.scans_table[scan_id]['progress']
267
268
    def get_target_progress(self, scan_id, target):
269
        """ Get a target's current progress value.
270
        The value is calculated with the progress of each single host
271
        in the target."""
272
273
        total_hosts = len(target_str_to_list(target))
274
        host_progresses = self.scans_table[scan_id]['target_progress'].get(
275
            target
276
        )
277
        try:
278
            t_prog = sum(host_progresses.values()) / total_hosts
279
        except ZeroDivisionError:
280
            LOGGER.error(
281
                "Zero division error in %s", self.get_target_progress.__name__
282
            )
283
            raise
284
        return t_prog
285
286
    def get_start_time(self, scan_id):
287
        """ Get a scan's start time. """
288
289
        return self.scans_table[scan_id]['start_time']
290
291
    def get_end_time(self, scan_id):
292
        """ Get a scan's end time. """
293
294
        return self.scans_table[scan_id]['end_time']
295
296
    def get_target_list(self, scan_id):
297
        """ Get a scan's target list. """
298
299
        target_list = []
300
        for target, _, _, _ in self.scans_table[scan_id]['targets']:
301
            target_list.append(target)
302
        return target_list
303
304
    def get_ports(self, scan_id, target):
305
        """ Get a scan's ports list. If a target is specified
306
        it will return the corresponding port for it. If not,
307
        it returns the port item of the first nested list in
308
        the target's list.
309
        """
310
        if target:
311
            for item in self.scans_table[scan_id]['targets']:
312
                if target == item[0]:
313
                    return item[1]
314
315
        return self.scans_table[scan_id]['targets'][0][1]
316
317
    def get_exclude_hosts(self, scan_id, target):
318
        """ Get an exclude host list for a given target.
319
        """
320
        if target:
321
            for item in self.scans_table[scan_id]['targets']:
322
                if target == item[0]:
323
                    return item[3]
324
325
    def get_credentials(self, scan_id, target):
326
        """ Get a scan's credential list. It return dictionary with
327
        the corresponding credential for a given target.
328
        """
329
        if target:
330
            for item in self.scans_table[scan_id]['targets']:
331
                if target == item[0]:
332
                    return item[2]
333
334
    def get_vts(self, scan_id):
335
        """ Get a scan's vts list. """
336
337
        return self.scans_table[scan_id]['vts']
338
339
    def id_exists(self, scan_id):
340
        """ Check whether a scan exists in the table. """
341
342
        return self.scans_table.get(scan_id) is not None
343
344
    def delete_scan(self, scan_id):
345
        """ Delete a scan if fully finished. """
346
347
        if self.get_status(scan_id) == ScanStatus.RUNNING:
348
            return False
349
        self.scans_table.pop(scan_id)
350
        if len(self.scans_table) == 0:
351
            del self.data_manager
352
            self.data_manager = None
353
        return True
354
355
356
class ResultType(object):
357
358
    """ Various scan results types values. """
359
360
    ALARM = 0
361
    LOG = 1
362
    ERROR = 2
363
    HOST_DETAIL = 3
364
365
    @classmethod
366
    def get_str(cls, result_type):
367
        """ Return string name of a result type. """
368
        if result_type == cls.ALARM:
369
            return "Alarm"
370
        elif result_type == cls.LOG:
371
            return "Log Message"
372
        elif result_type == cls.ERROR:
373
            return "Error Message"
374
        elif result_type == cls.HOST_DETAIL:
375
            return "Host Detail"
376
        else:
377
            assert False, "Erroneous result type {0}.".format(result_type)
378
379
    @classmethod
380
    def get_type(cls, result_name):
381
        """ Return string name of a result type. """
382
        if result_name == "Alarm":
383
            return cls.ALARM
384
        elif result_name == "Log Message":
385
            return cls.LOG
386
        elif result_name == "Error Message":
387
            return cls.ERROR
388
        elif result_name == "Host Detail":
389
            return cls.HOST_DETAIL
390
        else:
391
            assert False, "Erroneous result name {0}.".format(result_name)
392
393
394
def valid_uuid(value):
395
    """ Check if value is a valid UUID. """
396
397
    try:
398
        uuid.UUID(value, version=4)
399
        return True
400
    except (TypeError, ValueError, AttributeError):
401
        return False
402
403
404
def go_to_background():
405
    """ Daemonize the running process. """
406
    try:
407
        if os.fork():
408
            sys.exit()
409
    except OSError as errmsg:
410
        LOGGER.error('Fork failed: %s', errmsg)
411
        sys.exit(1)
412
413
def create_pid(pidfile):
414
    """ Check if there is an already running daemon and creates the pid file.
415
    Otherwise gives an error. """
416
417
    pid = str(os.getpid())
418
419
    if Path(pidfile).is_file():
420
        LOGGER.error("There is an already running process.")
421
        return False
422
423
    try:
424
        with open(pidfile, 'w') as f:
425
            f.write(pid)
426
    except (FileNotFoundError, PermissionError) as e:
427
        msg = "Failed to create pid file %s. %s" % (os.path.dirname(pidfile), e)
428
        LOGGER.error(msg)
429
        return False
430
431
    return True
432
433
def remove_pidfile(pidfile, signum=None, frame=None):
434
    """ Removes the pidfile before ending the daemon. """
435
    pidpath = Path(pidfile)
436
    if pidpath.is_file():
437
        LOGGER.debug("Finishing daemon process")
438
        pidpath.unlink()
439
        sys.exit()
440