events.py 9.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237
  1. # -*- encoding: utf8 -*-
  2. #
  3. # The Qubes OS Project, http://www.qubes-os.org
  4. #
  5. # Copyright (C) 2017 Marek Marczykowski-Górecki
  6. # <marmarek@invisiblethingslab.com>
  7. #
  8. # This program is free software; you can redistribute it and/or modify
  9. # it under the terms of the GNU Lesser General Public License as published by
  10. # the Free Software Foundation; either version 2.1 of the License, or
  11. # (at your option) any later version.
  12. #
  13. # This program is distributed in the hope that it will be useful,
  14. # but WITHOUT ANY WARRANTY; without even the implied warranty of
  15. # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  16. # GNU Lesser General Public License for more details.
  17. #
  18. # You should have received a copy of the GNU Lesser General Public License along
  19. # with this program; if not, see <http://www.gnu.org/licenses/>.
  20. import socket
  21. import subprocess
  22. import qubesadmin.tests
  23. import unittest
  24. try:
  25. # qubesadmin.events require python3, so this tests can also use python3 features
  26. import asyncio
  27. import unittest.mock
  28. import qubesadmin.events
  29. except ImportError:
  30. # don't run any tests on python2
  31. def load_tests(loader, tests, pattern):
  32. return unittest.TestSuite()
  33. # don't fail on coroutine decorator
  34. class asyncio(object):
  35. @staticmethod
  36. def coroutine(f):
  37. return f
  38. class TC_00_Events(qubesadmin.tests.QubesTestCase):
  39. def setUp(self):
  40. super().setUp()
  41. self.app = unittest.mock.MagicMock()
  42. self.dispatcher = qubesadmin.events.EventsDispatcher(self.app)
  43. def test_000_handler_specific(self):
  44. handler = unittest.mock.Mock()
  45. self.dispatcher.add_handler('some-event', handler)
  46. self.dispatcher.handle('', 'some-event', arg1='value1')
  47. handler.assert_called_once_with(None, 'some-event', arg1='value1')
  48. handler.reset_mock()
  49. self.dispatcher.handle('test-vm', 'some-event', arg1='value1')
  50. handler.assert_called_once_with(
  51. self.app.domains['test-vm'], 'some-event', arg1='value1')
  52. handler.reset_mock()
  53. self.dispatcher.handle('', 'other-event', arg1='value1')
  54. self.assertFalse(handler.called)
  55. self.dispatcher.remove_handler('some-event', handler)
  56. self.dispatcher.handle('', 'some-event', arg1='value1')
  57. self.assertFalse(handler.called)
  58. def test_001_handler_glob(self):
  59. handler = unittest.mock.Mock()
  60. self.dispatcher.add_handler('*', handler)
  61. self.dispatcher.handle('', 'some-event', arg1='value1')
  62. handler.assert_called_once_with(None, 'some-event', arg1='value1')
  63. handler.reset_mock()
  64. self.dispatcher.handle('test-vm', 'some-event', arg1='value1')
  65. handler.assert_called_once_with(
  66. self.app.domains['test-vm'], 'some-event', arg1='value1')
  67. handler.reset_mock()
  68. self.dispatcher.handle('', 'other-event', arg1='value1')
  69. handler.assert_called_once_with(None, 'other-event', arg1='value1')
  70. handler.reset_mock()
  71. self.dispatcher.remove_handler('*', handler)
  72. self.dispatcher.handle('', 'some-event', arg1='value1')
  73. self.assertFalse(handler.called)
  74. @asyncio.coroutine
  75. def mock_get_events_reader(self, stream, cleanup_func, expected_vm,
  76. vm=None):
  77. self.assertEqual(expected_vm, vm)
  78. return stream, cleanup_func
  79. @asyncio.coroutine
  80. def send_events(self, stream, events):
  81. for event in events:
  82. stream.feed_data(event)
  83. # don't use yield from...
  84. for x in asyncio.sleep(0.01):
  85. yield x
  86. stream.feed_eof()
  87. def test_010_listen_for_events(self):
  88. loop = asyncio.new_event_loop()
  89. asyncio.set_event_loop(loop)
  90. stream = asyncio.StreamReader()
  91. cleanup_func = unittest.mock.Mock()
  92. self.dispatcher._get_events_reader = \
  93. lambda vm: self.mock_get_events_reader(stream, cleanup_func,
  94. None, vm)
  95. handler = unittest.mock.Mock()
  96. self.dispatcher.add_handler('some-event', handler)
  97. events = [
  98. b'1\0\0some-event\0arg1\0value1\0\0',
  99. b'1\0some-vm\0some-event\0arg1\0value1\0\0',
  100. b'1\0some-vm\0some-event\0arg_without_value\0\0arg2\0value\0\0',
  101. b'1\0some-vm\0other-event\0\0',
  102. ]
  103. asyncio.ensure_future(self.send_events(stream, events))
  104. loop.run_until_complete(self.dispatcher.listen_for_events(
  105. reconnect=False))
  106. self.assertEqual(handler.mock_calls, [
  107. unittest.mock.call(None, 'some-event', arg1='value1'),
  108. unittest.mock.call(
  109. self.app.domains['some-vm'], 'some-event', arg1='value1'),
  110. unittest.mock.call(
  111. self.app.domains['some-vm'], 'some-event',
  112. arg_without_value='', arg2='value'),
  113. ])
  114. cleanup_func.assert_called_once_with()
  115. loop.close()
  116. def mock_open_unix_connection(self, expected_path, sock, path):
  117. self.assertEqual(expected_path, path)
  118. return asyncio.open_connection(sock=sock)
  119. def read_all(self, sock):
  120. buf = b''
  121. for data in iter(lambda: sock.recv(4096), b''):
  122. buf += data
  123. return buf
  124. def test_020_get_events_reader_local(self):
  125. self.app.qubesd_connection_type = 'socket'
  126. loop = asyncio.new_event_loop()
  127. asyncio.set_event_loop(loop)
  128. sock1, sock2 = socket.socketpair()
  129. with unittest.mock.patch('asyncio.open_unix_connection',
  130. lambda path: self.mock_open_unix_connection(
  131. qubesadmin.config.QUBESD_SOCKET, sock1, path)):
  132. task = asyncio.ensure_future(self.dispatcher._get_events_reader())
  133. reader = asyncio.ensure_future(loop.run_in_executor(None,
  134. self.read_all, sock2))
  135. loop.run_until_complete(asyncio.wait([task, reader]))
  136. self.assertEqual(reader.result(),
  137. b'dom0\0admin.Events\0dom0\0\0')
  138. self.assertIsInstance(task.result()[0], asyncio.StreamReader)
  139. cleanup_func = task.result()[1]
  140. cleanup_func()
  141. sock2.close()
  142. # run socket cleanup functions
  143. loop.stop()
  144. loop.run_forever()
  145. loop.close()
  146. def test_021_get_events_reader_local_vm(self):
  147. self.app.qubesd_connection_type = 'socket'
  148. loop = asyncio.new_event_loop()
  149. asyncio.set_event_loop(loop)
  150. sock1, sock2 = socket.socketpair()
  151. vm = unittest.mock.Mock()
  152. vm.name = 'test-vm'
  153. with unittest.mock.patch('asyncio.open_unix_connection',
  154. lambda path: self.mock_open_unix_connection(
  155. qubesadmin.config.QUBESD_SOCKET, sock1, path)):
  156. task = asyncio.ensure_future(self.dispatcher._get_events_reader(vm))
  157. reader = asyncio.ensure_future(loop.run_in_executor(None,
  158. self.read_all, sock2))
  159. loop.run_until_complete(asyncio.wait([task, reader]))
  160. self.assertEqual(reader.result(),
  161. b'dom0\0admin.Events\0test-vm\0\0')
  162. self.assertIsInstance(task.result()[0], asyncio.StreamReader)
  163. cleanup_func = task.result()[1]
  164. cleanup_func()
  165. sock2.close()
  166. # run socket cleanup functions
  167. loop.stop()
  168. loop.run_forever()
  169. loop.close()
  170. @asyncio.coroutine
  171. def mock_coroutine(self, mock, *args, **kwargs):
  172. return mock(*args, **kwargs)
  173. def test_022_get_events_reader_remote(self):
  174. self.app.qubesd_connection_type = 'qrexec'
  175. loop = asyncio.new_event_loop()
  176. asyncio.set_event_loop(loop)
  177. mock_proc = unittest.mock.Mock()
  178. with unittest.mock.patch('asyncio.create_subprocess_exec',
  179. lambda *args, **kwargs: self.mock_coroutine(mock_proc,
  180. *args, **kwargs)):
  181. task = asyncio.ensure_future(self.dispatcher._get_events_reader())
  182. loop.run_until_complete(task)
  183. self.assertEqual(mock_proc.mock_calls, [
  184. unittest.mock.call('qrexec-client-vm', 'dom0',
  185. 'admin.Events', stdin=subprocess.PIPE,
  186. stdout=subprocess.PIPE),
  187. unittest.mock.call().stdin.write_eof()
  188. ])
  189. self.assertEqual(task.result()[0], mock_proc().stdout)
  190. cleanup_func = task.result()[1]
  191. cleanup_func()
  192. unittest.mock.call().kill.assert_called_once_with()
  193. loop.close()
  194. def test_023_get_events_reader_remote_vm(self):
  195. self.app.qubesd_connection_type = 'qrexec'
  196. loop = asyncio.new_event_loop()
  197. asyncio.set_event_loop(loop)
  198. mock_proc = unittest.mock.Mock()
  199. vm = unittest.mock.Mock()
  200. vm.name = 'test-vm'
  201. with unittest.mock.patch('asyncio.create_subprocess_exec',
  202. lambda *args, **kwargs: self.mock_coroutine(mock_proc,
  203. *args, **kwargs)):
  204. task = asyncio.ensure_future(self.dispatcher._get_events_reader(vm))
  205. loop.run_until_complete(task)
  206. self.assertEqual(mock_proc.mock_calls, [
  207. unittest.mock.call('qrexec-client-vm', 'test-vm',
  208. 'admin.Events', stdin=subprocess.PIPE,
  209. stdout=subprocess.PIPE),
  210. unittest.mock.call().stdin.write_eof()
  211. ])
  212. self.assertEqual(task.result()[0], mock_proc().stdout)
  213. cleanup_func = task.result()[1]
  214. cleanup_func()
  215. unittest.mock.call().kill.assert_called_once_with()
  216. loop.close()