extra.py 8.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245
  1. #
  2. # The Qubes OS Project, https://www.qubes-os.org/
  3. #
  4. # Copyright (C) 2016
  5. # Marek Marczykowski-Górecki <marmarek@invisiblethingslab.com>
  6. #
  7. # This library is free software; you can redistribute it and/or
  8. # modify it under the terms of the GNU Lesser General Public
  9. # License as published by the Free Software Foundation; either
  10. # version 2.1 of the License, or (at your option) any later version.
  11. #
  12. # This library is distributed in the hope that it will be useful,
  13. # but WITHOUT ANY WARRANTY; without even the implied warranty of
  14. # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
  15. # Lesser General Public License for more details.
  16. #
  17. # You should have received a copy of the GNU Lesser General Public
  18. # License along with this library; if not, see <https://www.gnu.org/licenses/>.
  19. #
  20. import asyncio
  21. import os
  22. import subprocess
  23. import sys
  24. import pkg_resources
  25. import qubes.tests
  26. import qubes.vm.appvm
  27. class ProcessWrapper(object):
  28. def __init__(self, proc, loop=None):
  29. self._proc = proc
  30. self._loop = loop or asyncio.get_event_loop()
  31. def __getattr__(self, item):
  32. return getattr(self._proc, item)
  33. def __setattr__(self, key, value):
  34. if key.startswith('_'):
  35. return super(ProcessWrapper, self).__setattr__(key, value)
  36. return setattr(self._proc, key, value)
  37. def communicate(self, input=None):
  38. if self._proc.stdin is not None and input is None:
  39. input = b''
  40. return self._loop.run_until_complete(self._proc.communicate(input))
  41. def wait(self):
  42. return self._loop.run_until_complete(self._proc.wait())
  43. class VMWrapper(object):
  44. '''Wrap VM object to provide stable API for basic operations'''
  45. def __init__(self, vm, loop=None):
  46. self._vm = vm
  47. self._loop = loop or asyncio.get_event_loop()
  48. def __getattr__(self, item):
  49. return getattr(self._vm, item)
  50. def __setattr__(self, key, value):
  51. if key.startswith('_'):
  52. return super(VMWrapper, self).__setattr__(key, value)
  53. return setattr(self._vm, key, value)
  54. def __str__(self):
  55. return str(self._vm)
  56. def __eq__(self, other):
  57. return self._vm == other
  58. def __hash__(self):
  59. return hash(self._vm)
  60. def start(self, start_guid=True):
  61. return self._loop.run_until_complete(
  62. self._vm.start(start_guid=start_guid))
  63. def shutdown(self):
  64. return self._loop.run_until_complete(self._vm.shutdown())
  65. def run(self, command, wait=False, user=None, passio_popen=False,
  66. passio_stderr=False, **kwargs):
  67. if wait:
  68. try:
  69. self._loop.run_until_complete(
  70. self._vm.run_for_stdio(command, user=user))
  71. except subprocess.CalledProcessError as err:
  72. return err.returncode
  73. return 0
  74. elif passio_popen:
  75. p = self._loop.run_until_complete(self._vm.run(command, user=user,
  76. stdin=subprocess.PIPE,
  77. stdout=subprocess.PIPE,
  78. stderr=subprocess.PIPE if passio_stderr else None))
  79. return ProcessWrapper(p, self._loop)
  80. else:
  81. asyncio.ensure_future(self._vm.run_for_stdio(command, user=user),
  82. loop=self._loop)
  83. def run_service(self, service, wait=True, input=None, user=None,
  84. passio_popen=False,
  85. passio_stderr=False, **kwargs):
  86. if wait:
  87. try:
  88. if isinstance(input, str):
  89. input = input.encode()
  90. self._loop.run_until_complete(
  91. self._vm.run_service_for_stdio(service,
  92. input=input, user=user))
  93. except subprocess.CalledProcessError as err:
  94. return err.returncode
  95. return 0
  96. elif passio_popen:
  97. p = self._loop.run_until_complete(self._vm.run_service(service,
  98. user=user,
  99. stdin=subprocess.PIPE,
  100. stdout=subprocess.PIPE,
  101. stderr=subprocess.PIPE if passio_stderr else None))
  102. return ProcessWrapper(p, self._loop)
  103. class ExtraTestCase(qubes.tests.SystemTestCase):
  104. template = None
  105. def setUp(self):
  106. super(ExtraTestCase, self).setUp()
  107. self.init_default_template(self.template)
  108. if self.template is not None:
  109. # also use this template for DispVMs
  110. dispvm_base = self.app.add_new_vm('AppVM',
  111. name=self.make_vm_name('dvm'),
  112. template=self.template, label='red', template_for_dispvms=True)
  113. self.loop.run_until_complete(dispvm_base.create_on_disk())
  114. self.app.default_dispvm = dispvm_base
  115. def tearDown(self):
  116. self.app.default_dispvm = None
  117. super(ExtraTestCase, self).tearDown()
  118. def create_vms(self, names):
  119. """
  120. Create AppVMs for the duration of the test. Will be automatically
  121. removed after completing the test.
  122. :param names: list of VM names to create (each of them will be
  123. prefixed with some test specific string)
  124. :return: list of created VM objects
  125. """
  126. if self.template:
  127. template = self.app.domains[self.template]
  128. else:
  129. template = self.app.default_template
  130. for vmname in names:
  131. vm = self.app.add_new_vm(qubes.vm.appvm.AppVM,
  132. name=self.make_vm_name(vmname),
  133. template=template,
  134. label='red')
  135. self.loop.run_until_complete(vm.create_on_disk())
  136. self.app.save()
  137. # get objects after reload
  138. vms = []
  139. for vmname in names:
  140. vms.append(VMWrapper(self.app.domains[self.make_vm_name(vmname)],
  141. loop=self.loop))
  142. return vms
  143. def enable_network(self):
  144. """
  145. Enable access to the network. Must be called before creating VMs.
  146. """
  147. self.init_networking()
  148. def qrexec_policy(self, service, source, destination, allow=True):
  149. """
  150. Allow qrexec calls for duration of the test
  151. :param service: service name
  152. :param source: source VM name
  153. :param destination: destination VM name
  154. :return:
  155. """
  156. def add_remove_rule(add=True):
  157. with open('/etc/qubes-rpc/policy/{}'.format(service), 'r+') as policy:
  158. policy_rules = policy.readlines()
  159. rule = "{} {} {}\n".format(source, destination,
  160. 'allow' if allow else 'deny')
  161. if add:
  162. policy_rules.insert(0, rule)
  163. else:
  164. policy_rules.remove(rule)
  165. policy.truncate(0)
  166. policy.seek(0)
  167. policy.write(''.join(policy_rules))
  168. add_remove_rule(add=True)
  169. self.addCleanup(add_remove_rule, add=False)
  170. def load_tests(loader, tests, pattern):
  171. include_list = None
  172. if 'QUBES_TEST_EXTRA_INCLUDE' in os.environ:
  173. include_list = os.environ['QUBES_TEST_EXTRA_INCLUDE'].split()
  174. exclude_list = []
  175. if 'QUBES_TEST_EXTRA_EXCLUDE' in os.environ:
  176. exclude_list = os.environ['QUBES_TEST_EXTRA_EXCLUDE'].split()
  177. for entry in pkg_resources.iter_entry_points('qubes.tests.extra'):
  178. if include_list is not None and entry.name not in include_list:
  179. continue
  180. if entry.name in exclude_list:
  181. continue
  182. try:
  183. for test_case in entry.resolve()():
  184. tests.addTests(loader.loadTestsFromNames([
  185. '{}.{}'.format(test_case.__module__, test_case.__name__)]))
  186. except Exception as err: # pylint: disable=broad-except
  187. def runTest(self, err=err):
  188. raise err
  189. ExtraLoadFailure = type('ExtraLoadFailure',
  190. (qubes.tests.QubesTestCase,),
  191. {entry.name: runTest})
  192. tests.addTest(ExtraLoadFailure(entry.name))
  193. for entry in pkg_resources.iter_entry_points(
  194. 'qubes.tests.extra.for_template'):
  195. if include_list is not None and entry.name not in include_list:
  196. continue
  197. if entry.name in exclude_list:
  198. continue
  199. try:
  200. for test_case in entry.resolve()():
  201. tests.addTests(loader.loadTestsFromNames(
  202. qubes.tests.create_testcases_for_templates(
  203. test_case.__name__, test_case,
  204. module=sys.modules[test_case.__module__])))
  205. except Exception as err: # pylint: disable=broad-except
  206. def runTest(self, err=err):
  207. raise err
  208. ExtraForTemplateLoadFailure = type('ExtraForTemplateLoadFailure',
  209. (qubes.tests.QubesTestCase,),
  210. {entry.name: runTest})
  211. tests.addTest(ExtraForTemplateLoadFailure(entry.name))
  212. return tests