Exit code 1 if unsanitised vulnerabilities found

pyt/__main__.py

@@ -22,6 +22,7 @@
     get_vulnerabilities_not_in_baseline,
     UImode
 )
+from .vulnerabilities.vulnerability_helper import SanitisedVulnerability
 from .web_frameworks import (
     FrameworkAdaptor,
     is_django_view_function,
@@ -137,6 +138,10 @@ def main(command_line_args=sys.argv[1:]):  # noqa: C901
     else:
         text.report(vulnerabilities, args.output_file)
 
+    has_unsanitized_vulnerabilities = any(not isinstance(v, SanitisedVulnerability) for v in vulnerabilities)
+    if has_unsanitized_vulnerabilities:
+        sys.exit(1)
+
 
 if __name__ == '__main__':
     main()

setup.py

@@ -33,7 +33,7 @@
     install_requires=[],
     entry_points={
         'console_scripts': [
-            'pyt = pyt:main'
+            'pyt = pyt.__main__:main'
         ]
     }
 )

tests/main_test.py

@@ -16,22 +16,41 @@ def test_text_output(self, mock_text, mock_find_vulnerabilities, mock_parse_args
 
         mock_discover_files.return_value = [example_file]
         mock_parse_args.return_value = mock.Mock(
-            autospec=True,
             project_root=None,
             baseline=None,
             json=None,
             output_file=output_file
         )
-        main([
-            'parse_args is mocked'
-        ])
+        with self.assertRaises(SystemExit):
+            main(['parse_args is mocked'])
         assert mock_text.report.call_count == 1
-        # This with: makes no sense
-        with self.assertRaises(AssertionError):
-            assert mock_text.report.assert_called_with(
-                mock_find_vulnerabilities.return_value,
-                mock_parse_args.return_value.output_file
-            )
+        mock_text.report.assert_called_with(
+            mock_find_vulnerabilities.return_value,
+            mock_parse_args.return_value.output_file
+        )
+
+    @mock.patch('pyt.__main__.discover_files')
+    @mock.patch('pyt.__main__.parse_args')
+    @mock.patch('pyt.__main__.find_vulnerabilities')
+    @mock.patch('pyt.__main__.text')
+    def test_no_vulns_found(self, mock_text, mock_find_vulnerabilities, mock_parse_args, mock_discover_files):
+        mock_find_vulnerabilities.return_value = []
+        example_file = 'examples/vulnerable_code/inter_command_injection.py'
+        output_file = 'mocked_outfile'
+
+        mock_discover_files.return_value = [example_file]
+        mock_parse_args.return_value = mock.Mock(
+            project_root=None,
+            baseline=None,
+            json=None,
+            output_file=output_file
+        )
+        main(['parse_args is mocked'])  # No SystemExit
+        assert mock_text.report.call_count == 1
+        mock_text.report.assert_called_with(
+            mock_find_vulnerabilities.return_value,
+            mock_parse_args.return_value.output_file
+        )
 
     @mock.patch('pyt.__main__.discover_files')
     @mock.patch('pyt.__main__.parse_args')
@@ -44,22 +63,18 @@ def test_json_output(self, mock_json, mock_find_vulnerabilities, mock_parse_args
 
         mock_discover_files.return_value = [example_file]
         mock_parse_args.return_value = mock.Mock(
-            autospec=True,
             project_root=None,
             baseline=None,
             json=True,
             output_file=output_file
         )
-        main([
-            'parse_args is mocked'
-        ])
+        with self.assertRaises(SystemExit):
+            main(['parse_args is mocked'])
         assert mock_json.report.call_count == 1
-        # This with: makes no sense
-        with self.assertRaises(AssertionError):
-            assert mock_json.report.assert_called_with(
-                mock_find_vulnerabilities.return_value,
-                mock_parse_args.return_value.output_file
-            )
+        mock_json.report.assert_called_with(
+            mock_find_vulnerabilities.return_value,
+            mock_parse_args.return_value.output_file
+        )
 
 
 class DiscoverFilesTest(BaseTestCase):