CR changes

monkey/infection_monkey/network/postgresql_finger.py

@@ -29,8 +29,11 @@ class PostgreSQLFinger(HostFinger):
                                 "SSL usage is forced.\n",
             'only_selected': "Only selected hosts can make connections (SSL or non-SSL).\n"
         }
-    RELEVANT_EX_SUBSTRINGS = ["password authentication failed",
+    RELEVANT_EX_SUBSTRINGS =\
-                              "entry for host"]  # "no pg_hba.conf entry for host" but filename may be diff
+        {
+            'no_auth': "password authentication failed",
+            'no_entry': "entry for host"  # "no pg_hba.conf entry for host" but filename may be diff
+        }
 
     def get_host_fingerprint(self, host):
         try:
@@ -45,7 +48,7 @@ class PostgreSQLFinger(HostFinger):
             try:
                 exception_string = str(ex)
 
-                if not self.is_relevant_exception(exception_string):
+                if not self._is_relevant_exception(exception_string):
                     return False
 
                 # all's well; start analyzing errors
@@ -57,8 +60,8 @@ class PostgreSQLFinger(HostFinger):
 
             return False
 
-    def is_relevant_exception(self, exception_string):
+    def _is_relevant_exception(self, exception_string):
-        if not any(substr in exception_string for substr in self.RELEVANT_EX_SUBSTRINGS):
+        if not any(substr in exception_string for substr in self.RELEVANT_EX_SUBSTRINGS.values()):
             # OperationalError due to some other reason - irrelevant exception
             return False
         return True
@@ -71,32 +74,10 @@ class PostgreSQLFinger(HostFinger):
         ssl_connection_details = []
         ssl_conf_on_server = self.is_ssl_configured(exceptions)
 
-        # SSL configured
+        if ssl_conf_on_server:  # SSL configured
-        if ssl_conf_on_server:
+            self.get_connection_details_ssl_configured()
-            ssl_connection_details.append(self.CONNECTION_DETAILS['ssl_conf'])
+        else:  # SSL not configured
-            # SSL
+            self.get_connection_details_ssl_not_configured()
-            ssl_selected_comms_only = False
-            if self.found_entry_for_host_but_pwd_auth_failed(exceptions[0]):
-                ssl_connection_details.append(self.CONNECTION_DETAILS['all_ssl'])
-            else:
-                ssl_connection_details.append(self.CONNECTION_DETAILS['selected_ssl'])
-                ssl_selected_comms_only = True
-            # non-SSL
-            if self.found_entry_for_host_but_pwd_auth_failed(exceptions[1]):
-                ssl_connection_details.append(self.CONNECTION_DETAILS['all_non_ssl'])
-            else:
-                if ssl_selected_comms_only:  # if only selected SSL allowed and only selected non-SSL allowed
-                    ssl_connection_details[-1] = self.CONNECTION_DETAILS['only_selected']
-                else:
-                    ssl_connection_details.append(self.CONNECTION_DETAILS['selected_non_ssl'])
-
-        # SSL not configured
-        else:
-            ssl_connection_details.append(self.CONNECTION_DETAILS['ssl_not_conf'])
-            if self.found_entry_for_host_but_pwd_auth_failed(exceptions[0]):
-                ssl_connection_details.append(self.CONNECTION_DETAILS['all_non_ssl'])
-            else:
-                ssl_connection_details.append(self.CONNECTION_DETAILS['selected_non_ssl'])
 
         host.services[self._SCANNED_SERVICE]['communication_encryption_details'] = ''.join(ssl_connection_details)
 
@@ -109,8 +90,35 @@ class PostgreSQLFinger(HostFinger):
         elif len(exceptions) == 2:  # SSL configured so checks for both
             return True
 
+    def get_connection_details_ssl_configured(self):
+        ssl_connection_details.append(self.CONNECTION_DETAILS['ssl_conf'])
+        ssl_selected_comms_only = False
+
+        # check exception message for SSL connection
+        if self.found_entry_for_host_but_pwd_auth_failed(exceptions[0]):
+            ssl_connection_details.append(self.CONNECTION_DETAILS['all_ssl'])
+        else:
+            ssl_connection_details.append(self.CONNECTION_DETAILS['selected_ssl'])
+            ssl_selected_comms_only = True
+
+        # check exception message for non-SSL connection
+        if self.found_entry_for_host_but_pwd_auth_failed(exceptions[1]):
+            ssl_connection_details.append(self.CONNECTION_DETAILS['all_non_ssl'])
+        else:
+            if ssl_selected_comms_only:  # if only selected SSL allowed and only selected non-SSL allowed
+                ssl_connection_details[-1] = self.CONNECTION_DETAILS['only_selected']
+            else:
+                ssl_connection_details.append(self.CONNECTION_DETAILS['selected_non_ssl'])
+
+    def get_connection_details_ssl_not_configured(self):
+        ssl_connection_details.append(self.CONNECTION_DETAILS['ssl_not_conf'])
+        if self.found_entry_for_host_but_pwd_auth_failed(exceptions[0]):
+            ssl_connection_details.append(self.CONNECTION_DETAILS['all_non_ssl'])
+        else:
+            ssl_connection_details.append(self.CONNECTION_DETAILS['selected_non_ssl'])
+
     @staticmethod
     def found_entry_for_host_but_pwd_auth_failed(exception):
-        if PostgreSQLFinger.RELEVANT_EX_SUBSTRINGS[0] in exception:
+        if PostgreSQLFinger.RELEVANT_EX_SUBSTRINGS['no_auth'] in exception:
             return True  # entry found in pg_hba.conf file but password authentication failed
         return False  # entry not found in pg_hba.conf file

monkey/infection_monkey/network/test_postgresql_finger.py

@@ -1,7 +1,7 @@
 from unittest import TestCase
 from unittest.mock import Mock
 
-from infection_monkey.network.postgresql_fingerprint import PostgreSQLFinger
+from infection_monkey.network.postgresql_finger import PostgreSQLFinger
 
 IRRELEVANT_EXCEPTION_STRING = "This is an irrelevant exception string."
 
@@ -76,9 +76,9 @@ EXAMPLE_EXCEPTIONS_WITH_EXPECTED_RESULTS =\
 
 class TestPostgreSQLFinger(TestCase):
     def test_is_relevant_exception(self):
-        assert PostgreSQLFinger().is_relevant_exception(IRRELEVANT_EXCEPTION_STRING) is False
+        assert PostgreSQLFinger()._is_relevant_exception(IRRELEVANT_EXCEPTION_STRING) is False
         for exception_string in EXAMPLE_EXCEPTIONS_WITH_EXPECTED_RESULTS:
-            assert PostgreSQLFinger().is_relevant_exception(exception_string) is True
+            assert PostgreSQLFinger()._is_relevant_exception(exception_string) is True
 
     def test_analyze_operational_error(self):
         host = Mock(['services'])