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)