ospd.command.command.StartScan.is_new_scan_allowed() - Code Metrics - Inspection of "Merge pull request #245 from jjnicola/max-scan" - greenbone/ospd - Measure and Improve Code Quality continuously with Scrutinizer

Completed

Push — master ( f571c7...585a62 )

unknown

created 2020-04-09 10:49 UTC

StartScan.is_new_scan_allowed() A

↳ Parent: ospd.command.command

Complexity

Conditions

Size

Total Lines	12
Code Lines	5

Duplication

Lines	0
Ratio	0 %

Importance

Changes

Metric	Value
cc	3
eloc	5
nop	1
dl	0
loc	12
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 multiprocessing
import re
import logging
import subprocess

from decimal import Decimal
from typing import Optional, Dict, Any, Union, Iterator

from xml.etree.ElementTree import Element, SubElement

import psutil

from ospd.errors import OspdCommandError
from ospd.misc import valid_uuid, create_process
from ospd.protocol import OspRequest, OspResponse
from ospd.xml import (
    simple_response_str,
    get_elements_from_dict,
    XmlStringHelper,
)

from .initsubclass import InitSubclassMeta
from .registry import register_command

logger = logging.getLogger(__name__)


class BaseCommand(metaclass=InitSubclassMeta):

    name = None
    description = None
    attributes = None
    elements = None
    must_be_initialized = None

    def __init_subclass__(cls, **kwargs):
        super_cls = super()

        if hasattr(super_cls, '__init_subclass__'):
            super_cls.__init_subclass__(**kwargs)

        register_command(cls)

    def __init__(self, daemon):
        self._daemon = daemon

    def get_name(self) -> str:
        return self.name

    def get_description(self) -> str:
        return self.description

    def get_attributes(self) -> Optional[Dict[str, Any]]:
        return self.attributes

    def get_elements(self) -> Optional[Dict[str, Any]]:
        return self.elements

    def handle_xml(
        self, xml: Element, initialized: bool
    ) -> Union[bytes, Iterator[bytes]]:
        raise NotImplementedError()

    def as_dict(self):
        return {
            'name': self.get_name(),
            'attributes': self.get_attributes(),
            'description': self.get_description(),
            'elements': self.get_elements(),
        }

    def __repr__(self):
        return '<{} description="{}" attributes={} elements={}>'.format(
            self.name, self.description, self.attributes, self.elements
        )


class HelpCommand(BaseCommand):
    name = "help"
    description = 'Print the commands help.'
    attributes = {'format': 'Help format. Could be text or xml.'}
    must_be_initialized = False

    def handle_xml(self, xml: Element) -> bytes:
        help_format = xml.get('format')

        if help_format is None or help_format == "text":
            # Default help format is text.
            return simple_response_str(
                'help', 200, 'OK', self._daemon.get_help_text()
            )
        elif help_format == "xml":
            text = get_elements_from_dict(
                {k: v.as_dict() for k, v in self._daemon.commands.items()}
            )
            return simple_response_str('help', 200, 'OK', text)

        raise OspdCommandError('Bogus help format', 'help')


class GetVersion(BaseCommand):
    name = "get_version"
    description = 'Return various version information'
    must_be_initialized = False

    def handle_xml(self, xml: Element) -> bytes:
        """ Handles <get_version> command.

        Return:
            Response string for <get_version> command.
        """
        protocol = Element('protocol')

        for name, value in [
            ('name', 'OSP'),
            ('version', self._daemon.get_protocol_version()),
        ]:
            elem = SubElement(protocol, name)
            elem.text = value

        daemon = Element('daemon')
        for name, value in [
            ('name', self._daemon.get_daemon_name()),
            ('version', self._daemon.get_daemon_version()),
        ]:
            elem = SubElement(daemon, name)
            elem.text = value

        scanner = Element('scanner')
        for name, value in [
            ('name', self._daemon.get_scanner_name()),
            ('version', self._daemon.get_scanner_version()),
        ]:
            elem = SubElement(scanner, name)
            elem.text = value

        content = [protocol, daemon, scanner]

        vts_version = self._daemon.get_vts_version()
        if vts_version:
            vts = Element('vts')
            elem = SubElement(vts, 'version')
            elem.text = vts_version
            content.append(vts)

        return simple_response_str('get_version', 200, 'OK', content)


GVMCG_TITLES = [
    'cpu-*',
    'proc',
    'mem',
    'swap',
    'load',
    'df-*',
    'disk-sd[a-z][0-9]-rw',
    'disk-sd[a-z][0-9]-load',
    'disk-sd[a-z][0-9]-io-load',
    'interface-eth*-traffic',
    'interface-eth*-err-rate',
    'interface-eth*-err',
    'sensors-*_temperature-*',
    'sensors-*_fanspeed-*',
    'sensors-*_voltage-*',
    'titles',
]  # type: List[str]


