def _check_file() -> str:
"""Call set_id if image identifier is not set and in WinPE.
Returns:
Image identifier as a string if already set.
Raises:
Error: Could not locate build info file.
Error: Could not determine image identifier from file.
"""
# First boot into host needs to grab image_id from buildinfo file.
# It has not been written to registry yet.
path = os.path.join(constants.SYS_CACHE, 'build_info.yaml')
if os.path.exists(path):
with open(path) as handle:
try:
input_config = yaml.safe_load(handle)
image_id = input_config['BUILD']['image_id']
registry.set_value('image_id',
image_id,
path=constants.REG_ROOT)
return image_id
except registry.Error as e:
raise Error(e)
except KeyError as e:
raise Error('Could not determine %s from file: %s.' %
(e, path))
else:
raise Error('Could not locate build info file.')
def set_hostname(hostname: Optional[str] = None) -> str:
"""Sets the hostname in the registry.
Gets hostname from socket.hostname if no hostname is passed.
Args:
hostname: Value to set as the hostname in registry.
Returns:
hostname: The determined hostname.
Raise:
Error: Failed to set hostname in registry.
"""
if not hostname:
hostname = socket.gethostname()
hostname = hostname.strip()
try:
registry.set_value('Name', hostname, path=constants.REG_ROOT)
except registry.Error as e:
raise Error(e)
return hostname
def CheckBeyondCorp(self) -> bool:
"""Verify whether the image is running Beyond Corp.
Returns:
True if running beyond_corp.
False if not running beyond_corp.
"""
if FLAGS.use_signed_url:
try:
registry.set_value('beyond_corp', 'True', path=constants.REG_ROOT)
return True
except registry.Error as e:
raise BCError(str(e))
else:
try:
bc = registry.get_value('beyond_corp', path=constants.REG_ROOT)
if bc:
if bc.lower() == 'true':
return True
elif bc.lower() == 'false':
return False
except registry.Error as e:
logging.warning(str(e))
try:
registry.set_value('beyond_corp', 'False', path=constants.REG_ROOT)
except registry.Error as e:
raise BCError(str(e))
return False
def _set_id() -> str:
"""Set the image id registry key."""
image_id = _generate_id()
try:
registry.set_value('image_id', image_id, path=constants.REG_ROOT)
except registry.Error as e:
raise Error(e)
return image_id
def _set_id() -> Text:
"""Set the image id registry key."""
image_id = _generate_id()
try:
registry.set_value('image_id', image_id)
except registry.Error as e:
raise Error(str(e))
return image_id
def exit_stage(stage_id: int):
"""Exit the current stage by setting the End value."""
end = _utc_now().isoformat()
logging.info('Exiting stage %d as of %s', stage_id, end)
try:
registry.set_value('End', str(end), 'HKLM', _stage_root(stage_id))
registry.set_value(ACTIVE_KEY, '', 'HKLM', STAGES_ROOT)
except registry.Error as e:
raise Error(str(e))
def test_set_value(self, reg):
registry.set_value(self.name, self.value)
reg.assert_called_with('HKLM')
reg.return_value.SetKeyValue.assert_has_calls([
mock.call(key_path=registry.constants.REG_ROOT,
key_name=self.name,
key_value=self.value,
key_type='REG_SZ',
use_64bit=registry.constants.USE_REG_64),
])
def Run(self):
timer = str(self._args[0])
key_path = r'{0}\{1}'.format(constants.REG_ROOT, 'Timers')
self._build_info.TimerSet(timer)
value_name = 'TIMER_' + timer
value_data = str(self._build_info.TimerGet(timer))
try:
registry.set_value(value_name, value_data, 'HKLM', key_path, log=False)
logging.info('Set image timer: %s', timer)
except registry.Error as e:
raise Error(str(e))
def _WriteRegistry(self, reg_values):
"""Populates the registry with build_info settings for future reference.
Args:
reg_values: A dictionary of key/value pairs to be added to the registry.
"""
for value_name in reg_values:
key_path = constants.REG_ROOT
value_data = reg_values[value_name]
if 'TIMER_' in value_name:
key_path = r'{0}\{1}'.format(constants.REG_ROOT, 'Timers')
try:
registry.set_value(value_name, value_data, 'HKLM', key_path)
except registry.Error as e:
raise ActionError(str(e))
def set_stage(stage_id: int):
"""Sets or updates the current build stage."""
if not isinstance(stage_id, int):
raise Error('Invalid stage type; got: %s, want: int' % type(stage_id))
active = get_active_stage()
if active:
exit_stage(active)
start = _utc_now().isoformat()
try:
registry.set_value('Start', str(start), 'HKLM', _stage_root(stage_id))
registry.set_value(ACTIVE_KEY, str(stage_id), 'HKLM', STAGES_ROOT)
except registry.Error as e:
raise Error(str(e))
def Run(self):
use_64bit = constants.USE_REG_64
if len(self._args) > 5:
use_64bit = self._args[5]
try:
registry.set_value(self._args[2], self._args[3],
self._args[0], self._args[1],
self._args[4], use_64bit=use_64bit)
except registry.Error as e:
raise ActionError(str(e))
except IndexError:
raise ActionError('Unable to access all required arguments. [%s]' %
str(self._args))
def set_username(username: Optional[str] = None,
prompt: Optional[str] = None) -> str:
"""Sets the username in the registry.
Optionally prompts if there is no username supplied as a parameter.
Args:
username: Value to set as the username in registry.
prompt: Custom string to append to username prompt.
Returns:
username: The determined username.
Raises:
Error: Failed to set username in registry.
"""
if not username:
username = interact.GetUsername(prompt)
try:
registry.set_value('Username', username, path=constants.REG_ROOT)
except registry.Error as e:
raise Error(e)
return username
def test_set_value_silent(self, d, unused_reg):
registry.set_value(self.name, self.value, log=False)
self.assertFalse(d.called)
