ospd.command.command.GetMemoryUsage.handle_xml() - Code Metrics - Inspection of "Get memory usage" - greenbone/ospd - Measure and Improve Code Quality continuously with Scrutinizer

Passed

Pull Request — master (#207)

unknown

created 2020-02-21 13:23 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) 2020 Greenbone Networks GmbH
#
# SPDX-License-Identifier: GPL-2.0-or-later
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# 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 General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.

import multiprocessing
import re
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.network import target_str_to_list
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


class BaseCommand(metaclass=InitSubclassMeta):

    name = None
    description = None
    attributes = None
    elements = 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.'}

    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'

    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.',
    }

    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'

    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.'}

    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.',
    }

    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')

        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
            )

        # 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)

        yield begin_vts_tag

        for vt in self._daemon.get_vts_selection_list(vt_id, filtered_vts):
            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.'}

    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()
            exitcode = self._daemon.scan_processes[  # pylint: disable=unused-variable
                scan_id
            ].exitcode
        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.',
    }

    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',
    }

    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 handle_xml(self, xml: Element) -> bytes:
        """ Handles <start_scan> command.

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

        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_list = xml.find('targets')
            if target_list is None or len(target_list) == 0:
                raise OspdCommandError('No targets or ports', 'start_scan')
            else:
                scan_targets = OspRequest.process_targets_element(target_list)
        else:
            scan_targets = []
            for single_target in target_str_to_list(target_str):
                scan_targets.append([single_target, ports_str, '', '', '', ''])

        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')

        try:
            parallel = int(xml.get('parallel', '1'))
        except ValueError:
            raise OspdCommandError(
                'Invalid value for parallel scans. It must be a number',
                'start_scan',
            )

        if parallel < 1 or parallel > 20:
            parallel = 1

        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_targets, 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_targets, parallel)
        )

        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.'
    }

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

greenbone / ospd

Pull Request — master (#207)

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

Complexity

Size

Duplication

Importance

Duplication Side-by-Side

Filter issues like