Merge pull request #201 from jjnicola/yield-vts

Improve get_vts cmd response, sending the vts piece by piece.

CHANGELOG.md

@@ -16,6 +16,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 - Modify __init__() method and use new syntax for super(). [#186](https://github.com/greenbone/ospd/pull/186)
 - Create data manager and spawn new process to keep the vts dictionary. [#191](https://github.com/greenbone/ospd/pull/191)
 - Update daemon start sequence. Run daemon.check before daemon.init now. [#197](https://github.com/greenbone/ospd/pull/197)
+- Improve get_vts cmd response, sending the vts piece by piece.[#201](https://github.com/greenbone/ospd/pull/201)
 
 ## [2.0.1] (unreleased)
 

ospd/command/__init__.py

@@ -16,4 +16,5 @@
 # along with this program; if not, write to the Free Software
 # Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
 
+import ospd.command.command
 from .registry import get_commands

ospd/command/command.py

@@ -19,6 +19,7 @@
 import re
 import subprocess
 
+from types import GeneratorType
 from typing import Optional, Dict, Any
 
 from xml.etree.ElementTree import Element, SubElement
@@ -27,7 +28,11 @@
 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
+from ospd.xml import (
+    simple_response_str,
+    get_elements_from_dict,
+    XmlStringHelper,
+)
 
 from .initsubclass import InitSubclassMeta
 from .registry import register_command
@@ -287,35 +292,39 @@ class GetVts(BaseCommand):
 
     def handle_xml(self, xml: Element) -> str:
         """ 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.
+        @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)
-            return simple_response_str('get_vts', 404, text)
+            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
             )
 
-        responses = []
-
-        vts_xml = self._daemon.get_vts_xml(vt_id, filtered_vts)
+        # List of xml pieces with the generator to be iterated
+        yield xml_helper.create_response('get_vts')
+        yield xml_helper.create_element('vts')
 
-        responses.append(vts_xml)
+        for vt in self._daemon.get_vts_selection_list(vt_id, filtered_vts):
+            yield xml_helper.add_element(self._daemon.get_vt_xml(vt))
 
-        return simple_response_str('get_vts', 200, 'OK', responses)
+        yield xml_helper.create_element('vts', end=True)
+        yield xml_helper.create_response('get_vts', end=True)
 
 
 class StopScan(BaseCommand):

ospd/ospd.py

@@ -495,8 +495,9 @@ def handle_client_stream(self, stream) -> None:
             logger.debug("Empty client stream")
             return
 
+        response = None
         try:
-            response = self.handle_command(data)
+            self.handle_command(data, stream)
         except OspdCommandError as exception:
             response = exception.as_xml()
             logger.debug('Command error: %s', exception.message)
@@ -505,7 +506,8 @@ def handle_client_stream(self, stream) -> None:
             exception = OspdCommandError('Fatal error', 'error')
             response = exception.as_xml()
 
-        stream.write(response)
+        if response:
+            stream.write(response)
         stream.close()
 
     def parallel_scan(self, scan_id: str, target: str) -> None:
@@ -1140,8 +1142,11 @@ def get_vt_xml(self, vt_id: str):
 
         return vt_xml
 
-    def get_vts_xml(self, vt_id: str = None, filtered_vts: Dict = None):
-        """ Gets collection of vulnerability test information in XML format.
+    def get_vts_selection_list(
+        self, vt_id: str = None, filtered_vts: Dict = None
+    ) -> List:
+        """
+        Get list of VT's OID.
         If vt_id is specified, the collection will contain only this vt, if
         found.
         If no vt_id is specified or filtered_vts is None (default), the
@@ -1151,36 +1156,32 @@ def get_vts_xml(self, vt_id: str = None, filtered_vts: Dict = None):
 
         Arguments:
             vt_id (vt_id, optional): ID of the vt to get.
-            filtered_vts (dict, optional): Filtered VTs collection.
+            filtered_vts (list, optional): Filtered VTs collection.
 
         Return:
-            String of collection of vulnerability test information in
-            XML format.
+            List of selected VT's OID.
         """
-
-        vts_xml = Element('vts')
-
+        vts_xml = []
         if not self.vts:
             return vts_xml
 
+        # No match for the filter
         if filtered_vts is not None and len(filtered_vts) == 0:
             return vts_xml
 
         if filtered_vts:
-            for vt_id in filtered_vts:
-                vts_xml.append(self.get_vt_xml(vt_id))
+            vts_list = filtered_vts
         elif vt_id:
-            vts_xml.append(self.get_vt_xml(vt_id))
+            vts_list = [vt_id]
         else:
-            # Because DictProxy for python3.5 doesn't support
+            # TODO: Because DictProxy for python3.5 doesn't support
             # iterkeys(), itervalues(), or iteritems() either, the iteration
             # must be done as follow.
-            for vt_id in iter(self.vts.keys()):
-                vts_xml.append(self.get_vt_xml(vt_id))
+            vts_list = iter(self.vts.keys())
 
-        return vts_xml
+        return vts_list
 
-    def handle_command(self, command: str) -> str:
+    def handle_command(self, command: str, stream) -> str:
         """ Handles an osp command in a string.
 
         @return: OSP Response to command.
@@ -1195,7 +1196,13 @@ def handle_command(self, command: str) -> str:
         if not command and tree.tag != "authenticate":
             raise OspdCommandError('Bogus command name')
 
-        return command.handle_xml(tree)
+        response = command.handle_xml(tree)
+
+        if isinstance(response, bytes):
+            stream.write(response)
+        else:
+            for data in response:
+                stream.write(data)
 
     def check(self):
         """ Asserts to False. Should be implemented by subclass. """

ospd/vtfilter.py

@@ -98,7 +98,9 @@ def format_filter_value(self, element: str, value: Dict):
         format_func = self.allowed_filter.get(element)
         return format_func(value)
 
-    def get_filtered_vts_list(self, vts: Dict, vt_filter: str) -> Optional[Dict]:
+    def get_filtered_vts_list(
+        self, vts: Dict, vt_filter: str
+    ) -> Optional[Dict]:
         """ Gets a collection of vulnerability test from the vts dictionary,
         which match the filter.
 
@@ -107,7 +109,8 @@ def get_filtered_vts_list(self, vts: Dict, vt_filter: str) -> Optional[Dict]:
             vts (dictionary): The complete vts collection.
 
         Returns:
-            Dictionary with filtered vulnerability tests.
+            List with filtered vulnerability tests. The list can be empty.
+            None in case of filter parse failure.
         """
         if not vt_filter:
             raise OspdCommandError('vt_filter: A valid filter is required.')
@@ -116,17 +119,23 @@ def get_filtered_vts_list(self, vts: Dict, vt_filter: str) -> Optional[Dict]:
         if not filters:
             return None
 
-        _vts_aux = vts.copy()
+        vt_oid_list = list(vts.keys())
+
         for _element, _oper, _filter_val in filters:
-            for vt_id in _vts_aux.copy():
-                if not _vts_aux[vt_id].get(_element):
-                    _vts_aux.pop(vt_id)
+            # Use iter because python3.5 has no support for
+            # iteration over DictProxy.
+            vts_generator = (vt for vt in iter(vts.keys()))
+            for vt_oid in vts_generator:
+                if vt_oid not in vt_oid_list:
+                    continue
+                if not vts[vt_oid].get(_element):
+                    vt_oid_list.remove(vt_oid)
                     continue
-                _elem_val = _vts_aux[vt_id].get(_element)
+                _elem_val = vts[vt_oid].get(_element)
                 _val = self.format_filter_value(_element, _elem_val)
                 if self.filter_operator[_oper](_val, _filter_val):
                     continue
                 else:
-                    _vts_aux.pop(vt_id)
+                    vt_oid_list.remove(vt_oid)
 
-        return _vts_aux
+        return vt_oid_list

ospd/xml.py

@@ -133,3 +133,82 @@ def elements_as_text(
         text = ''.join([text, ele_txt])
 
     return text
+
+
+class XmlStringHelper:
+    """ Class with methods to help the creation of a xml object in
+    string format.
+    """
+
+    def create_element(self, elem_name: str, end: bool = False) -> bytes:
+        """ Get a name and create the open element of an entity.
+
+        Arguments:
+            elem_name (str): The name of the tag element.
+            end (bool): Create a initial tag if False, otherwise the end tag.
+
+        Return:
+            Encoded string representing a part of an xml element.
+        """
+        if end:
+            ret = "</%s>" % elem_name
+        else:
+            ret = "<%s>" % elem_name
+
+        return ret.encode()
+
+    def create_response(self, command: str, end: bool = False) -> bytes:
+        """ Create or end an xml response.
+
+        Arguments:
+            command (str): The name of the command for the response element.
+            end (bool): Create a initial tag if False, otherwise the end tag.
+
+        Return:
+            Encoded string representing a part of an xml element.
+        """
+        if not command:
+            return
+
+        if end:
+            return ('</%s_response>' % command).encode()
+
+        return (
+            '<%s_response status="200" status_text="OK">' % command
+        ).encode()
+
+    def add_element(
+        self,
+        content: Union[Element, str, list],
+        xml_str: bytes = None,
+        end: bool = False,
+    ) -> bytes:
+        """Create the initial or ending tag for a subelement, or add
+        one or many xml elements
+
+        Arguments:
+            content (Element, str, list): Content to add.
+            xml_str (bytes): Initial string where content to be added to.
+            end (bool): Create a initial tag if False, otherwise the end tag.
+                        It will be added to the xml_str.
+
+        Return:
+            Encoded string representing a part of an xml element.
+        """
+
+        if not xml_str:
+            xml_str = b''
+
+        if content:
+            if isinstance(content, list):
+                for elem in content:
+                    xml_str = xml_str + tostring(elem)
+            elif isinstance(content, Element):
+                xml_str = xml_str + tostring(content)
+            else:
+                if end:
+                    xml_str = xml_str + self.create_element(content, False)
+                else:
+                    xml_str = xml_str + self.create_element(content)
+
+        return xml_str

tests/command/test_commands.py

@@ -24,7 +24,7 @@
 from ospd.command.command import GetPerformance, StartScan, StopScan
 from ospd.errors import OspdCommandError, OspdError
 
-from ..helper import DummyWrapper, assert_called
+from ..helper import DummyWrapper, assert_called, FakeStream
 
 
 class GetPerformanceTestCase(TestCase):
@@ -261,13 +261,15 @@ def test_stop_scan(self, mock_create_process, mock_os):
         mock_process.is_alive.return_value = True
         mock_process.pid = "foo"
 
+        fs = FakeStream()
         daemon = DummyWrapper([])
         request = (
             '<start_scan target="localhost" ports="80, 443">'
             '<scanner_params />'
             '</start_scan>'
         )
-        response = et.fromstring(daemon.handle_command(request))
+        daemon.handle_command(request, fs)
+        response = fs.get_response()
 
         assert_called(mock_create_process)
         assert_called(mock_process.start)

tests/helper.py

@@ -20,6 +20,8 @@
 
 from unittest.mock import Mock
 
+from xml.etree import ElementTree as et
+
 from ospd.ospd import OSPDaemon
 
 
@@ -36,6 +38,17 @@ def assert_called(mock: Mock):
         raise AssertionError(msg)
 
 
+class FakeStream:
+    def __init__(self):
+        self.response = b''
+
+    def write(self, data):
+        self.response = self.response + data
+
+    def get_response(self):
+        return et.fromstring(self.response)
+
+
 class DummyWrapper(OSPDaemon):
     def __init__(self, results, checkresult=True):
         super().__init__()