ermine_ctl.py 5.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196
  1. # Copyright 2022 The Chromium Authors. All rights reserved.
  2. # Use of this source code is governed by a BSD-style license that can be
  3. # found in the LICENSE file.
  4. """Adds python interface to erminectl tools on workstation products."""
  5. from typing import Any, List, Tuple, Union
  6. import enum
  7. import logging
  8. import subprocess
  9. import time
  10. class ErmineCtl:
  11. """Tool for automating control of Ermine and its OOBE, if available.
  12. Must be used after checking if the tool exists.
  13. Usage:
  14. ctl = ermine_ctl.ErmineCtl(some_target)
  15. if ctl.exists:
  16. ctl.WaitUntilReady()
  17. ctl.TakeToShell()
  18. logging.info('In the shell')
  19. else:
  20. logging.info('Tool does not exist!')
  21. This is only necessary after a target reboot or provision (IE pave).
  22. """
  23. _OOBE_PASSWORD = 'some_test_password'
  24. _TOOL = 'erminectl'
  25. _OOBE_SUBTOOL = 'oobe'
  26. _MAX_STATE_TRANSITIONS = 5
  27. # Mapping between the current state and the next command to run
  28. # to move it to the next state.
  29. _STATE_TO_NEXT = {
  30. 'SetPassword': ['set_password', _OOBE_PASSWORD],
  31. 'Unknown': ['skip'],
  32. 'Shell': [],
  33. 'Login': ['login', _OOBE_PASSWORD],
  34. }
  35. _COMPLETE_STATE = 'Shell'
  36. _READY_TIMEOUT = 10
  37. _WAIT_ATTEMPTS = 10
  38. _WAIT_FOR_READY_SLEEP_SEC = 3
  39. def __init__(self, target: Any):
  40. self.target = target
  41. self._ermine_exists = False
  42. self._ermine_exists_check = False
  43. @property
  44. def exists(self) -> bool:
  45. """Returns the existence of the tool.
  46. Checks whether the tool exists on and caches the result.
  47. Returns:
  48. True if the tool exists, False if not.
  49. """
  50. if not self._ermine_exists_check:
  51. self._ermine_exists = self.target.RunCommand([self._TOOL, '--help']) == 0
  52. self._ermine_exists_check = True
  53. logging.debug('erminectl exists: %s',
  54. ('true' if self._ermine_exists else 'false'))
  55. return self._ermine_exists
  56. @property
  57. def status(self) -> Tuple[int, str]:
  58. """Returns the status of ermine.
  59. Note that if the tool times out or does not exist, a non-zero code
  60. is returned.
  61. Returns:
  62. Tuple of (return code, status as string). -1 for timeout, and
  63. -2 for no tool.
  64. """
  65. if self.exists:
  66. # Executes base command, which returns status.
  67. proc = self._ExecuteCommandAsync([])
  68. try:
  69. proc.wait(timeout=self._READY_TIMEOUT)
  70. except subprocess.TimeoutExpired:
  71. logging.warning('Timed out waiting for status')
  72. return -1, 'Timeout'
  73. stdout, _ = proc.communicate()
  74. return proc.returncode, stdout.strip()
  75. return -2, 'InvalidState'
  76. @property
  77. def ready(self) -> bool:
  78. """Indicates if the tool is ready for regular use.
  79. Returns:
  80. False if not ready or does not exist, and True if ready.
  81. """
  82. if self.exists:
  83. rc, _ = self.status
  84. return rc == 0
  85. return False
  86. def _ExecuteCommandAsync(self, command: List[str]) -> subprocess.Popen:
  87. """Executes a sub-command asynchronously.
  88. Args:
  89. command: list of strings to compose the command. Forwards to the
  90. command runner.
  91. Returns:
  92. Popen of the subprocess.
  93. """
  94. full_command = [self._TOOL, self._OOBE_SUBTOOL]
  95. full_command.extend(command)
  96. # Returns immediately with Popen.
  97. return self.target.RunCommandPiped(full_command,
  98. stdout=subprocess.PIPE,
  99. stderr=subprocess.STDOUT,
  100. text=True)
  101. def _ExecuteCommand(self, command: List[str]):
  102. """Executes a sub-command of the tool synchronously.
  103. Raises exception if non-zero returncode is given.
  104. Args:
  105. command: list of strings to compose the command. Forwards to the
  106. command runner.
  107. Raises:
  108. RuntimeError: if non-zero returncode is returned.
  109. """
  110. proc = self._ExecuteCommandAsync(command)
  111. proc.wait()
  112. stdout, stderr = proc.communicate()
  113. if proc.returncode != 0:
  114. raise RuntimeError(f'Command {" ".join(command)} failed.'
  115. f'\nSTDOUT: {stdout}\nSTDERR: {stderr}')
  116. def WaitUntilReady(self) -> bool:
  117. """Waits until the tool is ready through sleep-poll.
  118. The tool may not be ready after a pave or restart.
  119. This checks the status and exits after its ready or Timeout.
  120. Returns:
  121. True if the tool exists and is ready. False if the tool does not exist.
  122. Raises:
  123. TimeoutError: if tool is not ready after certain amount of attempts.
  124. """
  125. if self.exists:
  126. for _ in range(self._WAIT_ATTEMPTS):
  127. if self.ready:
  128. break
  129. time.sleep(self._WAIT_FOR_READY_SLEEP_SEC)
  130. else:
  131. raise TimeoutError('Timed out waiting for a valid status to return')
  132. return True
  133. return False
  134. def TakeToShell(self):
  135. """Takes device to shell after waiting for tool to be ready.
  136. Examines the current state of the device after waiting for it to be ready.
  137. Once ready, goes through the states of logging in. This is:
  138. - CreatePassword -> Skip screen -> Shell
  139. - Login -> Shell
  140. - Shell
  141. Regardless of starting state, this will exit once the shell state is
  142. reached.
  143. Raises:
  144. NotImplementedError: if an unknown state is reached.
  145. RuntimeError: If number of state transitions exceeds the max number that
  146. is expected.
  147. """
  148. assert self.WaitUntilReady(), 'erminectl does not exist, cannot be ready'
  149. _, state = self.status
  150. max_states = self._MAX_STATE_TRANSITIONS
  151. while state != self._COMPLETE_STATE and max_states:
  152. max_states -= 1
  153. command = self._STATE_TO_NEXT.get(state)
  154. logging.debug('Ermine state is: %s', state)
  155. if command is None:
  156. raise NotImplementedError('Encountered invalid state: %s' % state)
  157. self._ExecuteCommand(command)
  158. _, state = self.status
  159. if not max_states:
  160. raise RuntimeError('Did not transition to shell in %d attempts.'
  161. ' Please file a bug.' % self._MAX_STATE_TRANSITIONS)