Allowing to trace an already running program

dissect/__init__.py

@@ -17,53 +17,53 @@
 # along with dissect.  If not, see <http://www.gnu.org/licenses/>.
 #endregion
 
-from dissect.consolidation.data_sink import EntityDataSink
+from dissect.consolidation.async_publisher import AsyncPublisher
 from dissect.model.util.entity_id_generator import EntityIdGenerator
 from dissect.modeling.dynamic.driver import Driver \
                                            as DynamicModelingDriver
-from dissect.consolidation.data_source import DataSource
 from dissect.consolidation.observable_model import ObservableModel
 from dissect.modeling.orchestrator import Orchestrator
 from dissect.model.entity.module import Module
+from contextlib import contextmanager
 import logging
-import Queue
 import sys
 import os
 
 class Trace(object):
     def __init__(self, base_path, callback):
-        queue = Queue.Queue()
-        data_source = DataSource(queue)
-        self.model = ObservableModel(data_source)
-        self.data_sink = EntityDataSink(queue, self)
+        self.async_publisher = AsyncPublisher(callback)
+        self.model = ObservableModel(self.async_publisher)
 
         entity_id_generator = EntityIdGenerator(base_path)
         modeling_orchestrator = Orchestrator(base_path, self.model)
         self.dynamic_modeling_driver = DynamicModelingDriver(self,
                                                      entity_id_generator,
                                                      modeling_orchestrator)
-        self.callback = callback
-
-    def handle(self, entity):
-        self.callback(entity)
 
     def start(self):
-        self.data_sink.start()
+        self.async_publisher.start()
         self.dynamic_modeling_driver.start()
 
     def stop(self):
         self.dynamic_modeling_driver.stop()
-        self.data_sink.stop()
+        self.async_publisher.stop()
 
 def run(filepath, callback):
-    dir_path = os.path.dirname(filepath)
-    sys.path = [dir_path] + sys.path
-    trace = Trace(dir_path, callback)
-    trace.start()
+    root_path = os.path.dirname(filepath)
+    sys.path = [root_path] + sys.path
     globals_namespace = {'__file__': filepath,
                          '__name__': '__main__',
                          '__package__': None,
                          '__cached__': None
                         }
-    execfile(filepath, globals_namespace)
-    trace.stop()
+    with trace(root_path, callback):
+        execfile(filepath, globals_namespace)
+
+@contextmanager
+def trace(root_path, callback):
+    trace = Trace(root_path, callback)
+    try:
+        trace.start()
+        yield
+    finally:
+        trace.stop()

dissect/collection/static/source_code_parser.py

@@ -18,6 +18,7 @@
 #endregion
 
 import sys
+import os
 
 from astroid.exceptions import AstroidBuildingException
 from astroid.manager import AstroidManager
@@ -50,6 +51,8 @@ def add_files(self, paths):
         if not isinstance(paths, list):
             paths = [paths]
 
+        paths = [path for path in paths if os.path.isfile(path)]
+
         len_before = len(self.file_paths)
         self.file_paths.update(paths)
         return len(self.file_paths) != len_before
@@ -60,7 +63,7 @@ def register(self, modeler):
     def parse(self):
         manager = AstroidManager()
         project = manager.project_from_files(list(self.file_paths),
-                                func_wrapper=astroid_ignore_modname_wrapper)
+                            func_wrapper=astroid_ignore_modname_wrapper)
 
         # First collect all definitions (e.g. module X, function foo) before
         # trying to relate one definition with another (e.g. module X depends on

dissect/collection/static/test/test_source_code_parser.py

@@ -21,16 +21,19 @@
 from nose.tools import *
 from dissect.collection.static.source_code_parser import SourceCodeParser
 
+
 class TestSourceCodeParser():
 
     def setUp(self):
         self.ast_ng_manager_patcher = patch('dissect.collection.static.source_code_parser.AstroidManager')
         self.ast_ng_manager_mock = Mock()
         self.ast_ng_manager_class_mock = self.ast_ng_manager_patcher.start()
         self.ast_ng_manager_class_mock.return_value = self.ast_ng_manager_mock
