202 lines
7.0 KiB
Python
202 lines
7.0 KiB
Python
|
# Copyright 2022 The Chromium Authors
|
||
|
|
# Use of this source code is governed by a BSD-style license that can be
|
||
|
|
# found in the LICENSE file.
|
||
|
|
"""Adds python interface to erminectl tools on workstation products."""
|
||
|
|
|
||
|
|
import logging
|
||
|
|
import subprocess
|
||
|
|
import time
|
||
|
|
from typing import List, Tuple
|
||
|
|
|
||
|
|
|
||
|
|
class BaseErmineCtl:
|
||
|
|
"""Compatible class for automating control of Ermine and its OOBE.
|
||
|
|
|
||
|
|
Must be used after checking if the tool exists.
|
||
|
|
|
||
|
|
Usage:
|
||
|
|
ctl = base_ermine_ctl.BaseErmineCtl(some_target)
|
||
|
|
if ctl.exists:
|
||
|
|
ctl.take_to_shell()
|
||
|
|
|
||
|
|
logging.info('In the shell')
|
||
|
|
else:
|
||
|
|
logging.info('Tool does not exist!')
|
||
|
|
|
||
|
|
This is only necessary after a target reboot or provision (IE pave).
|
||
|
|
"""
|
||
|
|
|
||
|
|
_OOBE_PASSWORD = 'workstation_test_password'
|
||
|
|
_TOOL = 'erminectl'
|
||
|
|
_OOBE_SUBTOOL = 'oobe'
|
||
|
|
_MAX_STATE_TRANSITIONS = 5
|
||
|
|
|
||
|
|
# Mapping between the current state and the next command to run
|
||
|
|
# to move it to the next state.
|
||
|
|
_STATE_TO_NEXT = {
|
||
|
|
'SetPassword': ['set_password', _OOBE_PASSWORD],
|
||
|
|
'Unknown': ['skip'],
|
||
|
|
'Shell': [],
|
||
|
|
'Login': ['login', _OOBE_PASSWORD],
|
||
|
|
}
|
||
|
|
_COMPLETE_STATE = 'Shell'
|
||
|
|
|
||
|
|
_READY_TIMEOUT = 10
|
||
|
|
_WAIT_ATTEMPTS = 10
|
||
|
|
_WAIT_FOR_READY_SLEEP_SEC = 3
|
||
|
|
|
||
|
|
def __init__(self):
|
||
|
|
self._ermine_exists = False
|
||
|
|
self._ermine_exists_check = False
|
||
|
|
|
||
|
|
# pylint: disable=no-self-use
|
||
|
|
# Overridable method to determine how command gets executed.
|
||
|
|
def execute_command_async(self, args: List[str]) -> subprocess.Popen:
|
||
|
|
"""Executes command asynchronously, returning immediately."""
|
||
|
|
raise NotImplementedError
|
||
|
|
|
||
|
|
# pylint: enable=no-self-use
|
||
|
|
|
||
|
|
@property
|
||
|
|
def exists(self) -> bool:
|
||
|
|
"""Returns the existence of the tool.
|
||
|
|
|
||
|
|
Checks whether the tool exists on and caches the result.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
True if the tool exists, False if not.
|
||
|
|
"""
|
||
|
|
if not self._ermine_exists_check:
|
||
|
|
self._ermine_exists = self._execute_tool(['--help'],
|
||
|
|
can_fail=True) == 0
|
||
|
|
self._ermine_exists_check = True
|
||
|
|
logging.debug('erminectl exists: %s',
|
||
|
|
('true' if self._ermine_exists else 'false'))
|
||
|
|
return self._ermine_exists
|
||
|
|
|
||
|
|
@property
|
||
|
|
def status(self) -> Tuple[int, str]:
|
||
|
|
"""Returns the status of ermine.
|
||
|
|
|
||
|
|
Note that if the tool times out or does not exist, a non-zero code
|
||
|
|
is returned.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
Tuple of (return code, status as string). -1 for timeout.
|
||
|
|
Raises:
|
||
|
|
AssertionError: if the tool does not exist.
|
||
|
|
"""
|
||
|
|
assert self.exists, (f'Tool {self._TOOL} cannot have a status if'
|
||
|
|
' it does not exist')
|
||
|
|
# Executes base command, which returns status.
|
||
|
|
proc = self._execute_tool_async([])
|
||
|
|
try:
|
||
|
|
proc.wait(timeout=self._READY_TIMEOUT)
|
||
|
|
except subprocess.TimeoutExpired:
|
||
|
|
logging.warning('Timed out waiting for status')
|
||
|
|
return -1, 'Timeout'
|
||
|
|
stdout, _ = proc.communicate()
|
||
|
|
return proc.returncode, stdout.strip()
|
||
|
|
|
||
|
|
@property
|
||
|
|
def ready(self) -> bool:
|
||
|
|
"""Indicates if the tool is ready for regular use.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
False if not ready, and True if ready.
|
||
|
|
Raises:
|
||
|
|
AssertionError: if the tool does not exist.
|
||
|
|
"""
|
||
|
|
assert self.exists, (f'Tool {self._TOOL} cannot be ready if'
|
||
|
|
' it does not exist')
|
||
|
|
return_code, _ = self.status
|
||
|
|
return return_code == 0
|
||
|
|
|
||
|
|
def _execute_tool_async(self, command: List[str]) -> subprocess.Popen:
|
||
|
|
"""Executes a sub-command asynchronously.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
command: list of strings to compose the command. Forwards to the
|
||
|
|
command runner.
|
||
|
|
Returns:
|
||
|
|
Popen of the subprocess.
|
||
|
|
"""
|
||
|
|
full_command = [self._TOOL, self._OOBE_SUBTOOL]
|
||
|
|
full_command.extend(command)
|
||
|
|
|
||
|
|
# Returns immediately with Popen.
|
||
|
|
return self.execute_command_async(full_command)
|
||
|
|
|
||
|
|
def _execute_tool(self, command: List[str], can_fail: bool = False) -> int:
|
||
|
|
"""Executes a sub-command of the tool synchronously.
|
||
|
|
Raises exception if non-zero returncode is given and |can_fail| = False.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
command: list of strings to compose the command. Forwards to the
|
||
|
|
command runner.
|
||
|
|
can_fail: Whether or not the command can fail.
|
||
|
|
Raises:
|
||
|
|
RuntimeError: if non-zero returncode is returned and can_fail =
|
||
|
|
False.
|
||
|
|
Returns:
|
||
|
|
Return code of command execution if |can_fail| is True.
|
||
|
|
"""
|
||
|
|
proc = self._execute_tool_async(command)
|
||
|
|
stdout, stderr = proc.communicate()
|
||
|
|
if not can_fail and proc.returncode != 0:
|
||
|
|
raise RuntimeError(f'Command {" ".join(command)} failed.'
|
||
|
|
f'\nSTDOUT: {stdout}\nSTDERR: {stderr}')
|
||
|
|
return proc.returncode
|
||
|
|
|
||
|
|
def wait_until_ready(self) -> None:
|
||
|
|
"""Waits until the tool is ready through sleep-poll.
|
||
|
|
|
||
|
|
The tool may not be ready after a pave or restart.
|
||
|
|
This checks the status and exits after its ready or Timeout.
|
||
|
|
|
||
|
|
Raises:
|
||
|
|
TimeoutError: if tool is not ready after certain amount of attempts.
|
||
|
|
AssertionError: if tool does not exist.
|
||
|
|
"""
|
||
|
|
assert self.exists, f'Tool {self._TOOL} must exist to use it.'
|
||
|
|
for _ in range(self._WAIT_ATTEMPTS):
|
||
|
|
if self.ready:
|
||
|
|
return
|
||
|
|
time.sleep(self._WAIT_FOR_READY_SLEEP_SEC)
|
||
|
|
raise TimeoutError('Timed out waiting for a valid status to return')
|
||
|
|
|
||
|
|
def take_to_shell(self) -> None:
|
||
|
|
"""Takes device to shell after waiting for tool to be ready.
|
||
|
|
|
||
|
|
Examines the current state of the device after waiting for it to be
|
||
|
|
ready. Once ready, goes through the states of logging in. This is:
|
||
|
|
- CreatePassword -> Skip screen -> Shell
|
||
|
|
- Login -> Shell
|
||
|
|
- Shell
|
||
|
|
|
||
|
|
Regardless of starting state, this will exit once the shell state is
|
||
|
|
reached.
|
||
|
|
|
||
|
|
Raises:
|
||
|
|
NotImplementedError: if an unknown state is reached.
|
||
|
|
RuntimeError: If number of state transitions exceeds the max number
|
||
|
|
that is expected.
|
||
|
|
"""
|
||
|
|
self.wait_until_ready()
|
||
|
|
_, state = self.status
|
||
|
|
max_states = self._MAX_STATE_TRANSITIONS
|
||
|
|
while state != self._COMPLETE_STATE and max_states:
|
||
|
|
max_states -= 1
|
||
|
|
command = self._STATE_TO_NEXT.get(state)
|
||
|
|
logging.debug('Ermine state is: %s', state)
|
||
|
|
if command is None:
|
||
|
|
raise NotImplementedError('Encountered invalid state: %s' %
|
||
|
|
state)
|
||
|
|
self._execute_tool(command)
|
||
|
|
_, state = self.status
|
||
|
|
|
||
|
|
if not max_states:
|
||
|
|
raise RuntimeError('Did not transition to shell in %d attempts.'
|
||
|
|
' Please file a bug.' %
|
||
|
|
self._MAX_STATE_TRANSITIONS)
|