From e2841be764ff563b4ace44c88dc1ecd138862372 Mon Sep 17 00:00:00 2001 From: bcaller Date: Fri, 27 Jul 2018 18:03:42 +0100 Subject: [PATCH 1/3] Make assertion in main test work assert_called_with does the assertion. The function has return value None. So don't assert assert_called_with. Fixes the "This with: makes no sense" comment. --- tests/main_test.py | 20 ++++++++------------ 1 file changed, 8 insertions(+), 12 deletions(-) diff --git a/tests/main_test.py b/tests/main_test.py index 4839da51..b87e4dbf 100644 --- a/tests/main_test.py +++ b/tests/main_test.py @@ -26,12 +26,10 @@ def test_text_output(self, mock_text, mock_find_vulnerabilities, mock_parse_args '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') @@ -54,12 +52,10 @@ def test_json_output(self, mock_json, mock_find_vulnerabilities, mock_parse_args '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): From 92fe367cd4921f9af3872a69807a210fb64f0e8c Mon Sep 17 00:00:00 2001 From: bcaller Date: Fri, 27 Jul 2018 18:20:57 +0100 Subject: [PATCH 2/3] Exit with error code if there are unsanitised vulns To be useful as a linter process such as in a Continuous Integration system, pyt should finish with pass or fail exit codes. Saves having to grep the output. If there are unsanitised vulnerabilities, sys.exit(1). In the future we'll probably want a flag to not print sanitised vulns. --- pyt/__main__.py | 5 +++++ tests/main_test.py | 35 +++++++++++++++++++++++++++-------- 2 files changed, 32 insertions(+), 8 deletions(-) diff --git a/pyt/__main__.py b/pyt/__main__.py index 5eed4747..192eaf99 100644 --- a/pyt/__main__.py +++ b/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() diff --git a/tests/main_test.py b/tests/main_test.py index b87e4dbf..b01cc8ee 100644 --- a/tests/main_test.py +++ b/tests/main_test.py @@ -16,15 +16,36 @@ 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 + 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, @@ -42,15 +63,13 @@ 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 mock_json.report.assert_called_with( mock_find_vulnerabilities.return_value, From d4cbde04fc9c30c875244f17e0fe25aba48ffe17 Mon Sep 17 00:00:00 2001 From: bcaller Date: Fri, 27 Jul 2018 18:27:39 +0100 Subject: [PATCH 3/3] Fix pyt console script Can now use just `pyt` exactly like `python -m pyt`. Before there was an error: AttributeError: module 'pyt' has no attribute 'main' --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index 17c81594..8d139dc0 100644 --- a/setup.py +++ b/setup.py @@ -33,7 +33,7 @@ install_requires=[], entry_points={ 'console_scripts': [ - 'pyt = pyt:main' + 'pyt = pyt.__main__:main' ] } )