+        self.os_patcher = patch('dissect.collection.static.source_code_parser.os')
+        self.os_patcher.start()
 
     def tearDown(self):
-        self.ast_ng_manager_patcher.stop()
+        patch.stopall()
 
     @patch('dissect.collection.static.source_code_parser.sys')
     def test_adds_base_path_to_top_of_python_path(self, sys_mock):
@@ -107,3 +110,17 @@ def test_add_files_returns_false_no_file_was_added(self):
 
         # Assert
         assert_false(result)
+
+    def test_if_only_add_files_that_exist(self):
+        with patch('dissect.collection.static.source_code_parser.os.path.isfile',
+                   new=Mock(side_effect=lambda f: f == 'file-which-exists')):
+            # Given
+            parser = SourceCodeParser('.')
+
+            # When
+            parser.add_files(['file-which-exists', 'file-which-does-not-exist'])
+            parser.parse()
+
+            # Then
+            self.ast_ng_manager_mock.project_from_files.assert_called_once_with(
+                ['file-which-exists'], func_wrapper=ANY)

dissect/consolidation/async_publisher.py

@@ -0,0 +1,54 @@
+#region GPLv3 notice
+# Copyright 2014 Damian Quiroga
+#
+# This file is part of dissect.
+#
+# dissect is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# dissect is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with dissect.  If not, see <http://www.gnu.org/licenses/>.
+#endregion
+
+import threading
+from Queue import Queue
+
+from dissect.consolidation.util.json_to_entity import JsonToEntity
+
+
+class AsyncPublisher(object):
+
+    def __init__(self, callback):
+        self.queue = Queue()
+        self.callback = callback
+        self._stop_sentinel = None
+        self.json_to_entity = JsonToEntity()
+
+    def start(self):
+        self.thread = threading.Thread(target=self._process_queue)
+        self.thread.setDaemon(True)
+        self.thread.start()
+
+    def on_entity(self, entity):
+        self.queue.put(entity)
+
+    def stop(self):
+        self.queue.put_nowait(self._stop_sentinel)
+        self.thread.join()
+
+    def _process_queue(self):
+        while True:
+            try:
+                entity = self.queue.get()
+                if entity is self._stop_sentinel:
+                    return
+                self.callback(entity)
+            except KeyError:
+                pass

dissect/consolidation/test/test_async_publisher.py

@@ -0,0 +1,73 @@
+#region GPLv3 notice
+# Copyright 2014 Damian Quiroga
+#
+# This file is part of dissect.
+#
+# dissect is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# dissect is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with dissect.  If not, see <http://www.gnu.org/licenses/>.
+#endregion
+
+import threading
+
+from dissect.consolidation.async_publisher import AsyncPublisher
+from dissect.test.object_factory import fake
+from mock import Mock, call
+
+
+class TestAsyncPublisher(object):
+    def test_it_invokes_the_callback_for_one_entity(self):
+        # Given
+        callback = Mock()
+        fake_entity = fake('Function')
+        async_publisher = AsyncPublisher(callback)
+        async_publisher.start()
+
+        # When
+        async_publisher.on_entity(fake_entity)
+        async_publisher.stop()
+
+        # Then
+        callback.assert_called_once_with(fake_entity)
+
+    def test_it_invokes_the_callback_for_two_entities(self):
+        # Given
+        callback = Mock()
+        fake_class = fake('Class_')
+        fake_function = fake('Function')
+        async_publisher = AsyncPublisher(callback)
+        async_publisher.start()
+
+        # When
+        async_publisher.on_entity(fake_class)
+        async_publisher.on_entity(fake_function)
+        async_publisher.stop()
+
+        # Then
+        callback.assert_has_calls([call(fake_class),
+                                   call(fake_function)])
+
+    def test_on_entity_returns_before_the_callback_completes(self):
+        # Given
+        finish = threading.Event()
+        callback = Mock()
+        callback.side_effect = lambda entity: finish.wait()
+        fake_entity = fake('Function')
+        async_publisher = AsyncPublisher(callback)
+        async_publisher.start()
+
+        # When
+        async_publisher.on_entity(fake_entity)
+        finish.set()
+
+        # Then
+        async_publisher.stop()