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 test_get_value(self, reg): reg.return_value.GetKeyValue.return_value = self.value self.assertEqual(registry.get_value(self.name), self.value) reg.assert_called_with('HKLM') reg.return_value.GetKeyValue.assert_called_with( key_path=registry.constants.REG_ROOT, key_name=self.name, use_64bit=registry.constants.USE_REG_64)
def get_username() -> Optional[str]: """Gets the username from registry. Returns: username as a string. """ try: return registry.get_value('Username', path=constants.REG_ROOT) except registry.Error as e: logging.error(str(e))
def _load_time(stage_id: int, key: str) -> Optional[datetime.datetime]: """Load a time string and convert it into a native datetime value.""" val = None try: v = registry.get_value(key, 'HKLM', _stage_root(stage_id)) if v: val = datetime.datetime.strptime(v, '%Y-%m-%dT%H:%M:%S.%f') except (registry.Error, ValueError) as e: logging.error(str(e)) return None return val
def get_hostname() -> Optional[str]: """Gets the hostname value from the registry. Returns: The hostname as a string, obtained from the registry value 'name'. """ try: hostname = registry.get_value('Name', path=constants.REG_ROOT) except registry.Error as e: logging.error(str(e)) else: return hostname.strip() if hostname else hostname
def get_active_stage() -> Optional[int]: """Get the active build stage, if one exists.""" val = None try: val = registry.get_value(ACTIVE_KEY, 'HKLM', STAGES_ROOT) except registry.Error as e: logging.error(str(e)) if not val: return None val = int(val) _check_expiration(val) return val
def _load_time(stage_id, key): """Load a time string and convert it into a native datetime value.""" val = None try: val = registry.get_value(key, 'HLKM', _stage_root(stage_id)) except registry.Error as e: logging.error(str(e)) if val: try: val = datetime.datetime.strptime(val, '%Y-%m-%dT%H:%M:%S.%f') except ValueError as e: logging.error(str(e)) return None return val
def _GetLogFileName(self): """Creates the destination file name for a text log file. Returns: The full text file log name (string). """ try: hostname = registry.get_value('name', path=constants.REG_ROOT) except registry.Error as e: raise LogCopyError( 'Hostname could not be determined for log copy: %s' % str(e)) destination_file_date = gtime.now().replace(microsecond=0) destination_file_date = destination_file_date.isoformat() destination_file_date = destination_file_date.replace(':', '') return 'l:\\' + hostname + '-' + destination_file_date + '.log'
def check_winpe() -> bool: """Verify image environment is WinPE or Host. Returns: True for WinPE, else False. """ try: value = registry.get_value('EditionID', 'HKLM', r'SOFTWARE\Microsoft\Windows NT\CurrentVersion', log=False) if value == 'WindowsPE': return True else: return False except registry.Error as e: raise Error('Failed to detect image environment (%s)' % str(e))
def check_id() -> str: """Call set_id if image identifier is not set and in WinPE. Check build_info (dumped via buildinfodump) in host if image_id does not exist. Returns: Image identifier as a string if already set. """ image_id = None try: image_id = registry.get_value('image_id', path=constants.REG_ROOT) except registry.Error as e: logging.error(str(e)) if image_id: return image_id if winpe.check_winpe(): return _set_id() return _check_file()
def test_get_value_silent(self, d, reg): reg.return_value.GetKeyValue.return_value = self.value registry.get_value(self.name, log=False) self.assertFalse(d.called)
def test_get_value_none(self, reg): reg.return_value.GetKeyValue.side_effect = registry.registry.RegistryError self.assertEqual(registry.get_value(self.name), None)
def get_active_stage(): """Get the active build stage, if one exists.""" try: return registry.get_value(ACTIVE_KEY, 'HKLM', STAGES_ROOT) except registry.Error as e: logging.error(str(e))