From 7f1c157d7e5ecd4bced3643fae49ef8d83c3fd3d Mon Sep 17 00:00:00 2001
From: lachmanfrantisek <flachman@redhat.com>
Date: Tue, 6 Nov 2018 09:49:43 +0100
Subject: [PATCH] Add timeout for each check

Signed-off-by: lachmanfrantisek <flachman@redhat.com>
---
 colin/core/check_runner.py          |  9 ++++++-
 colin/core/checks/abstract_check.py |  1 +
 colin/core/constant.py              |  2 ++
 colin/utils/cmd_tools.py            | 37 ++++++++++++++++++++++++++++-
 4 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/colin/core/check_runner.py b/colin/core/check_runner.py
index de7c43b3..fddb9c10 100644
--- a/colin/core/check_runner.py
+++ b/colin/core/check_runner.py
@@ -17,7 +17,9 @@
 import logging
 import traceback
 
+from .constant import CHECK_TIMEOUT
 from .result import CheckResults, FailedCheckResult
+from ..utils.cmd_tools import exit_after
 
 logger = logging.getLogger(__name__)
 
@@ -34,7 +36,12 @@ def _result_generator(target, checks):
         for check in checks:
             logger.debug("Checking {}".format(check.name))
             try:
-                yield check.check(target)
+                timeout = check.timeout or CHECK_TIMEOUT
+                yield exit_after(timeout)(check.check)(target)
+            except TimeoutError as ex:
+                logger.warning(
+                    "The check hit the timeout.")
+                yield FailedCheckResult(check, logs=[str(ex)])
             except Exception as ex:
                 tb = traceback.format_exc()
                 logger.warning(
diff --git a/colin/core/checks/abstract_check.py b/colin/core/checks/abstract_check.py
index bf59bd4f..05a6d435 100644
--- a/colin/core/checks/abstract_check.py
+++ b/colin/core/checks/abstract_check.py
@@ -27,6 +27,7 @@ def __init__(self, message, description, reference_url, tags):
         self.description = description
         self.reference_url = reference_url
         self.tags = tags
+        self.timeout = None
 
     def check(self, target):
         pass
diff --git a/colin/core/constant.py b/colin/core/constant.py
index 85c62922..e89d02dc 100644
--- a/colin/core/constant.py
+++ b/colin/core/constant.py
@@ -35,3 +35,5 @@
 }
 
 COLIN_CHECKS_PATH = "CHECKS_PATH"
+
+CHECK_TIMEOUT = 10 * 60  # s
diff --git a/colin/utils/cmd_tools.py b/colin/utils/cmd_tools.py
index 6f4dbdf3..2da41187 100644
--- a/colin/utils/cmd_tools.py
+++ b/colin/utils/cmd_tools.py
@@ -13,8 +13,16 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 #
-
+import logging
 import subprocess
+import threading
+
+try:
+    import thread
+except ImportError:
+    import _thread as thread
+
+logger = logging.getLogger(__name__)
 
 
 def get_version_of_the_python_package(module):
@@ -85,3 +93,30 @@ def is_rpm_installed():
     except FileNotFoundError:
         rpm_installed = False
     return rpm_installed
+
+
+def exit_after(s):
+    """
+    Use as decorator to exit process if
+    function takes longer than s seconds.
+
+    Direct call is available via exit_after(TIMEOUT_IN_S)(fce)(args).
+
+    Inspired by https://stackoverflow.com/a/31667005
+    """
+
+    def outer(fn):
+        def inner(*args, **kwargs):
+            timer = threading.Timer(s, thread.interrupt_main)
+            timer.start()
+            try:
+                result = fn(*args, **kwargs)
+            except KeyboardInterrupt:
+                raise TimeoutError("Function '{}' hit the timeout ({}s).".format(fn.__name__, s))
+            finally:
+                timer.cancel()
+            return result
+
+        return inner
+
+    return outer