Improve generator

ospd/command/command.py

@@ -308,33 +308,23 @@ def handle_xml(self, xml: Element) -> str:
 
         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
             )
 
-        # Generator
-        vts_list = (vt for vt in self._daemon.get_vts_xml(vt_id, filtered_vts))
-
         # List of xml pieces with the generator to be iterated
-        response = [
-            xml_helper.create_response('get_vts'),
-            xml_helper.create_element('vts'),
-            vts_list,
-            xml_helper.create_element('vts', end=True),
-            xml_helper.create_response('get_vts', end=True),
-        ]
+        yield xml_helper.create_response('get_vts')
+        yield xml_helper.create_element('vts')
 
-        for elem in response:
-            if isinstance(elem, GeneratorType):
-                for vts_chunk in elem:
-                    yield xml_helper.add_element(
-                        self._daemon.get_vt_xml(vts_chunk)
-                    )
-            else:
-                yield elem
+        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):

ospd/ospd.py

@@ -1142,9 +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):
-        """ Python Generator for VTS.
-        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
@@ -1157,8 +1159,7 @@ def get_vts_xml(self, vt_id: str = None, filtered_vts: Dict = None):
             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 = []
         if not self.vts: