ospd.command.command.GetMemoryUsage.handle_xml() - Code Metrics - Inspection of "Improve handling of dead hosts and improve scan pr..." - greenbone/ospd - Measure and Improve Code Quality continuously with Scrutinizer

Passed

Pull Request — master (#266)

by Juan José

created 2020-05-11 12:31 UTC

ospd.command.command.GetMemoryUsage.handle_xml() A

↳ Parent: ospd.command.command

Complexity

Conditions

Size

Total Lines	21
Code Lines	13

Duplication

Lines	0
Ratio	0 %

Importance

Changes

Metric	Value
cc	2
eloc	13
nop	2
dl	0
loc	21
rs	9.75
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) -> 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.
        """
        self._daemon.vts.is_cache_available = False

        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 self._daemon.vts and vt_id and vt_id not in self._daemon.vts:
            self._daemon.vts.is_cache_available = True
            text = "Failed to find vulnerability test '{0}'".format(vt_id)
            raise OspdCommandError(text, 'get_vts', 404)

        filtered_vts = None
        if vt_filter:
            try:
                filtered_vts = self._daemon.vts_filter.get_filtered_vts_list(
                    self._daemon.vts, vt_filter
                )
            except OspdCommandError as filter_error:
                self._daemon.vts.is_cache_available = True
                raise OspdCommandError(filter_error)

        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)

        self._daemon.vts.is_cache_available = 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')
        progress = xml.get('progress')

        if details and details == '0':
            details = False
        else:
            details = True
            pop_res = pop_res and pop_res == '1'

            if max_res:
                max_res = int(max_res)

        progress = progress and progress == '1'

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

greenbone / ospd

Pull Request — master (#266)

ospd.command.command.GetMemoryUsage.handle_xml() A

Complexity

Size

Duplication

Importance

Duplication Side-by-Side

Filter issues like