qvm_template_postprocess.py 18 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406
  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 asyncio
  21. import os
  22. import subprocess
  23. import tempfile
  24. from unittest import mock
  25. import qubesadmin.tests
  26. import qubesadmin.tools.qvm_template_postprocess
  27. class QubesLocalMock(qubesadmin.tests.QubesTest):
  28. def __init__(self):
  29. super(QubesLocalMock, self).__init__()
  30. self.__class__ = qubesadmin.app.QubesLocal
  31. qubesd_call = qubesadmin.tests.QubesTest.qubesd_call
  32. run_service = qubesadmin.tests.QubesTest.run_service
  33. class TC_00_qvm_template_postprocess(qubesadmin.tests.QubesTestCase):
  34. def setUp(self):
  35. super(TC_00_qvm_template_postprocess, self).setUp()
  36. self.source_dir = tempfile.TemporaryDirectory()
  37. def tearDown(self):
  38. try:
  39. self.source_dir.cleanup()
  40. except FileNotFoundError:
  41. pass
  42. super(TC_00_qvm_template_postprocess, self).tearDown()
  43. def test_000_import_root_img_raw(self):
  44. root_img = os.path.join(self.source_dir.name, 'root.img')
  45. volume_data = b'volume data'
  46. with open(root_img, 'wb') as f:
  47. f.write(volume_data)
  48. self.app.expected_calls[('dom0', 'admin.vm.List', None, None)] = \
  49. b'0\0test-vm class=TemplateVM state=Halted\n'
  50. self.app.expected_calls[('test-vm', 'admin.vm.volume.List', None,
  51. None)] = \
  52. b'0\0root\nprivate\nvolatile\nkernel\n'
  53. self.app.expected_calls[('test-vm', 'admin.vm.volume.Resize', 'root',
  54. str(len(volume_data)).encode())] = \
  55. b'0\0'
  56. self.app.expected_calls[('test-vm', 'admin.vm.volume.Import', 'root',
  57. volume_data)] = b'0\0'
  58. vm = self.app.domains['test-vm']
  59. qubesadmin.tools.qvm_template_postprocess.import_root_img(
  60. vm, self.source_dir.name)
  61. self.assertAllCalled()
  62. def test_001_import_root_img_tar(self):
  63. root_img = os.path.join(self.source_dir.name, 'root.img')
  64. volume_data = b'volume data' * 1000
  65. with open(root_img, 'wb') as f:
  66. f.write(volume_data)
  67. subprocess.check_call(['tar', 'cf', 'root.img.tar', 'root.img'],
  68. cwd=self.source_dir.name)
  69. subprocess.check_call(['split', '-d', '-b', '1024', 'root.img.tar',
  70. 'root.img.part.'], cwd=self.source_dir.name)
  71. os.unlink(root_img)
  72. self.app.expected_calls[('dom0', 'admin.vm.List', None, None)] = \
  73. b'0\0test-vm class=TemplateVM state=Halted\n'
  74. self.app.expected_calls[('test-vm', 'admin.vm.volume.List', None,
  75. None)] = \
  76. b'0\0root\nprivate\nvolatile\nkernel\n'
  77. self.app.expected_calls[('test-vm', 'admin.vm.volume.Resize', 'root',
  78. str(len(volume_data)).encode())] = \
  79. b'0\0'
  80. self.app.expected_calls[('test-vm', 'admin.vm.volume.Import', 'root',
  81. volume_data)] = b'0\0'
  82. vm = self.app.domains['test-vm']
  83. qubesadmin.tools.qvm_template_postprocess.import_root_img(
  84. vm, self.source_dir.name)
  85. self.assertAllCalled()
  86. def test_002_import_root_img_no_overwrite(self):
  87. self.app.qubesd_connection_type = 'socket'
  88. template_dir = os.path.join(self.source_dir.name, 'vm-templates',
  89. 'test-vm')
  90. os.makedirs(template_dir)
  91. root_img = os.path.join(template_dir, 'root.img')
  92. volume_data = b'volume data'
  93. with open(root_img, 'wb') as f:
  94. f.write(volume_data)
  95. self.app.expected_calls[('dom0', 'admin.vm.List', None, None)] = \
  96. b'0\0test-vm class=TemplateVM state=Halted\n'
  97. self.app.expected_calls[
  98. ('test-vm', 'admin.vm.volume.List', None, None)] = \
  99. b'0\0root\nprivate\nvolatile\nkernel\n'
  100. self.app.expected_calls[
  101. ('test-vm', 'admin.vm.volume.Info', 'root', None)] = \
  102. b'0\0pool=default\nvid=vm-templates/test-vm/root\n'
  103. self.app.expected_calls[
  104. ('dom0', 'admin.pool.List', None, None)] = \
  105. b'0\0default\n'
  106. self.app.expected_calls[
  107. ('dom0', 'admin.pool.Info', 'default', None)] = \
  108. b'0\0driver=file\ndir_path=' + self.source_dir.name.encode() + b'\n'
  109. self.app.expected_calls[
  110. ('test-vm', 'admin.vm.volume.Resize', 'root',
  111. str(len(volume_data)).encode())] = \
  112. b'0\0'
  113. vm = self.app.domains['test-vm']
  114. qubesadmin.tools.qvm_template_postprocess.import_root_img(
  115. vm, template_dir)
  116. self.assertAllCalled()
  117. def test_010_import_appmenus(self):
  118. with open(os.path.join(self.source_dir.name,
  119. 'vm-whitelisted-appmenus.list'), 'w') as f:
  120. f.write('org.gnome.Terminal.desktop\n')
  121. f.write('firefox.desktop\n')
  122. with open(os.path.join(self.source_dir.name,
  123. 'whitelisted-appmenus.list'), 'w') as f:
  124. f.write('org.gnome.Terminal.desktop\n')
  125. f.write('org.gnome.Software.desktop\n')
  126. f.write('gnome-control-center.desktop\n')
  127. self.app.expected_calls[('dom0', 'admin.vm.List', None, None)] = \
  128. b'0\0test-vm class=TemplateVM state=Halted\n'
  129. vm = self.app.domains['test-vm']
  130. with mock.patch('subprocess.check_call') as mock_proc:
  131. qubesadmin.tools.qvm_template_postprocess.import_appmenus(
  132. vm, self.source_dir.name)
  133. self.assertEqual(mock_proc.mock_calls, [
  134. mock.call(['qvm-appmenus',
  135. '--set-default-whitelist=' + os.path.join(self.source_dir.name,
  136. 'vm-whitelisted-appmenus.list'), 'test-vm']),
  137. mock.call(['qvm-appmenus', '--set-whitelist=' + os.path.join(
  138. self.source_dir.name, 'whitelisted-appmenus.list'), 'test-vm']),
  139. ])
  140. self.assertAllCalled()
  141. @mock.patch('grp.getgrnam')
  142. @mock.patch('os.getuid')
  143. def test_011_import_appmenus_as_root(self, mock_getuid, mock_getgrnam):
  144. with open(os.path.join(self.source_dir.name,
  145. 'vm-whitelisted-appmenus.list'), 'w') as f:
  146. f.write('org.gnome.Terminal.desktop\n')
  147. f.write('firefox.desktop\n')
  148. with open(os.path.join(self.source_dir.name,
  149. 'whitelisted-appmenus.list'), 'w') as f:
  150. f.write('org.gnome.Terminal.desktop\n')
  151. f.write('org.gnome.Software.desktop\n')
  152. f.write('gnome-control-center.desktop\n')
  153. self.app.expected_calls[('dom0', 'admin.vm.List', None, None)] = \
  154. b'0\0test-vm class=TemplateVM state=Halted\n'
  155. mock_getuid.return_value = 0
  156. mock_getgrnam.configure_mock(**{
  157. 'return_value.gr_mem.__getitem__.return_value': 'user'
  158. })
  159. vm = self.app.domains['test-vm']
  160. with mock.patch('subprocess.check_call') as mock_proc:
  161. qubesadmin.tools.qvm_template_postprocess.import_appmenus(
  162. vm, self.source_dir.name)
  163. self.assertEqual(mock_proc.mock_calls, [
  164. mock.call(['runuser', '-u', 'user', '--', 'env', 'DISPLAY=:0',
  165. 'qvm-appmenus',
  166. '--set-default-whitelist=' + os.path.join(self.source_dir.name,
  167. 'vm-whitelisted-appmenus.list'), 'test-vm']),
  168. mock.call(['runuser', '-u', 'user', '--', 'env', 'DISPLAY=:0',
  169. 'qvm-appmenus', '--set-whitelist=' + os.path.join(
  170. self.source_dir.name, 'whitelisted-appmenus.list'), 'test-vm']),
  171. ])
  172. self.assertAllCalled()
  173. @mock.patch('grp.getgrnam')
  174. @mock.patch('os.getuid')
  175. def test_012_import_appmenus_missing_user(self, mock_getuid, mock_getgrnam):
  176. with open(os.path.join(self.source_dir.name,
  177. 'vm-whitelisted-appmenus.list'), 'w') as f:
  178. f.write('org.gnome.Terminal.desktop\n')
  179. f.write('firefox.desktop\n')
  180. with open(os.path.join(self.source_dir.name,
  181. 'whitelisted-appmenus.list'), 'w') as f:
  182. f.write('org.gnome.Terminal.desktop\n')
  183. f.write('org.gnome.Software.desktop\n')
  184. f.write('gnome-control-center.desktop\n')
  185. self.app.expected_calls[('dom0', 'admin.vm.List', None, None)] = \
  186. b'0\0test-vm class=TemplateVM state=Halted\n'
  187. mock_getuid.return_value = 0
  188. mock_getgrnam.side_effect = KeyError
  189. vm = self.app.domains['test-vm']
  190. with mock.patch('subprocess.check_call') as mock_proc:
  191. qubesadmin.tools.qvm_template_postprocess.import_appmenus(
  192. vm, self.source_dir.name)
  193. self.assertEqual(mock_proc.mock_calls, [])
  194. self.assertAllCalled()
  195. def add_new_vm_side_effect(self, *args, **kwargs):
  196. self.app.expected_calls[('dom0', 'admin.vm.List', None, None)] = \
  197. b'0\0test-vm class=TemplateVM state=Halted\n'
  198. self.app.domains.clear_cache()
  199. return self.app.domains['test-vm']
  200. @asyncio.coroutine
  201. def wait_for_shutdown(self, vm):
  202. pass
  203. @mock.patch('qubesadmin.tools.qvm_template_postprocess.import_appmenus')
  204. @mock.patch('qubesadmin.tools.qvm_template_postprocess.import_root_img')
  205. def test_020_post_install(self, mock_import_root_img,
  206. mock_import_appmenus):
  207. self.app.expected_calls[('dom0', 'admin.vm.List', None, None)] = \
  208. b'0\0'
  209. self.app.add_new_vm = mock.Mock(side_effect=self.add_new_vm_side_effect)
  210. self.app.expected_calls[
  211. ('test-vm', 'admin.vm.property.Set', 'netvm', b'')] = b'0\0'
  212. self.app.expected_calls[
  213. ('test-vm', 'admin.vm.property.Reset', 'netvm', None)] = b'0\0'
  214. self.app.expected_calls[
  215. ('test-vm', 'admin.vm.feature.Set', 'qrexec', b'1')] = b'0\0'
  216. self.app.expected_calls[
  217. ('test-vm', 'admin.vm.Start', None, None)] = b'0\0'
  218. self.app.expected_calls[
  219. ('test-vm', 'admin.vm.Shutdown', None, None)] = b'0\0'
  220. if qubesadmin.tools.qvm_template_postprocess.have_events:
  221. patch_domain_shutdown = mock.patch(
  222. 'qubesadmin.events.utils.wait_for_domain_shutdown')
  223. self.addCleanup(patch_domain_shutdown.stop)
  224. mock_domain_shutdown = patch_domain_shutdown.start()
  225. mock_domain_shutdown.side_effect = self.wait_for_shutdown
  226. else:
  227. self.app.expected_calls[
  228. ('test-vm', 'admin.vm.List', None, None)] = \
  229. b'0\0test-vm class=TemplateVM state=Halted\n'
  230. asyncio.set_event_loop(asyncio.new_event_loop())
  231. ret = qubesadmin.tools.qvm_template_postprocess.main([
  232. '--really', 'post-install', 'test-vm', self.source_dir.name],
  233. app=self.app)
  234. self.assertEqual(ret, 0)
  235. self.app.add_new_vm.assert_called_once_with('TemplateVM',
  236. name='test-vm', label='black')
  237. mock_import_root_img.assert_called_once_with(self.app.domains[
  238. 'test-vm'], self.source_dir.name)
  239. mock_import_appmenus.assert_called_once_with(self.app.domains[
  240. 'test-vm'], self.source_dir.name)
  241. if qubesadmin.tools.qvm_template_postprocess.have_events:
  242. mock_domain_shutdown.assert_called_once_with([self.app.domains[
  243. 'test-vm']])
  244. self.assertEqual(self.app.service_calls, [
  245. ('test-vm', 'qubes.PostInstall', {}),
  246. ('test-vm', 'qubes.PostInstall', b''),
  247. ])
  248. self.assertAllCalled()
  249. @mock.patch('qubesadmin.tools.qvm_template_postprocess.import_appmenus')
  250. @mock.patch('qubesadmin.tools.qvm_template_postprocess.import_root_img')
  251. def test_021_post_install_reinstall(self, mock_import_root_img,
  252. mock_import_appmenus):
  253. self.app.expected_calls[('dom0', 'admin.vm.List', None, None)] = \
  254. b'0\0test-vm class=TemplateVM state=Halted\n'
  255. self.app.add_new_vm = mock.Mock()
  256. self.app.expected_calls[
  257. ('test-vm', 'admin.vm.property.Set', 'netvm', b'')] = b'0\0'
  258. self.app.expected_calls[
  259. ('test-vm', 'admin.vm.property.Reset', 'netvm', None)] = b'0\0'
  260. self.app.expected_calls[
  261. ('test-vm', 'admin.vm.feature.Set', 'qrexec', b'1')] = b'0\0'
  262. self.app.expected_calls[
  263. ('test-vm', 'admin.vm.Start', None, None)] = b'0\0'
  264. self.app.expected_calls[
  265. ('test-vm', 'admin.vm.Shutdown', None, None)] = b'0\0'
  266. if qubesadmin.tools.qvm_template_postprocess.have_events:
  267. patch_domain_shutdown = mock.patch(
  268. 'qubesadmin.events.utils.wait_for_domain_shutdown')
  269. self.addCleanup(patch_domain_shutdown.stop)
  270. mock_domain_shutdown = patch_domain_shutdown.start()
  271. mock_domain_shutdown.side_effect = self.wait_for_shutdown
  272. else:
  273. self.app.expected_calls[
  274. ('test-vm', 'admin.vm.List', None, None)] = \
  275. b'0\0test-vm class=TemplateVM state=Halted\n'
  276. asyncio.set_event_loop(asyncio.new_event_loop())
  277. ret = qubesadmin.tools.qvm_template_postprocess.main([
  278. '--really', 'post-install', 'test-vm', self.source_dir.name],
  279. app=self.app)
  280. self.assertEqual(ret, 0)
  281. self.assertFalse(self.app.add_new_vm.called)
  282. mock_import_root_img.assert_called_once_with(self.app.domains[
  283. 'test-vm'], self.source_dir.name)
  284. mock_import_appmenus.assert_called_once_with(self.app.domains[
  285. 'test-vm'], self.source_dir.name)
  286. if qubesadmin.tools.qvm_template_postprocess.have_events:
  287. mock_domain_shutdown.assert_called_once_with([self.app.domains[
  288. 'test-vm']])
  289. self.assertEqual(self.app.service_calls, [
  290. ('test-vm', 'qubes.PostInstall', {}),
  291. ('test-vm', 'qubes.PostInstall', b''),
  292. ])
  293. self.assertAllCalled()
  294. @mock.patch('qubesadmin.tools.qvm_template_postprocess.import_appmenus')
  295. @mock.patch('qubesadmin.tools.qvm_template_postprocess.import_root_img')
  296. def test_022_post_install_skip_start(self, mock_import_root_img,
  297. mock_import_appmenus):
  298. self.app.expected_calls[('dom0', 'admin.vm.List', None, None)] = \
  299. b'0\0test-vm class=TemplateVM state=Halted\n'
  300. self.app.add_new_vm = mock.Mock()
  301. if qubesadmin.tools.qvm_template_postprocess.have_events:
  302. patch_domain_shutdown = mock.patch(
  303. 'qubesadmin.events.utils.wait_for_domain_shutdown')
  304. self.addCleanup(patch_domain_shutdown.stop)
  305. mock_domain_shutdown = patch_domain_shutdown.start()
  306. mock_domain_shutdown.side_effect = self.wait_for_shutdown
  307. asyncio.set_event_loop(asyncio.new_event_loop())
  308. ret = qubesadmin.tools.qvm_template_postprocess.main([
  309. '--really', '--skip-start', 'post-install', 'test-vm',
  310. self.source_dir.name],
  311. app=self.app)
  312. self.assertEqual(ret, 0)
  313. self.assertFalse(self.app.add_new_vm.called)
  314. mock_import_root_img.assert_called_once_with(self.app.domains[
  315. 'test-vm'], self.source_dir.name)
  316. mock_import_appmenus.assert_called_once_with(self.app.domains[
  317. 'test-vm'], self.source_dir.name)
  318. if qubesadmin.tools.qvm_template_postprocess.have_events:
  319. self.assertFalse(mock_domain_shutdown.called)
  320. self.assertEqual(self.app.service_calls, [])
  321. self.assertAllCalled()
  322. def test_030_pre_remove(self):
  323. self.app.expected_calls[('dom0', 'admin.vm.List', None, None)] = \
  324. b'0\0test-vm class=TemplateVM state=Halted\n'
  325. self.app.expected_calls[('test-vm', 'admin.vm.Remove', None, None)] = \
  326. b'0\0'
  327. self.app.expected_calls[
  328. ('test-vm', 'admin.vm.property.Get', 'template', None)] = \
  329. b'2\0QubesNoSuchPropertyError\0\0invalid property \'template\' of ' \
  330. b'test-vm\0'
  331. ret = qubesadmin.tools.qvm_template_postprocess.main([
  332. '--really', 'pre-remove', 'test-vm',
  333. self.source_dir.name],
  334. app=self.app)
  335. self.assertEqual(ret, 0)
  336. self.assertEqual(self.app.service_calls, [])
  337. self.assertAllCalled()
  338. def test_031_pre_remove_existing_appvm(self):
  339. self.app.expected_calls[('dom0', 'admin.vm.List', None, None)] = \
  340. b'0\0test-vm class=TemplateVM state=Halted\n' \
  341. b'test-vm2 class=AppVM state=Halted\n'
  342. self.app.expected_calls[
  343. ('test-vm', 'admin.vm.property.Get', 'template', None)] = \
  344. b'2\0QubesNoSuchPropertyError\0\0invalid property \'template\' of ' \
  345. b'test-vm\0'
  346. self.app.expected_calls[
  347. ('test-vm2', 'admin.vm.property.Get', 'template', None)] = \
  348. b'0\0default=no type=vm test-vm'
  349. with self.assertRaises(SystemExit):
  350. qubesadmin.tools.qvm_template_postprocess.main([
  351. '--really', 'pre-remove', 'test-vm',
  352. self.source_dir.name],
  353. app=self.app)
  354. self.assertEqual(self.app.service_calls, [])
  355. self.assertAllCalled()
  356. def test_040_missing_really(self):
  357. with self.assertRaises(SystemExit):
  358. qubesadmin.tools.qvm_template_postprocess.main([
  359. 'post-install', 'test-vm', self.source_dir.name],
  360. app=self.app)
  361. self.assertAllCalled()