Chrome updater test framework

pyautolib/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.
+
+import os
+import shutil
+import ctypes
+import urllib
+import socket
+import httplib
+import _winreg
+import tempfile
+import platform
+import urlparse
+import subprocess
+
+_PLATFORM = platform.system().lower()
+_CSIDL_COMMON_APPDATA = 28
+_CSIDL_PROGRAM_FILESX86 = 42
+
+# Import only on Windows. On other platforms it will invoke a ValueError. 
+if _PLATFORM == 'windows':
+  from ctypes import wintypes, windll
+
+
+class ChromeInstaller:
+  """This class is used to install or uninstall Chrome."""
+  
+  def __init__(self, url, build, dest, opts='', clean=True):
+    self._url = (lambda u: u if self._DoesUrlExist(u) else '')(url)
+    self._dest = dest
+    self._opts = opts
+    self._clean = clean
+    self._type = (lambda o: 'system' if o.find('system-level') != -1 else
+                  'user')(self._opts)
+    self._plat = 'win'
+    self._installer = 'mini_installer.exe'       
+    self.SetCurrentBuild(build)
+    self.c_install = ChromeVersion(self._type)
+    self._c_path = self.c_install.GetChromeExePath()
+
+  def SetCurrentBuild(self, bld):
+    """Sets the build number that is to be installed."""
+    if type(bld) == str:
+      self.build = (lambda b: (all(map(lambda x: x.isdigit(), b.split('.')))
+                               and len(b.split('.')) == 4) and b or '')(bld)
+    else:
+      self.build = ''
+          
+  def _DoesUrlExist(self, url):
+    """Checks if a url exists.
+
+    Args:
+      url: URL to be verified.
+
+    Returns:
+      Boolean - True if the URL exists, otherwise False.
+    """          
+    parse = urlparse.urlparse(url)    
+    if parse[0] == '' or parse[1] == '':
+      return False
+    try:
+      conn = httplib.HTTPConnection(parse.netloc)
+      conn.request('HEAD', parse.path)
+      res = conn.getresponse()      
+    except socket.gaierror:
+      return False
+    finally:
+      conn.close()
+    # Redirected - (301: permanently moved, 302: temporarily moved)
+    if res.status == 302 or res.status == 301:
+      return self._DoesUrlExist(res.getheader('location'))
+    return res.status == 200 
+
+  def _Download(self, _file, url, dest=None):
+    """Downloads a file from the specified URL
+
+    Args:
+      _file: Name of the file to download.
+      url: URL where the file is located.
+      dest: Where file will be downloaded (optional). Default location is CWD.
+
+    Returns:
+      An integer - zero if successful, otherwise a non-zero value.
+    """
+    f_name = (lambda d, f: os.path.join(d, f) if(d and os.path.exists(d)) else
+              os.path.join(tempfile.gettempdir(), f))(dest, _file)        
+    file_url = '%s/%s' % (url, _file)    
+    # Check link before downloading, urlretrieve *WILL NOT* report
+    # an error if url doesn't exist.
+    if not self._DoesUrlExist(file_url):
+      print "Could not locate the file specified."
+      return -1
+    try:
+      d = urllib.urlretrieve(file_url, f_name)
+    except IOError, err:
+      print 'CChromeInstaller._Download: ', err
+      return -1
+    if os.path.isfile(d[0]):
+      return 0
+    return 1
+
+  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:
+      print 'CChromeInstaller._Delete: ', err
+      return -1
+    
+  def _LaunchInstaller(self, opts=''):
+    """Launches the Chrome installer using specified options.
+
+    Args:
+      opts: Any options that are to be passed to the installer.
+
+    Returns:
+      An integer - zero if successful, otherwise a non-zero value.
+    """
+    print 'Downloading installer for build %s ...' % (self.build)
+    build_url = '%s/%s/%s' % (self._url, self.build, self._plat)
+    if self._Download(self._installer, build_url, self._dest) != 0:
+      print 'Failed to download installer, exiting test'
+      return -1    
+    print 'Launching installer...'
+    cmd = '%s %s' % (os.path.join(self._dest, self._installer), opts)
+    try:
+      return subprocess.Popen.wait(subprocess.Popen(cmd))      
+    except OSError, err:
+      print 'CChromeInstaller._LaunchInstaller: ', err
+      return -2
+
+  def _RemoveDuplicates(self, args):
+    """Removes duplicates from list of options.
+
+    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 _DoPreliminaryChecks(self, op='install'):
+    """Checks test parameters for validity and performs version checking."""
+    if _PLATFORM != 'windows':
+      print 'Unsupported platform, aborting operation...'
+      return False
+    if self._url == '':
+      print 'URL specified is not valid. Please check the URL and try again.'
+      return False     
+    ver = self.c_install.GetChromeVersion()    
+    c_type = self.c_install.GetInstallTypeString()
+    if op == 'install':
+      self.pyauto_url = '%s/%s/%s/%s' % (self._url, self.build, 'win',
+                                         'chrome-win32.test')      
+      if c_type == 'system' and self._type == 'user':       
+        print 'System level Chrome exists, aborting user level installation.'
+        return False
+      # Version installed is higher than or same as the one we're installing. 
+      elif self._type == c_type:
+        if self.build <= ver:
+          print 'Please specify a version higher than the one installed.'
+          return False
+    elif op == 'uninstall':
+      if not ver:
+        print 'No version of Chrome found, aborting uninstall...'
+        return False
+      if self._type != c_type:
+        print ('Invalid args passed. Please specify --opts=--system-level if '
+               'uninstalling system version of Chrome, or omit it otherwise.')
+        return False
+    return True
+
+  def InstallChrome(self):
+    """Installs specified chrome build."""
+    if not self._DoPreliminaryChecks('install'):
+      return None
+    opts = '--install --do-not-launch-chrome %s' % (self._opts)
+    opts = self._RemoveDuplicates(opts)    
+    ret = self._LaunchInstaller(opts)
+    self._Delete(os.path.join(self._dest, self._installer))       
+    if ret == 0 and os.path.exists(self._c_path):
+      print 'Installation complete...'
+      return self.c_install
+    else:
+      print 'Installation failed.'
+    return None
+
+  def UninstallChrome(self):
+    """Uninstalls Chrome."""
+    if not self._DoPreliminaryChecks('uninstall'):
+      return False
+    install_type = self.c_install.GetChromeInstallType()
+    reg_opts = self.c_install.GetUninstallArgs(install_type)
+    opts = '%s --force-uninstall %s' % (reg_opts, self._opts)    
+    opts = self._RemoveDuplicates(opts)
+    if not os.path.exists(self._c_path):
+      print 'Couldn\'t find Chrome. Verify Chrome is installed and try again.'
+      return False
+    ret = self._LaunchInstaller(opts)
+    if os.path.exists(os.path.join(self._dest, self._installer)):
+      self._Delete(os.path.join(self._dest, self._installer))    
+    if not os.path.exists(self._c_path):
+      print 'Chrome was successfully uninstalled...'
+      if self._clean:
+        print 'Removing registry settings...'        
+        self.c_install.DeleteChromeRegEntries()
+        print 'Uninstall complete.'
+        return True
+    else:
+      print 'Uninstall failed.'
+    return False
+
+
+class ChromeVersion():
+  """Contains utility functions that provide information about installed
+
+  version of Chrome. The type of Chrome version is passed as an argument
+  to the constructor (i.e. - user or system level version).
+  """
+  
+  def __init__(self, install_type):
+    self._type = install_type
+    self.HKEY_LOCAL = (r'SOFTWARE\Wow6432Node\Google\Update\ClientState'
+                       '\{8A69D345-D564-463C-AFF1-A69D9E530F96}')
+    self.HKEY_USER = self.HKEY_LOCAL.replace('\\Wow6432Node', '')      
+
+  def _OpenKey(self, key_type, key_name, key_access=_winreg.KEY_READ):
+    """Opens a registry key.
+
+    Args:
+      key_type: HKEY_CURRENT_USER or HKEY_LOCAL_MACHINE.
+      key_name: Name of the key that is to be opened.
+      key_access: Type of access required (i.e. - KEY_READ or KEY_ALL_ACCESS).
+      
+    Returns:
+      Handle to the key if successful, otherwise zero.
+    """      
+    try:
+      key = _winreg.OpenKey(key_type, key_name, 0, key_access)
+      return key
+    except(OSError, TypeError), err:      
+      return 0
+
+  def _GetKeyValue(self, key, name):
+    """Gets the value of the specified key.
+
+    Args:
+      key: Handle to the key.
+      name: Name of the key.
+
+    Returns:
+      A string representing key value if successful, otherwise empty string.
+    """
+    try:
+      val = str(_winreg.QueryValueEx(key, name)[0])
+      return val
+    except OSError, err:
+      print 'ChromeVersion._GetKeyValue: %s' % (err)
+      return ''
+
+  def _GetKeyName(self, install_type, b_replace=True):
+    """Returns Chrome registry key name. """
+    if install_type == _winreg.HKEY_CURRENT_USER:
+      if b_replace:
+        return self.HKEY_USER.replace('ClientState', 'Clients')
+      else:
+        return self.HKEY_USER
+    elif install_type == _winreg.HKEY_LOCAL_MACHINE:
+      if b_replace:
+        return self.HKEY_LOCAL.replace('ClientState', 'Clients')
+      else:
+        return self.HKEY_LOCAL
+
+  def _GetWinLocalFolder(self, ftype=_CSIDL_COMMON_APPDATA):
+    """Returns the full path of the 'Local' folder on Windows.
+
+    Args:
+      ftype: Location to look up, which could vary based on installation type.
+
+    Returns:
+      A string containing the 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 GetInstallTypeString(self):
+    """Returns a string representing the type of Chrome installation."""
+    install_type = ({_winreg.HKEY_CURRENT_USER: 'user',
+                     _winreg.HKEY_LOCAL_MACHINE: 'system'}).get(
+                       self.GetChromeInstallType())    
+    return install_type
+
+  def GetChromeInstallType(self):
+    """Determines Chrome installation type."""
+    c_type = None    
+    key = self._OpenKey(_winreg.HKEY_CURRENT_USER,
+                        self.HKEY_USER.replace('ClientState', 'Clients'))
+    if key:
+      if self._GetKeyValue(key, 'pv'):
+        c_type = _winreg.HKEY_CURRENT_USER    
+    else:
+      key = self._OpenKey(_winreg.HKEY_LOCAL_MACHINE,
+                          self.HKEY_USER.replace('ClientState', 'Clients'))
+      if key:
+        if self._GetKeyValue(key, 'pv'):
+          c_type = _winreg.HKEY_LOCAL_MACHINE
+    if key:
+      _winreg.CloseKey(key)
+    return c_type
+
+  def GetChromeVersion(self):
+    """Returns current Chrome version."""
+    chrome_ver = ''
+    key_type = ({'user': _winreg.HKEY_CURRENT_USER,
+                 'system' : _winreg.HKEY_LOCAL_MACHINE}).get(self._type)
+    key = self._OpenKey(key_type, self._GetKeyName(key_type))
+    if key:
+      chrome_ver = self._GetKeyValue(key, 'pv')
+      _winreg.CloseKey(key)
+    return chrome_ver
+
+  def GetUninstallArgs(self, install_type=_winreg.HKEY_CURRENT_USER):
+    """Gets chrome's uninstall arguments from the registry.
+
+    Args:
+      install_type: Type of chrome installation (user or system). 
+
+    Returns:
+      A string containing the arguments if successful, otherwise empty string.
+    """
+    c_args = ''
+    key = self._OpenKey(install_type, self._GetKeyName(install_type, False))
+    if key:
+      c_args = self._GetKeyValue(key, 'UninstallArguments')
+      _winreg.CloseKey(key)    
+    return c_args
+
+  def GetChromeExePath(self):
+    """Gets Chrome's location, based on install type (i.e. user or system)."""
+    c_path = ''
+    f_id = (lambda _type: _type == 'user' and _CSIDL_COMMON_APPDATA or
+            _CSIDL_PROGRAM_FILESX86)(self._type)
+    if _PLATFORM == 'windows':
+      c_path = os.path.join(self._GetWinLocalFolder(f_id),
+                            'Google', 'Chrome',
+                            'Application', 'chrome.exe')
+    return c_path
+
+  def DeleteChromeRegEntries(self):
+    """Deletes chrome registry settings."""
+    if self._type == '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
+    else:
+      p_key = self.HKEY_LOCAL[: self.HKEY_LOCAL.rfind('\\')]
+      key_name = self.HKEY_LOCAL[self.HKEY_LOCAL.rfind('\\')+1:]
+      _type = _winreg.HKEY_LOCAL_MACHINE
+    key = self._OpenKey(_type, p_key, _winreg.KEY_ALL_ACCESS)
+    try:        
+      _winreg.DeleteKey(key, key_name)
+      key.Close()
+      return 0
+    except _winreg.error, err:
+      print "Warning: could not delete registry settings - ", err
+      return -1