Chrome updater test framework

install_test/chrome_installer.py

+#!/usr/bin/python
+# Copyright (c) 2012 The Chromium Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Provides an interface for installing Chrome."""
+
+import _winreg
+import ctypes
+import httplib
+import logging
+import os
+import platform
+import shutil
+import socket
+import subprocess
+import tempfile
+import urllib
+
+
+_CSIDL_COMMON_APPDATA = 0x1C
+_CSIDL_PROGRAM_FILESX86 = 0x2A
+_PLATFORM = platform.system().lower()
+
+# Import only on Windows. On other platforms it will invoke a ValueError. 
+if _PLATFORM == 'windows':
+  from ctypes import wintypes, windll
+
+
+class InstallationType:
+  """Defines the Chrome installation types."""
+  SYSTEM = 0
+  USER = 1
+
+
+def _RemoveDuplicates(args):
+  """Removes duplicates options from a list.
+
+  Args:
+    args: A list containing the options.
+
+  Returns:
+    A string without any duplicate options.
+  """
+  return ((lambda arg: arg if arg else '')
+          (' '.join(map(lambda opt: opt, list(frozenset(args.split(',')))))))
+
+
+def Install(installer_path, chrome_type, build, options='', clean=True):  
+  """Installs the specified Chrome build."""
+  def DoPreliminaryChecks(chrome_install):    
+    """Validates the test parameters and Chrome version."""
+    assert(_PLATFORM == 'windows' and os.path.isfile(installer_path))
+    # Chrome installation type specified by the user.
+    install_type = ('system-level' in options and InstallationType.SYSTEM
+                    or InstallationType.USER)     
+    cur_build = chrome_install.GetChromeVersion()
+    # Chrome already installed, make sure new build can be installed over it.
+    if cur_build:            
+      current_type = chrome_install.GetChromeInstallationType()      
+      if(current_type == InstallationType.SYSTEM and install_type ==
+         InstallationType.USER):
+        raise RuntimeError('System level Chrome exists, aborting user level '
+                           'installation.')
+      # Installing a build that's older than the currently installed build.
+      elif current_type == install_type:                
+        if cur_build >= build:
+          raise RuntimeError('Please specify a newer version of Chrome.')
+        
+  chrome_install = Installation(chrome_type)
+  DoPreliminaryChecks(chrome_install)
+  options += ' --install --do-not-launch-chrome'
+  options = _RemoveDuplicates(options)
+  logging.log(logging.INFO, 'Launching Chrome installer...')
+  cmd = '%s %s' % (installer_path, options)
+  ret = subprocess.Popen(cmd, shell=True).wait()   
+  if ret == 0:
+    logging.log(logging.INFO, 'Installation complete.')
+    return chrome_install
+  logging.log(logging.ERROR, 'Installation failed.')
+  return None
+
+  
+class Installation(object):
+  """Provides pertinent information about the installed Chrome version.
+
+  The type of Chrome version must be passed as an argument to the constructor, 
+  (i.e. - user or system level).
+  """
+
+  HKEY_LOCAL = (r'SOFTWARE\Wow6432Node\Google\Update\ClientState'
+                '\{8A69D345-D564-463C-AFF1-A69D9E530F96}')
+  HKEY_USER  = HKEY_LOCAL.replace('\\Wow6432Node', '')                
+  
+  def __init__(self, install_type):
+    assert(install_type == InstallationType.SYSTEM or
+           install_type == InstallationType.USER)
+    self._type = install_type 
+    self._key_type = self._GetKeyType(self._type)
+
+  def _OpenKey(self, key_type, key_name, key_access=_winreg.KEY_READ):
+    """Opens a registry key and returns the key handle.
+
+    Args:
+      key_type: The type of key HKLM or HKCU.
+      key_name: Name of the key.
+      key_access: Type of access required. Be default its read only.
+
+    Returns:
+      Key handle if successful, otherwise None.
+    """
+    try:
+      key = _winreg.OpenKey(key_type, key_name, 0, key_access)
+      return key
+    except _winreg.error:
+      return None
+
+  def _GetWinLocalFolder(self, ftype=_CSIDL_COMMON_APPDATA):
+    """Returns full path of the 'Local' folder on Windows.
+
+    Args:
+      ftype: Location to look up, which could vary based on installation type.
+
+    Returns:
+      A String representing the folder path if successful, otherwise an empty
+      string.
+    """
+    if _PLATFORM != 'windows':      
+      return ''    
+    SHGetFolderPathW = windll.shell32.SHGetFolderPathW
+    SHGetFolderPathW.argtypes = [wintypes.HWND,
+                                 ctypes.c_int,
+                                 wintypes.HANDLE,
+                                 wintypes.DWORD,
+                                 wintypes.LPCWSTR]
+    path_buf = wintypes.create_unicode_buffer(wintypes.MAX_PATH)
+    result = SHGetFolderPathW(0, ftype, 0, 0, path_buf)    
+    return str(path_buf.value)
+
+  def _GetKeyType(self, install_type):
+    """Determines the registry key to use based on installation type."""
+    if install_type == InstallationType.SYSTEM:      
+      return _winreg.HKEY_LOCAL_MACHINE
+    elif install_type == InstallationType.USER:      
+      return _winreg.HKEY_CURRENT_USER    
+  
+  def _GetKeyName(self, install_type, replace=True):
+    """Returns Chrome's registry key name based on installation type."""
+    name = (install_type == InstallationType.SYSTEM and self.HKEY_LOCAL or
+            self.HKEY_USER)
+    return (replace and name.replace('ClientState', 'Clients') or name)
+
+  def _DeleteChromeRegEntries(self):
+    """Deletes chrome registry settings."""    
+    if self._type == InstallationType.USER:
+      p_key = self.HKEY_USER[: self.HKEY_USER.rfind('\\')]      
+      key_name = self.HKEY_USER[self.HKEY_USER.rfind('\\') + 1 :]
+      _type = _winreg.HKEY_CURRENT_USER    
+    elif self._type == InstallationType.SYSTEM:
+      p_key = self.HKEY_LOCAL[: self.HKEY_LOCAL.rfind('\\')]
+      key_name = self.HKEY_LOCAL[self.HKEY_LOCAL.rfind('\\') + 1 :]
+      _type = _winreg.HKEY_LOCAL_MACHINE    
+    try:
+      key = self._OpenKey(_type, p_key, _winreg.KEY_ALL_ACCESS)      
+      _winreg.DeleteKey(key, key_name)
+      key.Close()
+      return 0      
+    except _winreg.error, err:
+      logging.log(logging.ERROR, 'Could not delete registry entries: %s' % err)      
+      return -1
+
+  def _GetValueFromRegistry(self, reg_type, key, subkey):
+    """Gets value of the specified subkey from the registry.
+
+    Args:
+      reg_type: Type of key to use HKLM or HKCU.
+      key: Name of the key.
+      subkey: Name of the subkey whose value will be returned.
+
+    Returns:
+      A string representing the subkey value if successful, otherwise None.      
+    """    
+    hkey = self._OpenKey(reg_type, key)
+    if hkey:
+      try:
+        value = str(_winreg.QueryValueEx(hkey, subkey)[0])        
+      except _winreg.error:
+        value = None
+      _winreg.CloseKey(hkey)
+      return value
+    return None
+
+  def _GetUninstallString(self):
+    """Returns Chrome uninstall string from the registry."""
+    key_name = self._GetKeyName(self._type, False)    
+    return self._GetValueFromRegistry(self._key_type, key_name,
+                                      'UninstallString')
+
+  def _GetUninstallArguments(self):
+    """Returns Chrome uninstall arguments from the registry."""
+    key_name = self._GetKeyName(self._type, False)
+    return self._GetValueFromRegistry(self._key_type, key_name,
+                                      'UninstallArguments')
+
+  def GetChromeVersion(self):
+    """Returns the installed version of Chrome."""
+    key_name = self._GetKeyName(self._type, True)
+    return self._GetValueFromRegistry(self._key_type, key_name, 'pv')
+
+  def _LaunchInstaller(self, installer_path, options):
+    """Launches the Chrome installer.
+
+    Args:
+      installer_path: Path where the installer is located.
+      options: Any additional options to be used for installation.
+    """
+    cmd = '%s %s' % (installer_path, options)
+    try:
+      subprocess.Popen(cmd, shell=True).wait()
+    except OSError, err:
+      raise err
+
+  def _Delete(self, _path):
+    """Deletes a file or folder."""
+    try:
+      (lambda p: shutil.rmtree(p) if os.path.isdir(p) else os.remove(p))(_path)
+      return 0
+    except(OSError, IOError, TypeError), err:      
+      return -1
+    
+  def GetChromeExePath(self):
+    """Returns Chrome binary location based on installation type."""
+    chrome_path = ''
+    if self._type == InstallationType.USER:
+      folder_id = _CSIDL_COMMON_APPDATA
+    elif self._type == InstallationType.SYSTEM:
+      folder_id = _CSIDL_PROGRAM_FILESX86
+    if _PLATFORM == 'windows':
+      chrome_path = os.path.join(self._GetWinLocalFolder(folder_id), 'Google',
+                                 'Chrome', 'Application', 'chrome.exe')
+    return chrome_path
+
+  def GetChromeInstallationType(self):
+    """Determines Chrome installation type."""
+    reg_type = _winreg.HKEY_CURRENT_USER
+    key_name = self.HKEY_USER.replace('ClientState', 'Clients')
+    key = self._OpenKey(reg_type, key_name)    
+    if not key:      
+      reg_type = _winreg.HKEY_LOCAL_MACHINE
+      key_name = self.HKEY_LOCAL.replace('ClientState', 'Clients')
+      key = self._OpenKey(reg_type, key_name)
+      if not key:        
+        return None    
+    version = self._GetValueFromRegistry(reg_type, key_name, 'pv')                                  
+    key.Close()    
+    if version:      
+      return (reg_type == _winreg.HKEY_CURRENT_USER and
+              InstallationType.USER or InstallationType.SYSTEM)
+    return None
+  
+  def UninstallChrome(self):
+    """Uninstalls Chrome."""
+    if not self.GetChromeVersion():
+      raise RuntimeError('No Chrome version found on this system.')
+    chrome_path = self.GetChromeExePath()
+    install_type = self.GetChromeInstallationType()
+    reg_opts = self._GetUninstallArguments()
+    uninstall_str = self._GetUninstallString()
+    options = '%s --force-uninstall' % (reg_opts)
+    if self._type == InstallationType.SYSTEM:
+      options += ' --system-level'    
+    if not os.path.exists(chrome_path):
+      logging.log(logging.ERROR, 'Could not find chrome, aborting uninstall.')
+      return False
+    logging.log(logging.INFO, 'Launching Chrome installer...')
+    ret = self._LaunchInstaller(uninstall_str, options)    
+    if not os.path.exists(chrome_path):
+      logging.log(logging.INFO, 'Chrome was uninstalled successfully...')      
+      # TODO: add clean option like ChromeInstaller
+      logging.log(logging.INFO, 'Deleting registry entries...')        
+      self._DeleteChromeRegEntries()
+      logging.log(logging.INFO, 'Uninstall complete.')      
+      return True
+    else:
+      logging.log(logging.ERROR, 'Uninstall failed.')      
+    return False