From 670ab0a102c1f580ce13fc8eb1404f20113e18fe Mon Sep 17 00:00:00 2001
From: Juan Jose Nicola <juan.nicola@greenbone.net>
Date: Tue, 11 Feb 2020 09:17:18 +0100
Subject: [PATCH] Improve generator

---
 ospd/command/command.py | 26 ++++++++------------------
 ospd/ospd.py            | 11 ++++++-----
 2 files changed, 14 insertions(+), 23 deletions(-)

diff --git a/ospd/command/command.py b/ospd/command/command.py
index b0016346..fe926990 100644
--- a/ospd/command/command.py
+++ b/ospd/command/command.py
@@ -308,6 +308,7 @@ 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:
@@ -315,26 +316,15 @@ def handle_xml(self, xml: Element) -> str:
                 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):
diff --git a/ospd/ospd.py b/ospd/ospd.py
index 3ca40e07..1c1db5d2 100644
--- a/ospd/ospd.py
+++ b/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: