From 7364391a94695babe5f864c46c94efb9852496b7 Mon Sep 17 00:00:00 2001
From: lachmanfrantisek <flachman@redhat.com>
Date: Wed, 7 Nov 2018 13:59:19 +0100
Subject: [PATCH] Add timeout to CLI

Signed-off-by: lachmanfrantisek <flachman@redhat.com>
---
 colin/cli/colin.py         |  5 ++++-
 colin/core/check_runner.py | 10 ++++++----
 colin/core/colin.py        |  4 +++-
 3 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/colin/cli/colin.py b/colin/cli/colin.py
index e4d95aeb..1f181070 100644
--- a/colin/cli/colin.py
+++ b/colin/cli/colin.py
@@ -76,10 +76,12 @@ def cli():
 @click.option('--target-type', type=click.STRING, default="image",
               help="Type of selected target (one of image, dockerfile, "
                    "ostree). For ostree, please specify image name and path like this: image@path")
+@click.option('--timeout', type=click.INT,
+              help="Timeout for each check in seconds. (default=600)")
 @click.option('--insecure', is_flag=True, default=False,
               help="Pull from an insecure registry (HTTP or invalid TLS).")
 def check(target, ruleset, ruleset_file, debug, json, stat, skip, tag, verbose,
-          checks_paths, target_type, pull, insecure):
+          checks_paths, target_type, timeout, pull, insecure):
     """
     Check the image/dockerfile (default).
     """
@@ -106,6 +108,7 @@ def check(target, ruleset, ruleset_file, debug, json, stat, skip, tag, verbose,
             pull=pull,
             checks_paths=checks_paths,
             target_type=target_type,
+            timeout=timeout,
             insecure=insecure,
             skips=skip
         )
diff --git a/colin/core/check_runner.py b/colin/core/check_runner.py
index fddb9c10..0825660c 100644
--- a/colin/core/check_runner.py
+++ b/colin/core/check_runner.py
@@ -24,19 +24,21 @@
 logger = logging.getLogger(__name__)
 
 
-def go_through_checks(target, checks):
+def go_through_checks(target, checks, timeout=None):
     logger.debug("Going through checks.")
     results = _result_generator(target=target,
-                                checks=checks)
+                                checks=checks,
+                                timeout=timeout)
     return CheckResults(results=results)
 
 
-def _result_generator(target, checks):
+def _result_generator(target, checks, timeout=None):
     try:
         for check in checks:
             logger.debug("Checking {}".format(check.name))
             try:
-                timeout = check.timeout or CHECK_TIMEOUT
+                timeout = timeout or check.timeout or CHECK_TIMEOUT
+                logger.debug("Check timeout: {}".format(timeout))
                 yield exit_after(timeout)(check.check)(target)
             except TimeoutError as ex:
                 logger.warning(
diff --git a/colin/core/colin.py b/colin/core/colin.py
index 5e19192a..f57b0333 100644
--- a/colin/core/colin.py
+++ b/colin/core/colin.py
@@ -35,10 +35,12 @@ def run(
     pull=None,
     insecure=False,
     skips=None,
+    timeout=None,
 ):
     """
     Runs the sanity checks for the target.
 
+    :param timeout: timeout per-check (in seconds)
     :param skips: name of checks to skip
     :param target: str (image name, ostree or dockertar)
                     or Image (instance from conu)
@@ -72,7 +74,7 @@ def run(
         checks_paths=checks_paths,
         skips=skips,
     )
-    result = go_through_checks(target=target, checks=checks_to_run)
+    result = go_through_checks(target=target, checks=checks_to_run, timeout=timeout)
     return result