Add methods to set and get the vts feed version.

ospd/ospd.py

@@ -275,6 +275,7 @@ def __init__(self, certfile, keyfile, cafile):
             self.add_scanner_param(name, param)
         self.vts = dict()
         self.vt_id_pattern = re.compile("[0-9a-zA-Z_\-:.]{1,80}")
+        self.vts_version = None
 
     def set_command_attributes(self, name, attributes):
         """ Sets the xml attributes of a specified command. """
@@ -354,6 +355,23 @@ def add_vt(self, vt_id, name=None, vt_params=None, vt_refs=None,
 
         return len(self.vts)
 
+    def set_vts_version(self, vts_version):
+        """ Add into the vts dictionary an entry to identify the
+        vts version.
+
+        Parameters:
+            vts_version (str): Identifies a unique vts version.
+        """
+        if not vts_version:
+            raise OSPDError('A vts_version parameter is required',
+                            'set_vts_version')
+        self.vts_version = vts_version
+
+    def get_vts_version(self):
+        """Return the vts version.
+        """
+        return self.vts_version
+
     def command_exists(self, name):
         """ Checks if a commands exists. """
         return name in self.commands.keys()
@@ -1464,7 +1482,15 @@ def handle_get_version_command(self):
             elem = SubElement(scanner, name)
             elem.text = value
 
-        return simple_response_str('get_version', 200, 'OK', [protocol, daemon, scanner])
+        content = [protocol, daemon, scanner]
+
+        if self.get_vts_version():
+            vts = Element('vts')
+            elem = SubElement(vts, 'version')
+            elem.text = self.get_vts_version()
+            content.append(vts)
+
+        return simple_response_str('get_version', 200, 'OK', content)
 
     def handle_command(self, command):
         """ Handles an osp command in a string.

tests/testScanAndResult.py

@@ -650,3 +650,13 @@ def testProgress(self):
         daemon.set_scan_target_progress(scan_id, 'localhost1', 75)
         daemon.set_scan_target_progress(scan_id, 'localhost2', 25)
         self.assertEqual(daemon.calculate_progress(scan_id), 50)
+
+    def testSetGetVtsVersion(self):
+        daemon = DummyWrapper([])
+        daemon.set_vts_version('1234')
+        version = daemon.get_vts_version()
+        self.assertEqual('1234', version)
+
+    def testSetGetVtsVersion_Error(self):
+        daemon = DummyWrapper([])
+        self.assertRaises(TypeError, daemon.set_vts_version)