munki/code/tools/pkginfo_hash_updater.py

#!/usr/bin/python
#
# Copyright 2010 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


"""Updates all pkginfo plists '(un)installer_item_hash' key.

This script will recursively update the '(un)installer_item_hash' key in all
pkginfo plists in the pkgsinfo directory of a Munki repo with a SHA-256 hash of
the corresponding package.

This script will run from Mac OSX or Linux alike, and it is safe to run more
than once on any pkginfo plist(s). However, it is recommended that you backup
your plists before running this script!

Dependencies:
- Linux: Python2.6 or higher, or 2.4/2.5 with python-plistlib manually
         installed: http://docs.python.org/library/plistlib.html
- Mac OSX: Python2.4 or higher.

Created on 2010-09-02.
"""

import optparse
import os
import hashlib
import sys
try:
  import plistlib
except ImportError:
  print ('ERROR: could not import module "plistlib." Please install it and '
         'try again.')
  sys.exit(1)


MUNKI_ROOT_PATH = '/var/www/munki/repo'
MUNKI_PKGS_DIR_NAME = 'pkgs'
MUNKI_PKGSINFO_DIR_NAME = 'pkgsinfo'


def GetHash(filename, hash_function):
  """Calculates the hashvalue of the given file with the given hash_function.
  
  Args:
    filename: The file name to calculate the hash value of.
    hash_function: The hash function object to use, which was instanciated
        before calling this function, e.g. hashlib.md5().
  
  Returns:
    The hashvalue of the given file as hex string.
  """
  if not os.path.isfile(filename):
    return None
  
  f = open(filename, 'rb')
  while 1:
    chunk = f.read(2**16)
    if not chunk:
      break
    hash_function.update(chunk)
  f.close()
  return hash_function.hexdigest()


def GetSHA256Hash(filename):
  """Returns the SHA-256 hash value of a file as a hex string."""
  hash_function = hashlib.sha256()
  return GetHash(filename, hash_function)


def AddHashesToPkginfoPlists(pkgsinfo_path, pkgs_path, update_existing=False):
  """Recursively updates plists' '(un)installer_item_hash' kay with pkg hash.
  
  Args:
    pkgsinfo_path: root dir to start updating from.
    pkgs_path: root dir where Munki pkgs live.
  """
  for f_path in os.listdir(pkgsinfo_path):
    if f_path.startswith("._") or f_path == ".DS_Store":
      # don't process these
      continue
    f_path = os.path.join(pkgsinfo_path, f_path)
    # if a directory is found, recursively call this function on that dir.
    if os.path.isdir(f_path):
      AddHashesToPkginfoPlists(f_path, pkgs_path)
    elif os.path.islink(f_path):
      print 'WARNING: symlinks not supported; skipping: %s' % f_path
      continue
    elif os.path.isfile(f_path):
      # read plist
      try:
        plist = plistlib.readPlist(f_path)
      except IOError, e:
        print 'WARNING: pkginfo plist failed to open: %s\n%s' % (f_path,
                                                                 str(e))
        continue
      
      updated_hash = False
      # generate the package path from the pkginfo plist
      # installer_item_location.
      if not 'installer_item_hash' in plist or update_existing:
          pkg_path = os.path.join(pkgs_path, plist['installer_item_location'])
          # display warning for items that cannot be found.
          if not os.path.isfile(pkg_path):
            print >> sys.stderr, ('WARNING: Installer item (%s) not found '
                                  'as specified in %s') % (pkg_path, f_path)
            continue
          # generate the package hash
          plist['installer_item_hash'] = GetSHA256Hash(pkg_path)
          updated_hash = True
      
      # check for and update uninstaller_item_hash if needed
      if 'uninstaller_item_location' in plist:
        if not 'uninstaller_item_hash' in plist or update_existing:
          pkg_path = os.path.join(pkgs_path,
                                  plist['uninstaller_item_location'])
          # display warning for items that cannot be found.
          if not os.path.isfile(pkg_path):
              print >> sys.stderr, ('WARNING: Uninstaller item (%s) not '
                                    'found as specified in %s') % (pkg_path,
                                                                   f_path)
              continue
          
          # generate the package hash
          plist['uninstaller_item_hash'] = GetSHA256Hash(pkg_path)
          updated_hash = True
          
      # write the plist file.
      if updated_hash:
        plistlib.writePlist(plist, f_path)
        print '- Wrote hash to plist: %s' % f_path


def main():
  usage = 'usage: %prog [options]'
  p = optparse.OptionParser(usage=usage)
  p.add_option('-r', '--munki_root', default=MUNKI_ROOT_PATH,
               help='Munki repo root path where pkginfo and pkgs dirs live; '
               'default "/var/www/munki/repo"')
  p.add_option('-p', '--pkgsinfo_dir_name', default=MUNKI_PKGSINFO_DIR_NAME,
               help='Munki pkgsinfo dir name; default "pkgsinfo".')
  p.add_option('-f', '--pkgs_dir_name', default=MUNKI_PKGS_DIR_NAME,
               help='Munki packages dir name; default "pkgs".')
  p.add_option('-u', "--update_existing", action='store_true',
                help='Update existing hashes.')
  options, unused_arguments = p.parse_args()
  
  pkgsinfo_path = os.path.join(options.munki_root, options.pkgsinfo_dir_name)
  pkgs_path = os.path.join(options.munki_root, options.pkgs_dir_name)
  if not os.path.isdir(pkgsinfo_path) or not os.listdir(pkgsinfo_path):
    print 'Pkgsinfo directory not found or is empty: %s' % pkgsinfo_path
  elif not os.path.isdir(pkgs_path) or not os.listdir(pkgs_path):
    print 'Pkgs directory not found or is empty: %s' % pkgs_path
  else:
    AddHashesToPkginfoPlists(pkgsinfo_path, pkgs_path,
                             options.update_existing)
    print '\nYou must run makecatalogs to update catalogs with pkginfo changes.'


if __name__ == '__main__':
  main()