Merge pull request #3 from cisagov/improvement/retry-timeout-args

Add --retries and --timeout optional arguments to command

.isort.cfg

@@ -7,6 +7,6 @@ import_heading_thirdparty=Third-Party Libraries
 import_heading_firstparty=cisagov Libraries
 
 # Should be auto-populated by seed-isort-config hook
-known_third_party=docopt,pytest,setuptools,urllib3
+known_third_party=docopt,pytest,schema,setuptools,urllib3
 # These must be manually set to correctly separate them from third party libraries
 known_first_party=check_cve

README.md

@@ -18,7 +18,7 @@ Python versions 3.6 and above.  Note that Python 2 *is not* supported.
 From a release:
 
 ```console
-pip install https://github.com/cisagov/check-cve-2019-19781/releases/download/v1.0.0/cve_2019_19781-1.0.0-py3-none-any.whl
+pip install https://github.com/cisagov/check-cve-2019-19781/releases/download/v1.0.1/cve_2019_19781-1.0.1-py3-none-any.whl
 ```
 
 From source:
@@ -59,15 +59,18 @@ For more information about this vulnerability see:
     https://nvd.nist.gov/vuln/detail/CVE-2019-19781
 
 Usage:
-  cve-2019-19781 [--log-level=LEVEL] <host>
+  cve-2019-19781 [options] <host>
   cve-2019-19781 (-h | --help)
 
 Options:
   -h --help              Show this message.
   --log-level=LEVEL      If specified, then the log level will be set to
                          the specified value.  Valid values are "debug", "info",
                          "warning", "error", and "critical". [default: info]