class GetPerformance(BaseCommand):
    name = "get_performance"
    description = 'Return system report'
    attributes = {
        'start': 'Time of first data point in report.',
        'end': 'Time of last data point in report.',
        'title': 'Name of report.',
    }
    must_be_initialized = False

    def handle_xml(self, xml: Element) -> bytes:
        """ Handles <get_performance> command.

        @return: Response string for <get_performance> command.
        """
        start = xml.attrib.get('start')
        end = xml.attrib.get('end')
        titles = xml.attrib.get('titles')

        cmd = ['gvmcg']
        if start:
            try:
                int(start)
            except ValueError:
                raise OspdCommandError(
                    'Start argument must be integer.', 'get_performance'
                )

            cmd.append(start)

        if end:
            try:
                int(end)
            except ValueError:
                raise OspdCommandError(
                    'End argument must be integer.', 'get_performance'
                )

            cmd.append(end)

        if titles:
            combined = "(" + ")|(".join(GVMCG_TITLES) + ")"
            forbidden = "^[^|&;]+$"
            if re.match(combined, titles) and re.match(forbidden, titles):
                cmd.append(titles)
            else:
                raise OspdCommandError(
                    'Arguments not allowed', 'get_performance'
                )

        try:
            output = subprocess.check_output(cmd)
        except (subprocess.CalledProcessError, OSError) as e:
            raise OspdCommandError(
                'Bogus get_performance format. %s' % e, 'get_performance'
            )

        return simple_response_str(
            'get_performance', 200, 'OK', output.decode()
        )


class GetScannerDetails(BaseCommand):
    name = 'get_scanner_details'
    description = 'Return scanner description and parameters'
    must_be_initialized = True

    def handle_xml(self, xml: Element) -> bytes:
        """ Handles <get_scanner_details> command.

        @return: Response string for <get_scanner_details> command.
        """
        desc_xml = Element('description')
        desc_xml.text = self._daemon.get_scanner_description()
        scanner_params = self._daemon.get_scanner_params()
        details = [
            desc_xml,
            OspResponse.create_scanner_params_xml(scanner_params),
        ]
        return simple_response_str('get_scanner_details', 200, 'OK', details)


class DeleteScan(BaseCommand):
    name = 'delete_scan'
    description = 'Delete a finished scan.'
    attributes = {'scan_id': 'ID of scan to delete.'}
    must_be_initialized = False

    def handle_xml(self, xml: Element) -> bytes:
        """ Handles <delete_scan> command.

        @return: Response string for <delete_scan> command.
        """
        scan_id = xml.get('scan_id')
        if scan_id is None:
            return simple_response_str(
                'delete_scan', 404, 'No scan_id attribute'
            )

        if not self._daemon.scan_exists(scan_id):
            text = "Failed to find scan '{0}'".format(scan_id)
            return simple_response_str('delete_scan', 404, text)

        self._daemon.check_scan_process(scan_id)

        if self._daemon.delete_scan(scan_id):
            return simple_response_str('delete_scan', 200, 'OK')

        raise OspdCommandError('Scan in progress', 'delete_scan')


class GetVts(BaseCommand):
    name = 'get_vts'
    description = 'List of available vulnerability tests.'
    attributes = {
        'vt_id': 'ID of a specific vulnerability test to get.',
        'filter': 'Optional filter to get an specific vt collection.',
    }
    must_be_initialized = True

    def handle_xml(self, xml: Element) -> Iterator[bytes]:
        """ Handles <get_vts> command.
        Writes the vt collection on the stream.
        The <get_vts> element accept two optional arguments.
        vt_id argument receives a single vt id.
        filter argument receives a filter selecting a sub set of vts.
        If both arguments are given, the vts which match with the filter
        are return.

        @return: Response string for <get_vts> command on fail.
        """
        xml_helper = XmlStringHelper()

        vt_id = xml.get('vt_id')
        vt_filter = xml.get('filter')
        _details = xml.get('details')

        vt_details = False if _details == '0' else True

        if vt_id and vt_id not in self._daemon.vts:
            text = "Failed to find vulnerability test '{0}'".format(vt_id)
            raise OspdCommandError(text, 'get_vts', 404)

        filtered_vts = None
        if vt_filter:
            filtered_vts = self._daemon.vts_filter.get_filtered_vts_list(
                self._daemon.vts, vt_filter
            )
        vts_selection = self._daemon.get_vts_selection_list(vt_id, filtered_vts)
        # List of xml pieces with the generator to be iterated
        yield xml_helper.create_response('get_vts')

        begin_vts_tag = xml_helper.create_element('vts')
        val = len(self._daemon.vts)
        begin_vts_tag = xml_helper.add_attr(begin_vts_tag, "total", val)
        if filtered_vts:
            val = len(filtered_vts)
            begin_vts_tag = xml_helper.add_attr(begin_vts_tag, "sent", val)

        if self._daemon.vts.sha256_hash is not None:
            begin_vts_tag = xml_helper.add_attr(
                begin_vts_tag, "sha256_hash", self._daemon.vts.sha256_hash
            )

        yield begin_vts_tag

        for vt in self._daemon.get_vt_iterator(vts_selection, vt_details):
            yield xml_helper.add_element(self._daemon.get_vt_xml(vt))

        yield xml_helper.create_element('vts', end=True)
        yield xml_helper.create_response('get_vts', end=True)


