def setUp(self):
config = {
'hooks': [{
'name': 'see.test.environment_test.TestHook'
}, {
'name': 'see.test.environment_test.WrongHook'
}]
}
self.environment = Environment(context_factory, config)
self.environment.allocate()
def main():
arguments = parse_arguments()
context_factory = QEMUContextFactory(arguments.context)
with Environment(context_factory, arguments.hooks) as environment:
protocol(environment.context, arguments.sample, arguments.command)
def main(args):
vm_name = args['<vm_name>']
uri = args['--connection']
debug = args['--debug']
hooks_config_path = args['<plugins_configuration>']
init_logger(debug)
hooks_config = {}
with open(hooks_config_path) as f:
hooks_config = json.load(f)
logging.info('Connect to Neo4j DB')
graph = Graph(password=DB_PASSWORD)
if 'configuration' not in hooks_config:
hooks_config['configuration'] = {}
# use default desktop ready delay if unset
if "desktop_ready_delay" not in hooks_config['configuration']:
hooks_config['configuration'] = DESKTOP_READY_DELAY
# insert graph object into general hook configuration
hooks_config['configuration']['graph'] = graph
# insert vm_name object
hooks_config['configuration']['domain_name'] = vm_name
# insert debug flag
hooks_config['configuration']['debug'] = debug
# delete entire graph ?
try:
delete = hooks_config['configuration']['delete']
except KeyError:
pass
else:
if delete:
logging.info("Deleting all nodes in graph database")
graph.delete_all()
# reset GraphQL IDL
graph.run("CALL graphql.idl(null)")
# replace existing OS ?
os_match = OS.match(graph).where("_.name = '{}'".format(vm_name))
try:
replace = hooks_config['configuration']['replace']
except KeyError:
# assume replace = False
if os_match.first():
logging.info('OS already inserted, exiting')
return
else:
if not replace and os_match.first():
logging.info('OS already inserted, exiting')
return
elif os_match.first():
# replace = True and an OS already exists
logging.info('Deleting previous OS')
graph.run(SUBGRAPH_DELETE_OS.format(vm_name))
with QEMUDomainContextFactory(vm_name, uri) as context:
with Environment(context, hooks_config) as environment:
logging.info('Capturing %s', vm_name)
protocol(environment)
def test_no_context(self):
"""RuntimeError raised if the Environment has not been initialized."""
environment = Environment(context_factory, {})
with self.assertRaises(RuntimeError):
context = environment.context
context.poweron()
def test_cleanup_context_error(self):
"""Errors in context cleanup are handled."""
config = {'hooks': [{'name': 'see.test.environment_test.WrongHook'}]}
with Environment(wrong_context_factory, config) as environment:
context = environment.context
self.assertTrue(context.clean)
class EnvironmentTest(unittest.TestCase):
def setUp(self):
config = {
'hooks': [{
'name': 'see.test.environment_test.TestHook'
}, {
'name': 'see.test.environment_test.WrongHook'
}]
}
self.environment = Environment(context_factory, config)
self.environment.allocate()
def tearDown(self):
self.environment.deallocate()
def test_trigger_simple_event(self):
"""Environment's events are propagated to Hooks."""
self.environment.context.trigger('event1')
self.assertTrue(self.environment._hookmanager.hooks[0].called1)
def test_trigger_event_async_handler(self):
"""Environment's events are propagated to Hooks."""
self.environment.context.trigger('event1')
hook = self.environment._hookmanager.hooks[0]
hook.async_called.wait()
self.assertTrue(hook.async_called.is_set())
def test_trigger_complex_event(self):
"""Environment's events attributes are propagated to Hooks."""
self.environment.context.trigger('event2', arg='foo')
self.assertEqual(self.environment._hookmanager.hooks[0].event_arg,
'foo')
def test_sync_cascade_event(self):
"""Events which synchronous handlers trigger other ones."""
self.environment.context.trigger('cascade_sync')
hook = self.environment._hookmanager.hooks[0]
hook.async_called.wait()
self.assertTrue(hook.async_called.is_set())
self.assertTrue(self.environment._hookmanager.hooks[0].called1)
def test_async_cascade_event(self):
"""Events which asynchronous handlers trigger other ones."""
self.environment.context.trigger('cascade_async')
hook = self.environment._hookmanager.hooks[0]
hook.async_called.wait()
self.assertTrue(hook.async_called.is_set())
self.assertTrue(self.environment._hookmanager.hooks[0].called1)
def test_trigger_handler_error(self):
"""Exceptions within synchronous handlers won't stop the execution."""
self.environment.context.trigger('error_event')
self.assertTrue(
self.environment._hookmanager.hooks[1].error_handler_called)
def test_trigger_async_handler_error(self):
"""Exceptions within asynchronous handlers won't stop the execution."""
self.environment.context.trigger('async_error_event')
hook = self.environment._hookmanager.hooks[1]
hook.error_async_handler_called.wait()
self.assertTrue(self.environment._hookmanager.hooks[1].
error_async_handler_called.is_set())
def test_cleanup(self):
"""Environment's hooks are cleaned up on deallocate."""
hook = self.environment._hookmanager.hooks[0]
self.environment.deallocate()
self.assertTrue(hook.clean)
def test_cleanup_hook_error(self):
"""Errors in hooks cleanup are handled."""
hook = self.environment._hookmanager.hooks[1]
self.environment.deallocate()
self.assertTrue(hook.clean)
def test_cleanup_context_error(self):
"""Errors in context cleanup are handled."""
config = {'hooks': [{'name': 'see.test.environment_test.WrongHook'}]}
with Environment(wrong_context_factory, config) as environment:
context = environment.context
self.assertTrue(context.clean)
def test_no_context(self):
"""RuntimeError raised if the Environment has not been initialized."""
environment = Environment(context_factory, {})
with self.assertRaises(RuntimeError):
context = environment.context
context.poweron()
def test_unsubscribe(self):
"""Handler is unsubscribed."""
hook = self.environment._hookmanager.hooks[0]
hook.context.unsubscribe('event1', hook.handler1)
self.environment.context.trigger('event1')
self.assertFalse(hook.called1)
def test_unsubscribe_async(self):
"""Async handler is unsubscribed."""
hook = self.environment._hookmanager.hooks[0]
hook.context.unsubscribe('event1', hook.async_handler)
self.environment.context.trigger('event1')
time.sleep(0.1)
self.assertFalse(hook.async_called.is_set())
def test_unsubscribe_no_such_event(self):
"""Error is raised if the given event is not registered."""
hook = self.environment._hookmanager.hooks[0]
with self.assertRaises(ValueError):
hook.context.unsubscribe('event42', hook.handler1)
def test_unsubscribe_no_such_handler(self):
"""Error is raised if the given handler is not registered."""
hook = self.environment._hookmanager.hooks[0]
with self.assertRaises(ValueError):
hook.context.unsubscribe('event1', hook.cleanup)
def setUp(self):
config = {'hooks': [{'name': 'see.test.environment_test.TestHook'},
{'name': 'see.test.environment_test.WrongHook'}]}
self.environment = Environment(context_factory, config)
self.environment.allocate()
class EnvironmentTest(unittest.TestCase):
def setUp(self):
config = {'hooks': [{'name': 'see.test.environment_test.TestHook'},
{'name': 'see.test.environment_test.WrongHook'}]}
self.environment = Environment(context_factory, config)
self.environment.allocate()
def tearDown(self):
self.environment.deallocate()
def test_trigger_simple_event(self):
"""Environment's events are propagated to Hooks."""
self.environment.context.trigger('event1')
self.assertTrue(self.environment._hookmanager.hooks[0].called1)
def test_trigger_event_async_handler(self):
"""Environment's events are propagated to Hooks."""
self.environment.context.trigger('event1')
hook = self.environment._hookmanager.hooks[0]
hook.async_called.wait()
self.assertTrue(hook.async_called.is_set())
def test_trigger_complex_event(self):
"""Environment's events attributes are propagated to Hooks."""
self.environment.context.trigger('event2', arg='foo')
self.assertEqual(self.environment._hookmanager.hooks[0].event_arg,
'foo')
def test_sync_cascade_event(self):
"""Events which synchronous handlers trigger other ones."""
self.environment.context.trigger('cascade_sync')
hook = self.environment._hookmanager.hooks[0]
hook.async_called.wait()
self.assertTrue(hook.async_called.is_set())
self.assertTrue(self.environment._hookmanager.hooks[0].called1)
def test_async_cascade_event(self):
"""Events which asynchronous handlers trigger other ones."""
self.environment.context.trigger('cascade_async')
hook = self.environment._hookmanager.hooks[0]
hook.async_called.wait()
self.assertTrue(hook.async_called.is_set())
self.assertTrue(self.environment._hookmanager.hooks[0].called1)
def test_trigger_handler_error(self):
"""Exceptions within synchronous handlers won't stop the execution."""
self.environment.context.trigger('error_event')
self.assertTrue(self.environment._hookmanager.hooks[1].error_handler_called)
def test_trigger_async_handler_error(self):
"""Exceptions within asynchronous handlers won't stop the execution."""
self.environment.context.trigger('async_error_event')
hook = self.environment._hookmanager.hooks[1]
hook.error_async_handler_called.wait()
self.assertTrue(self.environment._hookmanager.hooks[1].error_async_handler_called.is_set())
def test_cleanup(self):
"""Environment's hooks are cleaned up on deallocate."""
hook = self.environment._hookmanager.hooks[0]
self.environment.deallocate()
self.assertTrue(hook.clean)
def test_cleanup_hook_error(self):
"""Errors in hooks cleanup are handled."""
hook = self.environment._hookmanager.hooks[1]
self.environment.deallocate()
self.assertTrue(hook.clean)
def test_cleanup_context_error(self):
"""Errors in context cleanup are handled."""
config = {'hooks': [{'name': 'see.test.environment_test.WrongHook'}]}
with Environment(wrong_context_factory, config) as environment:
context = environment.context
self.assertTrue(context.clean)
def test_no_context(self):
"""RuntimeError raised if the Environment has not been initialized."""
environment = Environment(context_factory, {})
with self.assertRaises(RuntimeError):
context = environment.context
context.poweron()
def test_unsubscribe(self):
"""Handler is unsubscribed."""
hook = self.environment._hookmanager.hooks[0]
hook.context.unsubscribe('event1', hook.handler1)
self.environment.context.trigger('event1')
self.assertFalse(hook.called1)
def test_unsubscribe_async(self):
"""Async handler is unsubscribed."""
hook = self.environment._hookmanager.hooks[0]
hook.context.unsubscribe('event1', hook.async_handler)
self.environment.context.trigger('event1')
time.sleep(0.1)
self.assertFalse(hook.async_called.is_set())
def test_unsubscribe_no_such_event(self):
"""Error is raised if the given event is not registered."""
hook = self.environment._hookmanager.hooks[0]
with self.assertRaises(ValueError):
hook.context.unsubscribe('event42', hook.handler1)
def test_unsubscribe_no_such_handler(self):
"""Error is raised if the given handler is not registered."""
hook = self.environment._hookmanager.hooks[0]
with self.assertRaises(ValueError):
hook.context.unsubscribe('event1', hook.cleanup)
def capture_main(args):
vm_name = args['<vm_name>']
uri = args['--connection']
debug = args['--debug']
hooks_config_path = args['<plugins_configuration>']
init_logger(debug)
# load hooks.json
hooks_config = {}
with open(hooks_config_path) as f:
try:
hooks_config = json.load(f)
except json.JSONDecodeError:
logging.error("Failed to parse %s: Invalid JSON",
hooks_config_path)
return
if 'configuration' not in hooks_config:
hooks_config['configuration'] = {}
# use default desktop ready delay if unset
if "desktop_ready_delay" not in hooks_config['configuration']:
hooks_config['configuration'][
'desktop_ready_delay'] = DESKTOP_READY_DELAY
# insert vm_name object
hooks_config['configuration']['domain_name'] = vm_name
# insert debug flag
hooks_config['configuration']['debug'] = debug
# Neo4j required ?
neo4j = hooks_config['configuration'].get('neo4j', {})
if neo4j.get('enabled'):
logging.info('Connect to Neo4j DB')
graph = Graph(password=DB_PASSWORD)
# handle 'delete' key
# delete entire graph ?
delete = neo4j.get('delete')
if delete:
logging.info("Deleting all nodes in graph database")
graph.delete_all()
# handle 'replace' key
# replace existing OS in Neo4j ?
replace = neo4j.get('replace', False)
os_match = OS.match(graph).where("_.name = '{}'".format(vm_name))
if not replace and os_match.first():
logging.info('OS %s already inserted, exiting', vm_name)
return
elif os_match.first():
# replace = True and an OS already exists
logging.info('Deleting previous OS %s', vm_name)
graph.run(SUBGRAPH_DELETE_OS.format(vm_name))
# init new OS node
os_node = OS(vm_name)
# create it already, so transactions will work in hooks
logging.info('Creating OS node %s', os_node.name)
graph.create(os_node)
neo4j['OS'] = os_node
neo4j['graph'] = graph
# Run the protocol
try:
with QEMUDomainContextFactory(vm_name, uri) as context:
with Environment(context, hooks_config) as environment:
logging.info('Capturing %s', vm_name)
start = timer()
protocol(environment)
end = timer()
delta = timedelta(seconds=end - start)
# remove microseconds
duration = str(delta).split('.')[0]
logging.info('Capture duration: %s', duration)
except KeyboardInterrupt:
# cleanup
if neo4j.get('enabled'):
logging.warning("SIGINT received")
logging.info("cleanup: removing OS node")
graph: Graph = neo4j.get('graph', None)
os_node: OS = neo4j.get('OS', None)
if graph is not None and os_node is not None:
graph.run(SUBGRAPH_DELETE_OS.format(vm_name))
if neo4j.get('enabled'):
# push OS node updates
os_node: OS = neo4j['OS']
graph: Graph = neo4j['graph']
graph.push(os_node)