-
+  -r --retries=count     Number of times to retry a failed connection attempt before
+                         giving up. [default: 2]
+  -t --timeout=seconds   Number of seconds to wait during each connection attempt.
+                         [default: 10]
 ```
 
 ## Contributing ##

setup.py

@@ -51,7 +51,7 @@ def package_vars(version_file):
         #   3 - Alpha
         #   4 - Beta
         #   5 - Production/Stable
-        "Development Status :: 3 - Alpha",
+        "Development Status :: 5 - Production/Stable",
         # Indicate who your project is intended for
         "Intended Audience :: Developers",
         # Pick your license as you wish (should match "license" above)
@@ -69,7 +69,7 @@ def package_vars(version_file):
     package_dir={"": "src"},
     py_modules=[splitext(basename(path))[0] for path in glob("src/*.py")],
     include_package_data=True,
-    install_requires=["docopt", "setuptools", "urllib3"],
+    install_requires=["docopt", "setuptools", "schema", "urllib3"],
     extras_require={
         "test": [
             "pre-commit",
@@ -83,6 +83,6 @@ def package_vars(version_file):
         ]
     },
     python_requires=">=3.6",
-    # Conveniently allows one to run the CLI tool as `example`
+    # Conveniently allows one to run the CLI tool as `cve-2019-19781`
     entry_points={"console_scripts": ["cve-2019-19781=check_cve.check:main"]},
 )

src/check_cve/_version.py

@@ -1,2 +1,2 @@
 """This file defines the version of this module."""
-__version__ = "1.0.0"
+__version__ = "1.0.1"

src/check_cve/check.py

@@ -13,14 +13,18 @@
     https://nvd.nist.gov/vuln/detail/CVE-2019-19781
 
 Usage:
-  cve-2019-19781 [--log-level=LEVEL] <host>
+  cve-2019-19781 [options] <host>
   cve-2019-19781 (-h | --help)
 
 Options:
   -h --help              Show this message.
   --log-level=LEVEL      If specified, then the log level will be set to
                          the specified value.  Valid values are "debug", "info",
                          "warning", "error", and "critical". [default: info]
+  -r --retries=count     Number of times to retry a failed connection attempt before
+                         giving up. [default: 2]
+  -t --timeout=seconds   Number of seconds to wait during each connection attempt.
+                         [default: 10]
 """
 
 # Standard Python Libraries
@@ -29,14 +33,15 @@
 
 # Third-Party Libraries
 import docopt
+from schema import And, Schema, SchemaError, Use
 import urllib3
 
 from ._version import __version__
 
 INSECURE_CONTENT = r"You don't have permission to access /vpns/"
 
 
-def is_vulnerable(host):
+def is_vulnerable(host, retries=2, timeout=10):
     """Make an http request to a host to see if it is vulnerable."""
     # Suppress insecure request warning
     logging.debug("Disabling insecure request warnings.")
@@ -46,8 +51,8 @@ def is_vulnerable(host):
     logging.debug("Disabling URL normalization.")
     urllib3.util.url.NORMALIZABLE_SCHEMES = ()
 
-    logging.debug("Creating connection pool")
-    pool = urllib3.PoolManager(cert_reqs="CERT_NONE")
+    logging.debug(f"Creating connection pool, retries={retries}, timeout={timeout}s")
+    pool = urllib3.PoolManager(cert_reqs="CERT_NONE", retries=retries, timeout=timeout)
 
     # Build URL to solicit a telling response
     url = f"https://{host}/vpn/../vpns/"
@@ -65,24 +70,52 @@ def is_vulnerable(host):
 def main():
     """Parse arguments, handle logging, and call vulnerability check."""
     args = docopt.docopt(__doc__, version=__version__)
-    # Set up logging
-    log_level = args["--log-level"]
+
+    # Validate and convert arguments as needed
+    schema = Schema(
+        {
+            "--log-level": And(
+                str,
+                Use(str.lower),
+                lambda n: n in ("debug", "info", "warning", "error", "critical"),
+                error="Possible values for --log-level are "
+                + "debug, info, warning, error, and critical.",
+            ),
+            "--retries": And(
+                Use(int),
+                lambda n: n >= 0,
+                error="--retries must be an integer greater than or equal to 0.",
+            ),
+            "--timeout": And(
+                Use(int),
+                lambda n: n > 0,
+                error="--timeout must be an integer greater than 0.",
+            ),
+            str: object,  # Don't care about other keys, if any
+        }
+    )
+
     try:
-        logging.basicConfig(
-            format="%(asctime)-15s %(levelname)s %(message)s", level=log_level.upper()
-        )
-    except ValueError:
-        logging.critical(
-            f'"{log_level}" is not a valid logging level.  Possible values '
-            "are debug, info, warning, and error."
-        )
-        return 254
+        args = schema.validate(args)
+    except SchemaError as err:
+        # Exit because one or more of the arguments were invalid
+        print(err, file=sys.stderr)
+        return 1
 
+    # Assign validated arguments to variables
     host = args["<host>"]
+    log_level = args["--log-level"]
+    retries = args["--retries"]
+    timeout = args["--timeout"]
+
+    # Set up logging
+    logging.basicConfig(
+        format="%(asctime)-15s %(levelname)s %(message)s", level=log_level.upper()
+    )
 
     try:
-        if is_vulnerable(host):
-            logging.warn(f"{host} appears to be vulnerable.")
+        if is_vulnerable(host, retries, timeout):
+            logging.warning(f"{host} appears to be vulnerable.")
             result = 2
         else:
             logging.info(f"{host} does not appear to be vulnerable.")

tests/test_check_cve.py

@@ -26,6 +26,7 @@
 # define sources of version strings
 RELEASE_TAG = os.getenv("RELEASE_TAG")
 PROJECT_VERSION = check_cve.__version__
+INSECURE_BYTES = br"You don't have permission to access /vpns/"
 
 
 def test_stdout_version(capsys):
@@ -63,26 +64,67 @@ def test_log_levels(level):
             assert (
                 logging.root.hasHandlers() is True
             ), "root logger should now have a handler"
-            assert return_code == 3, "main() should return a connection error (3)"
+            assert return_code == 3, "main() should return a connection error"
 
 
 def test_missing_host_arg():
     """Verify the utility handles missing host argument."""
     with pytest.raises(docopt.DocoptExit):
         with patch.object(sys, "argv", ["exe_name"]):
             return_code = check_cve.check.main()
-            assert return_code == 254, "main() should return an error (254)"
+            assert return_code == 254, "main() should return an error"
 
 
 def test_connection_error():
     """Verify the utility handles missing host argument."""
     with patch.object(sys, "argv", ["exe_name", "bogus.bogus.bogus"]):
         return_code = check_cve.check.main()
-        assert return_code == 3, "main() should return an error (2)"
+        assert return_code == 3, "main() should return an error"
 
 
-def test_non_vuln_host():
-    """Verify the utility handles non-vulnerable host."""
-    with patch.object(sys, "argv", ["exe_name", "github.com"]):
+def test_valid_timeout():
+    """Test a valid timeout argument."""
+    with patch.object(sys, "argv", ["exe_name", "--timeout=5", "github.com"]):
         return_code = check_cve.check.main()
-        assert return_code == 0, "main() should return success (0)"
+        assert return_code == 0, "main() should return success"
+
+
+def test_invalid_timeout():
+    """Test an invalid timeout argument."""
+    with patch.object(sys, "argv", ["exe_name", "--timeout=0", "bogus.bogus.bogus"]):
+        return_code = check_cve.check.main()
+        assert return_code == 1, "main() should return an error"
+
+
+def test_valid_retries():
+    """Test a valid timeout argument."""
+    with patch.object(sys, "argv", ["exe_name", "--retries=0", "github.com"]):
+        return_code = check_cve.check.main()
+        assert return_code == 0, "main() should return success"
+
+
+def test_invalid_retries():
+    """Test an invalid timeout argument."""
+    with patch.object(sys, "argv", ["exe_name", "--retries=-1", "bogus.bogus.bogus"]):
+        return_code = check_cve.check.main()
+        assert return_code == 1, "main() should return an error"
+
+
+def test_vuln_host():
+    """Verify the utility handles vulnerable host."""
+    with patch.object(sys, "argv", ["exe_name", "--log-level=debug", "github.com"]):
+        with patch("check_cve.check.urllib3") as mock_lib:
+            mock_lib.PoolManager().request().data = INSECURE_BYTES
+            mock_lib.PoolManager().request().status = 403
+            return_code = check_cve.check.main()
+            assert return_code == 2, "main() should return error signaling detection"
+
+
+def test_non_vuln_host():
+    """Verify the utility handles a non-vulnerable host."""
+    with patch.object(sys, "argv", ["exe_name", "--log-level=debug", "github.com"]):
+        with patch("check_cve.check.urllib3") as mock_lib:
+            mock_lib.PoolManager().request().data = b""
+            mock_lib.PoolManager().request().status = 403
+            return_code = check_cve.check.main()
+            assert return_code == 0, "main() should return success"