class StopScan(BaseCommand):
    name = 'stop_scan'
    description = 'Stop a currently running scan.'
    attributes = {'scan_id': 'ID of scan to stop.'}
    must_be_initialized = True

    def handle_xml(self, xml: Element) -> bytes:
        """ Handles <stop_scan> command.

        @return: Response string for <stop_scan> command.
        """

        scan_id = xml.get('scan_id')
        if scan_id is None or scan_id == '':
            raise OspdCommandError('No scan_id attribute', 'stop_scan')

        self._daemon.stop_scan(scan_id)

        # Don't send response until the scan is stopped.
        try:
            self._daemon.scan_processes[scan_id].join()
        except KeyError:
            pass

        return simple_response_str('stop_scan', 200, 'OK')


class GetScans(BaseCommand):
    name = 'get_scans'
    description = 'List the scans in buffer.'
    attributes = {
        'scan_id': 'ID of a specific scan to get.',
        'details': 'Whether to return the full scan report.',
        'pop_results': 'Whether to remove the fetched results.',
        'max_results': 'Maximum number of results to fetch.',
    }
    must_be_initialized = False

    def handle_xml(self, xml: Element) -> bytes:
        """ Handles <get_scans> command.

        @return: Response string for <get_scans> command.
        """

        scan_id = xml.get('scan_id')
        details = xml.get('details')
        pop_res = xml.get('pop_results')
        max_res = xml.get('max_results')

        if details and details == '0':
            details = False
        else:
            details = True
            if pop_res and pop_res == '1':
                pop_res = True
            else:
                pop_res = False
            if max_res:
                max_res = int(max_res)

        responses = []
        if scan_id and scan_id in self._daemon.scan_collection.ids_iterator():
            self._daemon.check_scan_process(scan_id)
            scan = self._daemon.get_scan_xml(scan_id, details, pop_res, max_res)
            responses.append(scan)
        elif scan_id:
            text = "Failed to find scan '{0}'".format(scan_id)
            return simple_response_str('get_scans', 404, text)
        else:
            for scan_id in self._daemon.scan_collection.ids_iterator():
                self._daemon.check_scan_process(scan_id)
                scan = self._daemon.get_scan_xml(
                    scan_id, details, pop_res, max_res
                )
                responses.append(scan)

        return simple_response_str('get_scans', 200, 'OK', responses)


class StartScan(BaseCommand):
    name = 'start_scan'
    description = 'Start a new scan.'
    attributes = {
        'target': 'Target host to scan',
        'ports': 'Ports list to scan',
        'scan_id': 'Optional UUID value to use as scan ID',
        'parallel': 'Optional nummer of parallel target to scan',
    }
    must_be_initialized = True

    def get_elements(self):
        elements = {}

        if self.elements:
            elements.update(self.elements)

        scanner_params = elements.get('scanner_params', {}).copy()
        elements['scanner_params'] = scanner_params

        scanner_params.update(
            {
                k: v['description']
                for k, v in self._daemon.scanner_params.items()
            }
        )

        return elements

    def is_new_scan_allowed(self) -> bool:
        """ Check if max_scans has been reached.

        Return:
            True if a new scan can be launch.
        """
        if (self._daemon.max_scans == 0) or (
            len(self._daemon.scan_processes) < self._daemon.max_scans
        ):
            return True

        return False

    def is_enough_free_memory(self) -> bool:
        """ Check if there is enough free memory in the system to run
        a new scan. The necessary memory is a rough calculation and very
        conservative.

        Return:
            True if there is enough memory for a new scan.
        """

        ps_process = psutil.Process()
        proc_memory = ps_process.memory_info().rss

        free_mem = psutil.virtual_memory().free

        if free_mem > (4 * proc_memory):
            return True

        return False

    def handle_xml(self, xml: Element) -> bytes:
        """ Handles <start_scan> command.

        Return:
            Response string for <start_scan> command.
        """

        if self._daemon.check_free_memory and not self.is_enough_free_memory():
            raise OspdCommandError(
                'Not possible to run a new scan. Not enough free memory.',
                'start_scan',
            )

        if not self.is_new_scan_allowed():
            raise OspdCommandError(
                'Not possible to run a new scan. Max scan limit reached.',
                'start_scan',
            )

        target_str = xml.get('target')
        ports_str = xml.get('ports')

        # For backward compatibility, if target and ports attributes are set,
        # <targets> element is ignored.
        if target_str is None or ports_str is None:
            target_element = xml.find('targets/target')
            if target_element is None:
                raise OspdCommandError('No targets or ports', 'start_scan')
            else:
                scan_target = OspRequest.process_target_element(target_element)
        else:
            scan_target = {
                'hosts': target_str,
                'ports': ports_str,
                'credentials': {},
                'exclude_hosts': '',
                'finished_hosts': '',
                'options': {},
            }
            logger.warning(
                "Legacy start scan command format is beeing used, which "
                "is deprecated since 20.04. Please read the documentation "
                "for start scan command."
            )

        scan_id = xml.get('scan_id')
        if scan_id is not None and scan_id != '' and not valid_uuid(scan_id):
            raise OspdCommandError('Invalid scan_id UUID', 'start_scan')

        if xml.get('parallel'):
            logger.warning(
                "parallel attribute of start_scan will be ignored, sice "
                "parallel scan is not supported by OSPd."
            )

        scanner_params = xml.find('scanner_params')
        if scanner_params is None:
            raise OspdCommandError('No scanner_params element', 'start_scan')

        params = self._daemon.preprocess_scan_params(scanner_params)

        # VTS is an optional element. If present should not be empty.
        vt_selection = {}  # type: Dict
        scanner_vts = xml.find('vt_selection')
        if scanner_vts is not None:
            if len(scanner_vts) == 0:
                raise OspdCommandError('VTs list is empty', 'start_scan')
            else:
                vt_selection = OspRequest.process_vts_params(scanner_vts)

        # Dry run case.
        if 'dry_run' in params and int(params['dry_run']):
            scan_func = self._daemon.dry_run_scan
            scan_params = None
        else:
            scan_func = self._daemon.start_scan
            scan_params = self._daemon.process_scan_params(params)

        scan_id_aux = scan_id
        scan_id = self._daemon.create_scan(
            scan_id, scan_target, scan_params, vt_selection
        )

        if not scan_id:
            id_ = Element('id')
            id_.text = scan_id_aux
            return simple_response_str('start_scan', 100, 'Continue', id_)

        scan_process = create_process(
            func=scan_func, args=(scan_id, scan_target)
        )

        self._daemon.scan_processes[scan_id] = scan_process

        scan_process.start()

        id_ = Element('id')
        id_.text = scan_id

        return simple_response_str('start_scan', 200, 'OK', id_)


