From ce2f8fdc8846c5494a08284e8cb143ac4780e7c1 Mon Sep 17 00:00:00 2001
From: Juan Jose Nicola <juan.nicola@greenbone.net>
Date: Fri, 29 Jan 2021 15:41:36 +0100
Subject: [PATCH] Fix get_count_total(). Accept -1 value set by the server.

The value set by the server has priority over the value
calculated from the original target list by ospd.
As ospd is not intelligent enough to check the amount of valid
hosts, check for duplicated or invalid hosts, consider a negative
value set for the server, in case it detects an invalid target string
or a different amount than the orignal amount in the target list.
---
 CHANGELOG.md                  |  1 +
 ospd/scan.py                  | 15 +++++++++++++--
 tests/test_scan_and_result.py | 30 ++++++++++++++++++++++++++++++
 3 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index fe5d3419..8c4008be 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -17,6 +17,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 - Fix target_to_ipv4_short(). [#338](https://github.com/greenbone/ospd/pull/338)
 - Fix malformed target. [#341](https://github.com/greenbone/ospd/pull/341)
 - Initialize end_time with create_scan. [#354](https://github.com/greenbone/ospd/pull/354)
+- Fix get_count_total(). Accept -1 value set by the server. [#355](https://github.com/greenbone/ospd/pull/355)
 
 [20.8.2]: https://github.com/greenbone/ospd/compare/v20.8.1...ospd-20.08
 
diff --git a/ospd/scan.py b/ospd/scan.py
index 77e6d7c7..8c7eab19 100644
--- a/ospd/scan.py
+++ b/ospd/scan.py
@@ -393,7 +393,18 @@ def get_count_total(self, scan_id: str) -> int:
         """ Get a scan's total host count. """
 
         count_total = self.scans_table[scan_id]['count_total']
-        if not count_total:
+
+        # The value set by the server has priority over the value
+        # calculated from the original target list by ospd.
+        # As ospd is not intelligent enough to check the amount of valid
+        # hosts, check for duplicated or invalid hosts, consider a negative
+        # value set for the server, in case it detects an invalid target string
+        # or a different amount than the orignal amount in the target list.
+        if count_total == -1:
+            count_total = 0
+        # If the server does not set the total host count
+        # ospd set the amount of host from the original host list.
+        elif not count_total:
             count_total = self.get_host_count(scan_id)
             self.update_count_total(scan_id, count_total)
 
@@ -458,7 +469,7 @@ def calculate_target_progress(self, scan_id: str) -> int:
             )
         except ZeroDivisionError:
             # Consider the case in which all hosts are dead or excluded
-            LOGGER.debug('%s: All hosts dead or excluded.')
+            LOGGER.debug('%s: All hosts dead or excluded.', scan_id)
             t_prog = ScanProgress.FINISHED.value
 
         return t_prog
diff --git a/tests/test_scan_and_result.py b/tests/test_scan_and_result.py
index b11f3108..b3cbbf9f 100644
--- a/tests/test_scan_and_result.py
+++ b/tests/test_scan_and_result.py
@@ -1162,6 +1162,36 @@ def test_set_scan_total_hosts(self):
         count = self.daemon.scan_collection.get_count_total(scan_id)
         self.assertEqual(count, 3)
 
+        self.daemon.set_scan_total_hosts(scan_id, 3)
+        count = self.daemon.scan_collection.get_count_total(scan_id)
+        self.assertEqual(count, 3)
+
+    def test_set_scan_total_hosts_invalid_target(self):
+
+        fs = FakeStream()
+        self.daemon.handle_command(
+            '<start_scan parallel="2">'
+            '<scanner_params />'
+            '<targets><target>'
+            '<hosts>localhost1, localhost2, localhost3, localhost4</hosts>'
+            '<ports>22</ports>'
+            '</target></targets>'
+            '</start_scan>',
+            fs,
+        )
+        self.daemon.start_queued_scans()
+
+        response = fs.get_response()
+        scan_id = response.findtext('id')
+
+        count = self.daemon.scan_collection.get_count_total(scan_id)
+        self.assertEqual(count, 4)
+
+        # The total host is set by the server as -1, because invalid target
+        self.daemon.set_scan_total_hosts(scan_id, -1)
+        count = self.daemon.scan_collection.get_count_total(scan_id)
+        self.assertEqual(count, 0)
+
     def test_get_scan_progress_xml(self):
 
         fs = FakeStream()