From 7db11971e9ee8141e446ab1e890aefb93862776e Mon Sep 17 00:00:00 2001
From: Travis McPeak <travis.mcpeak@hpe.com>
Date: Thu, 25 Feb 2016 16:41:18 -0800
Subject: [PATCH] Adding profile generation to config generator

This enables config generator to create and save a profile
by embedding the possible plugin settings to a template which
includes some comments explaining the layout of the profile.

It is expected a user will delete the parts of the profile they
don't need and leave the parts they want.

Change-Id: Ie00aedfcf463685e0f15b793af9e2d9490a13180
---
 bandit/cli/config_generator.py          | 62 ++++++++++++++++++++++---
 tests/unit/cli/test_config_generator.py | 16 +++----
 2 files changed, 62 insertions(+), 16 deletions(-)

diff --git a/bandit/cli/config_generator.py b/bandit/cli/config_generator.py
index fabee449..a7477411 100644
--- a/bandit/cli/config_generator.py
+++ b/bandit/cli/config_generator.py
@@ -16,6 +16,7 @@ from __future__ import print_function
 import argparse
 import importlib
 import logging
+import os
 import sys
 
 from stevedore import extension
@@ -25,10 +26,26 @@ PROG_NAME = 'bandit_conf_generator'
 logger = logging.getLogger(__name__)
 
 
+template = """
+### profile may optionally select or skip tests
+
+# (optional) list included tests here:
+# tests: B101,B102
+
+# (optional) list skipped tests here:
+# skip: B201, B202
+
+
+### override settings - used to set settings for plugins to non-default values
+
+{settings}
+"""
+
+
 def init_logger():
     logger.handlers = []
     log_level = logging.INFO
-    log_format_string = "[bandit-config-generator] %(message)s"
+    log_format_string = "[%(levelname)5s]: %(message)s"
     logging.captureWarnings(True)
     logger.setLevel(log_level)
     handler = logging.StreamHandler(sys.stdout)
@@ -37,13 +54,28 @@ def init_logger():
 
 
 def parse_args():
-    parser = argparse.ArgumentParser(
-        description='Tool to display Bandit config options')
+    help_description = """Bandit Config Generator
 
-    parser.add_argument('--show-defaults', dest='show_defaults',
+    This tool is used to generate an optional profile.  The profile may be used
+    to include or skip tests and override values for plugins.
+
+    When used to store an output profile, this tool will output a template that
+    includes all plugins and their default settings.  Any settings which aren't
+    being overridden can be safely removed from the profile and default values
+    will be used.  Bandit will prefer settings from the profile over the built
+    in values."""
+
+    parser = argparse.ArgumentParser(
+        description=help_description,
+        formatter_class=argparse.RawTextHelpFormatter)
+
+    parser.add_argument('-s', '--show-defaults', dest='show_defaults',
                         action='store_true',
                         help='show the default settings values for each '
-                             'plugin')
+                             'plugin but do not output a profile')
+    parser.add_argument('-o', '--out', dest='output_file',
+                        action='store',
+                        help='output file to save profile')
 
     args = parser.parse_args()
     return args
@@ -78,8 +110,26 @@ def main():
     init_logger()
     args = parse_args()
 
+    yaml_settings = get_config_settings()
+
     if args.show_defaults:
-        print(get_config_settings())
+        print(yaml_settings)
+
+    if args.output_file:
+        if os.path.exists(os.path.abspath(args.output_file)):
+            logger.error("File %s already exists, exiting", args.output_file)
+            sys.exit(2)
+
+        try:
+            with open(args.output_file, 'w') as f:
+                contents = template.format(settings=yaml_settings)
+                f.write(contents)
+
+        except IOError:
+            logger.error("Unable to open %s for writing", args.output_file)
+
+        else:
+            logger.info("Successfully wrote profile: %s", args.output_file)
 
     return 0
 
diff --git a/tests/unit/cli/test_config_generator.py b/tests/unit/cli/test_config_generator.py
index 0a39b2e2..81ec0218 100644
--- a/tests/unit/cli/test_config_generator.py
+++ b/tests/unit/cli/test_config_generator.py
@@ -84,6 +84,12 @@ class BanditConfigGeneratorTests(testtools.TestCase):
         return_value = config_generator.parse_args()
         self.assertTrue(return_value.show_defaults)
 
+    @mock.patch('sys.argv', ['bandit-config-generator', '--out', 'dummyfile'])
+    def test_parse_args_out_file(self):
+        # Test config generator get proper output file when specified
+        return_value = config_generator.parse_args()
+        self.assertEqual('dummyfile', return_value.output_file)
+
     def test_get_config_settings(self):
         settings = config_generator.get_config_settings()
         self.assertEqual(settings, "test: {test: test data}\n")
@@ -97,13 +103,3 @@ class BanditConfigGeneratorTests(testtools.TestCase):
             # The get_config_settings function should have been called
             self.assertTrue(mock_config_settings.called)
             self.assertEqual(0, return_value)
-
-    @mock.patch('sys.argv', ['bandit-config-generator'])
-    def test_main_no_defaults(self):
-        # Test that the config generator does not show defaults and returns 0
-        with mock.patch('bandit.cli.config_generator.get_config_settings'
-                        ) as mock_config_settings:
-            return_value = config_generator.main()
-            # The get_config_settings function should not have been called
-            self.assertFalse(mock_config_settings.called)
-            self.assertEqual(0, return_value)