class GetMemoryUsage(BaseCommand):

    name = "get_memory_usage"
    description = "print the memory consumption of all processes"
    attributes = {
        'unit': 'Unit for displaying memory consumption (b = bytes, '
        'kb = kilobytes, mb = megabytes). Defaults to b.'
    }
    must_be_initialized = False

    @staticmethod
    def _get_memory(value: int, unit: str = None) -> str:
        if not unit:
            return str(value)

        unit = unit.lower()

        if unit == 'kb':
            return str(Decimal(value) / 1024)

        if unit == 'mb':
            return str(Decimal(value) / (1024 * 1024))

        return str(value)

    @staticmethod
    def _create_process_element(name: str, pid: int):
        process_element = Element('process')
        process_element.set('name', name)
        process_element.set('pid', str(pid))

        return process_element

    @classmethod
    def _add_memory_info(
        cls, process_element: Element, pid: int, unit: str = None
    ):
        try:
            ps_process = psutil.Process(pid)
        except psutil.NoSuchProcess:
            return

        memory = ps_process.memory_info()

        rss_element = Element('rss')
        rss_element.text = cls._get_memory(memory.rss, unit)

        process_element.append(rss_element)

        vms_element = Element('vms')
        vms_element.text = cls._get_memory(memory.vms, unit)

        process_element.append(vms_element)

        shared_element = Element('shared')
        shared_element.text = cls._get_memory(memory.shared, unit)

        process_element.append(shared_element)

    def handle_xml(self, xml: Element) -> bytes:
        processes_element = Element('processes')
        unit = xml.get('unit')

        current_process = multiprocessing.current_process()
        process_element = self._create_process_element(
            current_process.name, current_process.pid
        )

        self._add_memory_info(process_element, current_process.pid, unit)

        processes_element.append(process_element)

        for proc in multiprocessing.active_children():
            process_element = self._create_process_element(proc.name, proc.pid)

            self._add_memory_info(process_element, proc.pid, unit)

            processes_element.append(process_element)

        return simple_response_str('get_memory', 200, 'OK', processes_element)


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 multiprocessing
19			import re
20			import logging
21			import subprocess
22
23			from decimal import Decimal
24			from typing import Optional, Dict, Any, Union, Iterator
25
26			from xml.etree.ElementTree import Element, SubElement
27
28			import psutil
29
30			from ospd.errors import OspdCommandError
31			from ospd.misc import valid_uuid, create_process
32			from ospd.protocol import OspRequest, OspResponse
33			from ospd.xml import (
34			simple_response_str,
35			get_elements_from_dict,
36			XmlStringHelper,
37			)
38
39			from .initsubclass import InitSubclassMeta
40			from .registry import register_command
41
42			logger = logging.getLogger(__name__)
43
44
45			class BaseCommand(metaclass=InitSubclassMeta):
46
47			name = None
48			description = None
49			attributes = None
50			elements = None
51			must_be_initialized = None
52
53			def __init_subclass__(cls, **kwargs):
54			super_cls = super()
55
56			if hasattr(super_cls, '__init_subclass__'):
57			super_cls.__init_subclass__(**kwargs)
58
59			register_command(cls)
60
61			def __init__(self, daemon):
62			self._daemon = daemon
63
64			def get_name(self) -> str:
65			return self.name
66
67			def get_description(self) -> str:
68			return self.description
69
70			def get_attributes(self) -> Optional[Dict[str, Any]]:
71			return self.attributes
72
73			def get_elements(self) -> Optional[Dict[str, Any]]:
74			return self.elements
75
76			def handle_xml(
77			self, xml: Element, initialized: bool
78			) -> Union[bytes, Iterator[bytes]]:
79			raise NotImplementedError()
80
81			def as_dict(self):
82			return {
83			'name': self.get_name(),
84			'attributes': self.get_attributes(),
85			'description': self.get_description(),
86			'elements': self.get_elements(),
87			}
88
89			def __repr__(self):
90			return '<{} description="{}" attributes={} elements={}>'.format(
91			self.name, self.description, self.attributes, self.elements
92			)
93
94
95			class HelpCommand(BaseCommand):
96			name = "help"
97			description = 'Print the commands help.'
98			attributes = {'format': 'Help format. Could be text or xml.'}
99			must_be_initialized = False
100
101			def handle_xml(self, xml: Element) -> bytes:
102			help_format = xml.get('format')
103
104			if help_format is None or help_format == "text":
105			# Default help format is text.
106			return simple_response_str(
107			'help', 200, 'OK', self._daemon.get_help_text()
108			)
109			elif help_format == "xml":
110			text = get_elements_from_dict(
111			{k: v.as_dict() for k, v in self._daemon.commands.items()}
112			)
113			return simple_response_str('help', 200, 'OK', text)
114
115			raise OspdCommandError('Bogus help format', 'help')
116
117
118			class GetVersion(BaseCommand):
119			name = "get_version"
120			description = 'Return various version information'
121			must_be_initialized = False
122
123			def handle_xml(self, xml: Element) -> bytes:
124			""" Handles <get_version> command.
125
126			Return:
127			Response string for <get_version> command.
128			"""
129			protocol = Element('protocol')
130
131			for name, value in [
132			('name', 'OSP'),
133			('version', self._daemon.get_protocol_version()),
134			]:
135			elem = SubElement(protocol, name)
136			elem.text = value
137
138			daemon = Element('daemon')
139			for name, value in [
140			('name', self._daemon.get_daemon_name()),
141			('version', self._daemon.get_daemon_version()),
142			]:
143			elem = SubElement(daemon, name)
144			elem.text = value
145
146			scanner = Element('scanner')
147			for name, value in [
148			('name', self._daemon.get_scanner_name()),
149			('version', self._daemon.get_scanner_version()),
150			]:
151			elem = SubElement(scanner, name)
152			elem.text = value
153
154			content = [protocol, daemon, scanner]
155
156			vts_version = self._daemon.get_vts_version()
157			if vts_version:
158			vts = Element('vts')
159			elem = SubElement(vts, 'version')
160			elem.text = vts_version
161			content.append(vts)
162
163			return simple_response_str('get_version', 200, 'OK', content)
164
165
166			GVMCG_TITLES = [
167			'cpu-*',
168			'proc',
169			'mem',
170			'swap',
171			'load',
172			'df-*',
173			'disk-sd[a-z][0-9]-rw',
174			'disk-sd[a-z][0-9]-load',
175			'disk-sd[a-z][0-9]-io-load',
176			'interface-eth*-traffic',
177			'interface-eth*-err-rate',
178			'interface-eth*-err',
179			'sensors-_temperature-',
180			'sensors-_fanspeed-',
181			'sensors-_voltage-',
182			'titles',
183			] # type: List[str]
184
185
186			class GetPerformance(BaseCommand):
187			name = "get_performance"
188			description = 'Return system report'
189			attributes = {
190			'start': 'Time of first data point in report.',
191			'end': 'Time of last data point in report.',
192			'title': 'Name of report.',
193			}
194			must_be_initialized = False
195
196			def handle_xml(self, xml: Element) -> bytes:
197			""" Handles <get_performance> command.
198
199			@return: Response string for <get_performance> command.
200			"""
201			start = xml.attrib.get('start')
202			end = xml.attrib.get('end')
203			titles = xml.attrib.get('titles')
204
205			cmd = ['gvmcg']
206			if start:
207			try:
208			int(start)
209			except ValueError:
210			raise OspdCommandError(
211			'Start argument must be integer.', 'get_performance'
212			)
213
214			cmd.append(start)
215
216			if end:
217			try:
218			int(end)
219			except ValueError:
220			raise OspdCommandError(
221			'End argument must be integer.', 'get_performance'
222			)
223
224			cmd.append(end)
225
226			if titles:
227			combined = "(" + ")\|(".join(GVMCG_TITLES) + ")"
228			forbidden = "^[^\|&;]+$"
229			if re.match(combined, titles) and re.match(forbidden, titles):
230			cmd.append(titles)
231			else:
232			raise OspdCommandError(
233			'Arguments not allowed', 'get_performance'
234			)
235
236			try:
237			output = subprocess.check_output(cmd)
238			except (subprocess.CalledProcessError, OSError) as e:
239			raise OspdCommandError(
240			'Bogus get_performance format. %s' % e, 'get_performance'
241			)
242
243			return simple_response_str(
244			'get_performance', 200, 'OK', output.decode()
245			)
246
247
248			class GetScannerDetails(BaseCommand):
249			name = 'get_scanner_details'
250			description = 'Return scanner description and parameters'
251			must_be_initialized = True
252
253			def handle_xml(self, xml: Element) -> bytes:
254			""" Handles <get_scanner_details> command.
255
256			@return: Response string for <get_scanner_details> command.
257			"""
258			desc_xml = Element('description')
259			desc_xml.text = self._daemon.get_scanner_description()
260			scanner_params = self._daemon.get_scanner_params()
261			details = [
262			desc_xml,
263			OspResponse.create_scanner_params_xml(scanner_params),
264			]
265			return simple_response_str('get_scanner_details', 200, 'OK', details)
266
267
268			class DeleteScan(BaseCommand):
269			name = 'delete_scan'
270			description = 'Delete a finished scan.'
271			attributes = {'scan_id': 'ID of scan to delete.'}
272			must_be_initialized = False
273
274			def handle_xml(self, xml: Element) -> bytes:
275			""" Handles <delete_scan> command.
276
277			@return: Response string for <delete_scan> command.
278			"""
279			scan_id = xml.get('scan_id')
280			if scan_id is None:
281			return simple_response_str(
282			'delete_scan', 404, 'No scan_id attribute'
283			)
284
285			if not self._daemon.scan_exists(scan_id):
286			text = "Failed to find scan '{0}'".format(scan_id)
287			return simple_response_str('delete_scan', 404, text)
288
289			self._daemon.check_scan_process(scan_id)
290
291			if self._daemon.delete_scan(scan_id):
292			return simple_response_str('delete_scan', 200, 'OK')
293
294			raise OspdCommandError('Scan in progress', 'delete_scan')
295
296
297			class GetVts(BaseCommand):
298			name = 'get_vts'
299			description = 'List of available vulnerability tests.'
300			attributes = {
301			'vt_id': 'ID of a specific vulnerability test to get.',
302			'filter': 'Optional filter to get an specific vt collection.',
303			}
304			must_be_initialized = True
305
306			def handle_xml(self, xml: Element) -> Iterator[bytes]:
307			""" Handles <get_vts> command.
308			Writes the vt collection on the stream.
309			The <get_vts> element accept two optional arguments.
310			vt_id argument receives a single vt id.
311			filter argument receives a filter selecting a sub set of vts.
312			If both arguments are given, the vts which match with the filter
313			are return.
314
315			@return: Response string for <get_vts> command on fail.
316			"""
317			xml_helper = XmlStringHelper()
318
319			vt_id = xml.get('vt_id')
320			vt_filter = xml.get('filter')
321			_details = xml.get('details')
322
323			vt_details = False if _details == '0' else True
324
325			if vt_id and vt_id not in self._daemon.vts:
326			text = "Failed to find vulnerability test '{0}'".format(vt_id)
327			raise OspdCommandError(text, 'get_vts', 404)
328
329			filtered_vts = None
330			if vt_filter:
331			filtered_vts = self._daemon.vts_filter.get_filtered_vts_list(
332			self._daemon.vts, vt_filter
333			)
334			vts_selection = self._daemon.get_vts_selection_list(vt_id, filtered_vts)
335			# List of xml pieces with the generator to be iterated
336			yield xml_helper.create_response('get_vts')
337
338			begin_vts_tag = xml_helper.create_element('vts')
339			val = len(self._daemon.vts)
340			begin_vts_tag = xml_helper.add_attr(begin_vts_tag, "total", val)
341			if filtered_vts:
342			val = len(filtered_vts)
343			begin_vts_tag = xml_helper.add_attr(begin_vts_tag, "sent", val)
344
345			if self._daemon.vts.sha256_hash is not None:
346			begin_vts_tag = xml_helper.add_attr(
347			begin_vts_tag, "sha256_hash", self._daemon.vts.sha256_hash
348			)
349
350			yield begin_vts_tag
351
352			for vt in self._daemon.get_vt_iterator(vts_selection, vt_details):
353			yield xml_helper.add_element(self._daemon.get_vt_xml(vt))
354
355			yield xml_helper.create_element('vts', end=True)
356			yield xml_helper.create_response('get_vts', end=True)
357
358
359			class StopScan(BaseCommand):
360			name = 'stop_scan'
361			description = 'Stop a currently running scan.'
362			attributes = {'scan_id': 'ID of scan to stop.'}
363			must_be_initialized = True
364
365			def handle_xml(self, xml: Element) -> bytes:
366			""" Handles <stop_scan> command.
367
368			@return: Response string for <stop_scan> command.
369			"""
370
371			scan_id = xml.get('scan_id')
372			if scan_id is None or scan_id == '':
373			raise OspdCommandError('No scan_id attribute', 'stop_scan')
374
375			self._daemon.stop_scan(scan_id)
376
377			# Don't send response until the scan is stopped.
378			try:
379			self._daemon.scan_processes[scan_id].join()
380			except KeyError:
381			pass
382
383			return simple_response_str('stop_scan', 200, 'OK')
384
385
386			class GetScans(BaseCommand):
387			name = 'get_scans'
388			description = 'List the scans in buffer.'
389			attributes = {
390			'scan_id': 'ID of a specific scan to get.',
391			'details': 'Whether to return the full scan report.',
392			'pop_results': 'Whether to remove the fetched results.',
393			'max_results': 'Maximum number of results to fetch.',
394			}
395			must_be_initialized = False
396
397			def handle_xml(self, xml: Element) -> bytes:
398			""" Handles <get_scans> command.
399
400			@return: Response string for <get_scans> command.
401			"""
402
403			scan_id = xml.get('scan_id')
404			details = xml.get('details')
405			pop_res = xml.get('pop_results')
406			max_res = xml.get('max_results')
407
408			if details and details == '0':
409			details = False
410			else:
411			details = True
412			if pop_res and pop_res == '1':
413			pop_res = True
414			else:
415			pop_res = False
416			if max_res:
417			max_res = int(max_res)
418
419			responses = []
420			if scan_id and scan_id in self._daemon.scan_collection.ids_iterator():
421			self._daemon.check_scan_process(scan_id)
422			scan = self._daemon.get_scan_xml(scan_id, details, pop_res, max_res)
423			responses.append(scan)
424			elif scan_id:
425			text = "Failed to find scan '{0}'".format(scan_id)
426			return simple_response_str('get_scans', 404, text)
427			else:
428			for scan_id in self._daemon.scan_collection.ids_iterator():
429			self._daemon.check_scan_process(scan_id)
430			scan = self._daemon.get_scan_xml(
431			scan_id, details, pop_res, max_res
432			)
433			responses.append(scan)
434
435			return simple_response_str('get_scans', 200, 'OK', responses)
436
437
438			class StartScan(BaseCommand):
439			name = 'start_scan'
440			description = 'Start a new scan.'
441			attributes = {
442			'target': 'Target host to scan',
443			'ports': 'Ports list to scan',
444			'scan_id': 'Optional UUID value to use as scan ID',
445			'parallel': 'Optional nummer of parallel target to scan',
446			}
447			must_be_initialized = True
448
449			def get_elements(self):
450			elements = {}
451
452			if self.elements:
453			elements.update(self.elements)
454
455			scanner_params = elements.get('scanner_params', {}).copy()
456			elements['scanner_params'] = scanner_params
457
458			scanner_params.update(
459			{
460			k: v['description']
461			for k, v in self._daemon.scanner_params.items()
462			}
463			)
464
465			return elements
466
467			def is_new_scan_allowed(self) -> bool:
468			""" Check if max_scans has been reached.
469
470			Return:
471			True if a new scan can be launch.
472			"""
473			if (self._daemon.max_scans == 0) or (
474			len(self._daemon.scan_processes) < self._daemon.max_scans
475			):
476			return True
477
478			return False
479
480			def is_enough_free_memory(self) -> bool:
481			""" Check if there is enough free memory in the system to run
482			a new scan. The necessary memory is a rough calculation and very
483			conservative.
484
485			Return:
486			True if there is enough memory for a new scan.
487			"""
488
489			ps_process = psutil.Process()
490			proc_memory = ps_process.memory_info().rss
491
492			free_mem = psutil.virtual_memory().free
493
494			if free_mem > (4 * proc_memory):
495			return True
496
497			return False
498
499			def handle_xml(self, xml: Element) -> bytes:
500			""" Handles <start_scan> command.
501
502			Return:
503			Response string for <start_scan> command.
504			"""
505
506			if self._daemon.check_free_memory and not self.is_enough_free_memory():
507			raise OspdCommandError(
508			'Not possible to run a new scan. Not enough free memory.',
509			'start_scan',
510			)
511
512			if not self.is_new_scan_allowed():
513			raise OspdCommandError(
514			'Not possible to run a new scan. Max scan limit reached.',
515			'start_scan',
516			)
517
518			target_str = xml.get('target')
519			ports_str = xml.get('ports')
520
521			# For backward compatibility, if target and ports attributes are set,
522			# <targets> element is ignored.
523			if target_str is None or ports_str is None:
524			target_element = xml.find('targets/target')
525			if target_element is None:
526			raise OspdCommandError('No targets or ports', 'start_scan')
527			else:
528			scan_target = OspRequest.process_target_element(target_element)
529			else:
530			scan_target = {
531			'hosts': target_str,
532			'ports': ports_str,
533			'credentials': {},
534			'exclude_hosts': '',
535			'finished_hosts': '',
536			'options': {},
537			}
538			logger.warning(
539			"Legacy start scan command format is beeing used, which "
540			"is deprecated since 20.04. Please read the documentation "
541			"for start scan command."
542			)
543
544			scan_id = xml.get('scan_id')
545			if scan_id is not None and scan_id != '' and not valid_uuid(scan_id):
546			raise OspdCommandError('Invalid scan_id UUID', 'start_scan')
547
548			if xml.get('parallel'):
549			logger.warning(
550			"parallel attribute of start_scan will be ignored, sice "
551			"parallel scan is not supported by OSPd."
552			)
553
554			scanner_params = xml.find('scanner_params')
555			if scanner_params is None:
556			raise OspdCommandError('No scanner_params element', 'start_scan')
557
558			params = self._daemon.preprocess_scan_params(scanner_params)
559
560			# VTS is an optional element. If present should not be empty.
561			vt_selection = {} # type: Dict
562			scanner_vts = xml.find('vt_selection')
563			if scanner_vts is not None:
564			if len(scanner_vts) == 0:
565			raise OspdCommandError('VTs list is empty', 'start_scan')
566			else:
567			vt_selection = OspRequest.process_vts_params(scanner_vts)
568
569			# Dry run case.
570			if 'dry_run' in params and int(params['dry_run']):
571			scan_func = self._daemon.dry_run_scan
572			scan_params = None
573			else:
574			scan_func = self._daemon.start_scan
575			scan_params = self._daemon.process_scan_params(params)
576
577			scan_id_aux = scan_id
578			scan_id = self._daemon.create_scan(
579			scan_id, scan_target, scan_params, vt_selection
580			)
581
582			if not scan_id:
583			id_ = Element('id')
584			id_.text = scan_id_aux
585			return simple_response_str('start_scan', 100, 'Continue', id_)
586
587			scan_process = create_process(
588			func=scan_func, args=(scan_id, scan_target)
589			)
590
591			self._daemon.scan_processes[scan_id] = scan_process
592
593			scan_process.start()
594
595			id_ = Element('id')
596			id_.text = scan_id
597
598			return simple_response_str('start_scan', 200, 'OK', id_)
599
600
601			class GetMemoryUsage(BaseCommand):
602
603			name = "get_memory_usage"
604			description = "print the memory consumption of all processes"
605			attributes = {
606			'unit': 'Unit for displaying memory consumption (b = bytes, '
607			'kb = kilobytes, mb = megabytes). Defaults to b.'
608			}
609			must_be_initialized = False
610
611			@staticmethod
612			def _get_memory(value: int, unit: str = None) -> str:
613			if not unit:
614			return str(value)
615
616			unit = unit.lower()
617
618			if unit == 'kb':
619			return str(Decimal(value) / 1024)
620
621			if unit == 'mb':
622			return str(Decimal(value) / (1024 * 1024))
623
624			return str(value)
625
626			@staticmethod
627			def _create_process_element(name: str, pid: int):
628			process_element = Element('process')
629			process_element.set('name', name)
630			process_element.set('pid', str(pid))
631
632			return process_element
633
634			@classmethod
635			def _add_memory_info(
636			cls, process_element: Element, pid: int, unit: str = None
637			):
638			try:
639			ps_process = psutil.Process(pid)
640			except psutil.NoSuchProcess:
641			return
642
643			memory = ps_process.memory_info()
644
645			rss_element = Element('rss')
646			rss_element.text = cls._get_memory(memory.rss, unit)
647
648			process_element.append(rss_element)
649
650			vms_element = Element('vms')
651			vms_element.text = cls._get_memory(memory.vms, unit)
652
653			process_element.append(vms_element)
654
655			shared_element = Element('shared')
656			shared_element.text = cls._get_memory(memory.shared, unit)
657
658			process_element.append(shared_element)
659
660			def handle_xml(self, xml: Element) -> bytes:
661			processes_element = Element('processes')
662			unit = xml.get('unit')
663
664			current_process = multiprocessing.current_process()
665			process_element = self._create_process_element(
666			current_process.name, current_process.pid
667			)
668
669			self._add_memory_info(process_element, current_process.pid, unit)
670
671			processes_element.append(process_element)
672
673			for proc in multiprocessing.active_children():
674			process_element = self._create_process_element(proc.name, proc.pid)
675
676			self._add_memory_info(process_element, proc.pid, unit)
677
678			processes_element.append(process_element)
679
680			return simple_response_str('get_memory', 200, 'OK', processes_element)
681

greenbone / ospd

Push — master ( f571c7...585a62 )

StartScan.is_new_scan_allowed() A

Complexity

Size

Duplication

Importance

Duplication Side-by-Side

Filter issues like