Import only .py files when importing rules (#575)

.pylintrc

@@ -358,7 +358,7 @@ valid-metaclass-classmethod-first-arg=mcs
 max-args=5
 
 # Maximum number of attributes for a class (see R0902).
-max-attributes=10
+max-attributes=15
 
 # Maximum number of boolean expressions in a if statement
 max-bool-expr=5

stream_alert/rule_processor/main.py

@@ -18,20 +18,40 @@
 
 from stream_alert.rule_processor.handler import StreamAlert
 
-modules_to_import = set()
-# walk the rules directory to dymanically import
-for folder in ('matchers', 'rules'):
-    for root, dirs, files in os.walk(folder):
-        filtered_files = [rule_file for rule_file in files if not (rule_file.startswith((
-            '.', '__init__')) or rule_file.endswith('.pyc'))]
-        package_path = root.replace('/', '.')
-        for import_file in filtered_files:
-            import_module = os.path.splitext(import_file)[0]
-            if package_path and import_module:
-                modules_to_import.add('{}.{}'.format(package_path, import_module))
-
-for module_name in modules_to_import:
-    importlib.import_module(module_name)
+
+def _python_rule_paths():
+    """Yields all .py files in matchers/ and rules/."""
+    # 'matchers' and 'rules' are top-level folders, both in the repo (when testing)
+    # and in the generated Lambda packages.
+    for folder in ('matchers', 'rules'):
+        for root, _, files in os.walk(folder):
+            for file_name in files:
+                if file_name.endswith('.py') and not file_name.startswith('__'):
+                    yield os.path.join(root, file_name)
+
+
+def _path_to_module(path):
+    """Convert a Python rules file path to an importable module name.
+
+    For example, "rules/community/cloudtrail_critical_api_calls.py" becomes
+    "rules.community.cloudtrail_critical_api_calls"
+
+    Raises:
+        NameError if a '.' appears anywhere in the path except the file extension.
+    """
+    base_name = os.path.splitext(path)[0]
+    if '.' in base_name:
+        raise NameError('Python file {} cannot be imported because of "." in the name', path)
+    return os.path.splitext(path)[0].replace('/', '.')
+
+
+def _import_rules():
+    """Dynamically import all rules files."""
+    for path in _python_rule_paths():
+        importlib.import_module(_path_to_module(path))
+
+
+_import_rules()
 
 
 def handler(event, context):

tests/unit/stream_alert_rule_processor/test_init.py

@@ -0,0 +1,44 @@
+"""
+Copyright 2017-present, Airbnb Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+   http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+import os
+
+from mock import call, patch
+from nose.tools import assert_equal
+
+import stream_alert.rule_processor as rp
+
+
+@patch('stream_alert.rule_processor.LOGGER.error')
+@patch.dict(os.environ, {'LOGGER_LEVEL': 'INVALID'})
+def test_init_logging_bad(log_mock):
+    """Rule Processor Init - Logging, Bad Level"""
+    # Force reload the rule_processor package to trigger the init
+    reload(rp)
+
+    message = str(call('Defaulting to INFO logging: %s',
+                       ValueError('Unknown level: \'INVALID\'',)))
+
+    assert_equal(str(log_mock.call_args_list[0]), message)
+
+
+@patch('stream_alert.rule_processor.LOGGER.setLevel')
+@patch.dict(os.environ, {'LOGGER_LEVEL': '10'})
+def test_init_logging_int_level(log_mock):
+    """Rule Processor Init - Logging, Integer Level"""
+    # Force reload the rule_processor package to trigger the init
+    reload(rp)
+
+    log_mock.assert_called_with(10)

tests/unit/stream_alert_rule_processor/test_main.py

@@ -13,55 +13,74 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
-import os
-
 from mock import call, patch
-from nose.tools import assert_equal, with_setup, nottest
-
-import stream_alert.rule_processor as rp
-from tests.unit.stream_alert_rule_processor.test_helpers import get_mock_context
-
-
-def _teardown_env():
-    """Helper method to reset environment variables"""
-    if 'LOGGER_LEVEL' in os.environ:
-        del os.environ['LOGGER_LEVEL']
-
-# TODO(Jack) Investigate flakey test
-@nottest
-@patch('stream_alert.rule_processor.main.StreamAlert.run')
-def test_handler(mock_runner):
-    """Rule Processor Main - Test Handler"""
-    rp.main.handler('event', get_mock_context())
-    mock_runner.assert_called_with('event')
-
-
-@with_setup(setup=None, teardown=_teardown_env)
-@patch('stream_alert.rule_processor.LOGGER.error')
-def test_init_logging_bad(log_mock):
-    """Rule Processor Init - Logging, Bad Level"""
-    level = 'IFNO'
-
-    os.environ['LOGGER_LEVEL'] = level
-
-    # Force reload the rule_processor package to trigger the init
-    reload(rp)
-
-    message = str(call('Defaulting to INFO logging: %s',
-                       ValueError('Unknown level: \'IFNO\'',)))
-
-    assert_equal(str(log_mock.call_args_list[0]), message)
-
-
-@with_setup(setup=None, teardown=_teardown_env)
-@patch('stream_alert.rule_processor.LOGGER.setLevel')
-def test_init_logging_int_level(log_mock):
-    """Rule Processor Init - Logging, Integer Level"""
-    level = '10'
-
-    os.environ['LOGGER_LEVEL'] = level
-
-    # Force reload the rule_processor package to trigger the init
-    reload(rp)
-
-    log_mock.assert_called_with(10)
+from nose.tools import assert_equal, assert_raises
+from pyfakefs import fake_filesystem_unittest
+
+from stream_alert.rule_processor import main
+
+
+class RuleImportTest(fake_filesystem_unittest.TestCase):
+    """Test rule import logic with a mocked filesystem."""
+    # pylint: disable=protected-access
+
+    def setUp(self):
+        self.setUpPyfakefs()
+
+        # Add rules files which should be imported.
+        self.fs.CreateFile('matchers/matchers.py')
+        self.fs.CreateFile('rules/example.py')
+        self.fs.CreateFile('rules/community/cloudtrail/critical_api.py')
+
+        # Add other files which should NOT be imported.
+        self.fs.CreateFile('matchers/README')
+        self.fs.CreateFile('rules/__init__.py')
+        self.fs.CreateFile('rules/example.pyc')
+        self.fs.CreateFile('rules/community/REVIEWERS')
+
+    def tearDown(self):
+        self.tearDownPyfakefs()
+
+    @staticmethod
+    def test_python_rule_paths():
+        """Rule Processor Main - Find rule paths"""
+        result = set(main._python_rule_paths())
+        expected = {
+            'matchers/matchers.py',
+            'rules/example.py',
+            'rules/community/cloudtrail/critical_api.py'
+        }
+        assert_equal(expected, result)
+
+    @staticmethod
+    def test_path_to_module():
+        """Rule Processor Main - Convert rule path to module name"""
+        assert_equal('name', main._path_to_module('name.py'))
+        assert_equal('a.b.c.name', main._path_to_module('a/b/c/name.py'))
+
+    @staticmethod
+    def test_path_to_module_invalid():
+        """Rule Processor Main - Raise NameError for invalid Python filename."""
+        assert_raises(NameError, main._path_to_module, 'a.b.py')
+        assert_raises(NameError, main._path_to_module, 'a/b/old.name.py')
+
+    @staticmethod
+    @patch.object(main, 'importlib')
+    def test_import_rules(mock_importlib):
+        """Rule Processor Main - Import all rule modules."""
+        main._import_rules()
+        mock_importlib.assert_has_calls([
+            call.import_module('matchers.matchers'),
+            call.import_module('rules.example'),
+            call.import_module('rules.community.cloudtrail.critical_api')
+        ], any_order=True)
+
+    @staticmethod
+    @patch.object(main, 'StreamAlert')
+    def test_handler(mock_stream_alert):
+        """Rule Processor Main - Handler is invoked"""
+        main.handler('event', 'context')
+        mock_stream_alert.assert_has_calls([
+            call('context'),
+            call().run('event')
+        ])