# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Library General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
# Copyright 2007, 2008, 2009 Red Hat, Inc. 
# Written by skvidal@fedoraproject.org, wwoods@redhat.com

# TODO: gpgcheck downloaded pkgs
# TODO: probably have some defensive checks before we start anything
#       system sanity checks, depsolver checks against current rpmdb

import sys
import os
import os.path
import time
import shutil
import ConfigParser
from iniparse.compat import ConfigParser as yumConfigParser
from urlgrabber.grabber import URLGrabber, URLGrabError

import preupgrade.misc as misc
import preupgrade.dev as dev
from preupgrade.error import PUError

import yum, rpm
import yum.Errors
import rpmUtils.arch
from yum.misc import checksum
from yum.parser import varReplace
from yum.constants import *
from yum.Errors import RepoError

__version__ = "1.1.0"

# FIXME: this should be configurable somewhere.
if os.path.exists("releases.txt"):
    releasesfn = "releases.txt"
elif os.path.exists("data/releases.txt"):
    releasesfn = "data/releases.txt"
else:
    releasesfn = "http://mirrors.fedoraproject.org/releases.txt"

runfile = "/var/run/preupgrade"
repofile = "/var/run/preupgrade.repo"

class PreUpgrade(yum.YumBase):
    def __init__(self, release=None, configfile=releasesfn):
        yum.YumBase.__init__(self)
        self.release_map = {}
        self.configfile = configfile
        self.bootpath='/boot/upgrade'
        self.image_csums = {} # filename = (csumtype, csum)
        # Noisy if these plugins don't exist, but doesn't raise an exception
        # TODO: substitute this for yum 3.2.21 and higher
        #self.preconf.enabled_plugins = ['blacklist','whiteout']
        self._getConfig(enabled_plugins=['blacklist','whiteout'])
        self.conf.disable_excludes = ['all']
        self.orig_enabled_repos = self.repos.listEnabled()
        if release:
            self.setup(release)

    def upgrade_available(self, want_stable=True):
        '''If there is a newer release available in the release map, return its
        name/key. If there are multiple newer releases, it returns the newest
        one (the one with the highest version number).
        Otherwise returns None.
        You can set stable to False if you want to consider non-stable releases
        as well (but right now this always gets you Rawhide).'''
        current = self.conf.yumvar['releasever']
        target_ver = current
        target_key = None
        for release_key,release in self.release_map.iteritems():
            # skip anything that's not marked preupgrade-ok
            if not self.release_preupgrade_ok(release):
                continue
            # skip unstable releases (unless the user has requested otherwise)
            if want_stable and not self.release_is_stable(release):
                continue
            # FIXME: unless we return a list or something, we'll always return 
            # Rawhide if stable=False, since it's always available and always
            # considered newer than whatever you've got.
            ver = release['version']
            if ver > target_ver: 
                target_ver = ver
                target_key = release_key
        return target_key

    def setup(self, release_key, fake=False, download_progressbar=None):
        '''Given a release_key, will configure repos to upgrade to that
        release. Must be run after retrieve_release_info().'''
        if fake and (release_key not in self.release_map):
            self.release_map = {release_key:{'baseurl':'http://fake.url/',
                                             'version':'0'}}
        if self.release_map.has_key(release_key):
            rel = self.release_map[release_key]
            self.pu_release = release_key
            self.target_version = rel.get('version')
            self.pu_baseurl = rel.get('baseurl')
            self.pu_mirrorlist = rel.get('mirrorlist')
            self.pu_installurl = rel.get('installurl')
            self.pu_installmirrorlist = rel.get('installmirrorlist')
        else:
            raise PUError, 'Release does not match any in configuration.'
        # Set up the repos
        if self.pu_installurl or self.pu_installmirrorlist:
            # We've got a separate install repo - set both repos up
            self.setup_main_repo(download_progressbar)
            self.setup_install_repo(download_progressbar)
        else:
            # Only one repo - set it up as the install repo
            self.pu_installurl = self.pu_baseurl
            self.pu_installmirrorlist = self.pu_mirrorlist
            self.setup_install_repo(download_progressbar)
        # Attempt to set up any other enabled repo on the system
        self.setup_new_repos(download_progressbar)
        # Complete repo setup
        self.complete_repo_setup()

    def retrieve_release_info(self, configfile=None):
        if not configfile:
            configfile = self.configfile
        releases = {}
        cp = ConfigParser.ConfigParser()
        if '://' in configfile:
            grabber = URLGrabber()
            configdata = grabber.urlopen(configfile)
        else:
            configdata = open(configfile)
        cp.readfp(configdata)
        for sect in cp.sections():
            rel = dict(cp.items(sect))
            if ('baseurl' in rel) or ('mirrorlist' in rel):
                releases[sect] = rel
                
        self.release_map = releases
        return releases
    
    def __release_key_is_true(self,release,key):
        '''Return True if the given key in the given release dict is a True
        value - something like 'true' or 'yes' or '1'.'''
        r = False
        val = release.get(key)
        if type(val) == str:
            b = val.lower()
            if b == '1' or b == 'true' or b == 'yes':
                r = True
        return r

    def release_is_stable(self,release):
        '''Return True if the given release is a stable release'''
        # this is an abstraction - just in case we change the structure of the
        # 'release' data.
        return self.__release_key_is_true(release,'stable')

    def release_preupgrade_ok(self,release):
        '''Return True if the given release is OK for preupgrade'''
        return self.__release_key_is_true(release,'preupgrade-ok')

    def _setup_repo(self, id, url, urltype, download_progressbar=None):
        '''Set up an additional repo to fetch packages from during preupgrade.
        id:  string representing the repo (e.g. "updates")
        url: string, url of the repo/mirrorlist/metalink
        urltype: string, either 'baseurl', 'mirrorlist', or 'metalink'. 
        download_progressbar: callback object for the created repo object

        Adds the new YumRepository object to self.repos, and returns it.

        Note that this does *NOT* enable the repo - to do that, you must
        run self.repos.enableRepo(id) yourself.
        '''
        repo = yum.yumRepo.YumRepository(id)
        repo.name = id
        repo.baseurl = []
        repo.mirrorlist = None 
        repo.metalink = None
        repo.basecachedir = '/var/cache/yum/'
        repo.enablegroups = 1
        # FIXME ugh REAL LOGGING DAMMIT
        print "%s (%s) " % (id, urltype)
        print "  url: %s" % url
        if download_progressbar:
            repo.setCallback(download_progressbar)

        url = varReplace(url, self.conf.yumvar)
        print "  now: %s" % url
        if urltype == 'mirrorlist':
            repo.mirrorlist = url
        elif urltype == 'metalink':
            repo.metalink = url
        elif urltype == 'baseurl':
            repo.baseurl = [url]
        else:
            raise ValueError, "urltype must be mirrorlist, metalink, or baseurl"

        self.repos.add(repo)
        self.repos.doSetup(thisrepo=repo.id)
        return self.repos.findRepos(repo.id)[0]
        
    def setup_new_repos(self, download_progressbar=None):
        '''Setup repo objects for newer versions of all the enabled repos that
        have $version in 'em. This will only work for final releases.

        For unstable releases (Rawhide, pre-releases, etc) we enable any
        repo that ends with '-rawhide'.'''
        # TODO: This won't give us weirdo one-off repos like the "newkey"
        # repos from F8/F9 - we need a better way for that.

        # setup a yumvar that has $version = $target_version
        new_yumvar = self.conf.yumvar
        new_yumvar['releasever'] = str(self.target_version)
        # Disable everything so we can start fresh
        self.repos.disableRepo('*')

        # Rawhide is kinda special, so here's a special case for it.
        # Unstable (pre-release) trees are basically rawhide snapshots, so they
        # get the same treatment.
        if self.pu_release == 'Rawhide' or not \
            self.release_is_stable(self.release_map[self.pu_release]):
            # Make sure we have an install repo set up. We need that!
            if not hasattr(self,'instrepo'):
                self.setup_install_repo(download_progressbar)
            # Re-enable our main repo - or fall back to instrepo
            if hasattr(self,'mainrepo'):
                self.mainrepo.enable()
            else:
                self.instrepo.enable()
            # Look for matching repos that end with '-rawhide'
            orig_repo_idlist = [r.id for r in self.orig_enabled_repos]
            self.orig_enabled_repos = []
            for r in self.repos.findRepos('*-rawhide'):
                if r.id[:-8] in orig_repo_idlist:
                    self.orig_enabled_repos.append(r)
                    # TODO: remove original repo

        # Now let's enable the new version of all active repos
        for r in self.orig_enabled_repos:
            # Unfortunately we have to reparse the config file ourselves to
            # increment $version. So let's do that.
            # Read the config file for this repo
            cfg = yumConfigParser()
            cfg.read(r.repofile)
            # Find the section for this particular repo
            if not cfg.has_section(r.id):
                # FIXME non-fatal error message here..
                continue
            repoconfig = {'id': 'preupgrade-%s' % r.id,
                          'download_progressbar': download_progressbar}
            # Read the config items
            for k,v in cfg.items(r.id):
                if k in ('metalink', 'mirrorlist', 'baseurl'):
                    # found our URL!
                    repoconfig['urltype'] = k
                    repoconfig['url'] = varReplace(v, new_yumvar)
                    newrepo = self._setup_repo(**repoconfig)
                    try:
                        if newrepo.ready():
                            self.repos.enableRepo(newrepo.id)
                        # FIXME else: raise an error or something
                    except RepoError:
                        id = repoconfig['id']
                        print "Can't set up new repo %s - removing" % id
                        self.repos.delete(id)
                    break

    def setup_install_repo(self, download_progressbar=None):
        repoconfig = {'id': 'preupgrade',
                      'download_progressbar': download_progressbar}
        # install image repo (Fedora)
        if self.pu_installmirrorlist:
            repoconfig['urltype'] = 'mirrorlist'
            repoconfig['url'] = self.pu_installmirrorlist
        if self.pu_installurl:
            repoconfig['urltype'] = 'baseurl'
            repoconfig['url'] = self.pu_installurl
        instrepo = self._setup_repo(**repoconfig)
        self.pu_installmirrorlist = instrepo.mirrorlist
        self.pu_installurl = instrepo.baseurl
        
        # save repo object
        self.instrepo = instrepo
        # This is a little dicey, using private stuff from yumRepo. Oh well.
        self.instgrab = self.instrepo._getgrab()

    def setup_main_repo(self, download_progressbar=None):
        repoconfig = {'id': 'preupgrade-main',
                      'download_progressbar': download_progressbar}
        # main repo (Everything)
        if self.pu_mirrorlist:
            repoconfig['urltype'] = 'mirrorlist'
            repoconfig['url'] = self.pu_mirrorlist
        if self.pu_baseurl:
            repoconfig['urltype'] = 'baseurl'
            repoconfig['url'] = self.pu_baseurl
        mainrepo = self._setup_repo(**repoconfig)
        self.pu_mirrorlist = mainrepo.mirrorlist
        self.pu_baseurl = mainrepo.baseurl
        self.repos.enableRepo(mainrepo.id)
        self.mainrepo = mainrepo

    def complete_repo_setup(self):
        '''A few extra bits to configure after the repos are set up.'''
        # turn on caching
        self.conf.keepcache = 1
        # but force a check of the cached data
        self._cleanFiles(['cachecookie','mirrorlist.txt'],'cachedir','metadata')
        # write self.pu_release to a file that resuming_run can pick up
        rf = open(runfile,"w")
        rf.write(self.pu_release + "\n")
        rf.close()
        # write repo config to a file so we can easily restore it later 
        rf = open(repofile,"w")
        for repo in [self.instrepo] + self.repos.listEnabled():
            # adapted from YumRepository.write()
            yc = yumConfigParser()
            yc.add_section(repo.id)
            for k,v in repo.iteritems():
                option = repo.optionobj(k)
                if v != option.default:
                    yc.set(repo.id, k, option.tostring(v))
            yc.write(rf)
            rf.write('\n')
        rf.close()

    def clear_all_repos(self):
        self.clear_install_repo()
        # clear all other enabled repos
        for r in self.repos.listEnabled():
            self.repos.disableRepo(r.id)
            self.repos.delete(r.id)

    def clear_install_repo(self):
        '''Clear out previously-set up repos (so we can retry)'''
        self.instrepo = None
        id = 'preupgrade'
        if self.repos.findRepos(id):
            self.repos.disableRepo(id)
            self.repos.delete(id)
    
    def retrieve_treeinfo(self):
        # take our repo
        # find a .treeinfo in them
        # download it
        # parse it
        # store relative path to the right kernel and initrd
        # there can be only one
        # info we need:
        # initrd.img path
        # kernel path
        # stage2.img path
        # min image
        # example:
        #   http://download.fedora.redhat.com/pub/fedora/linux/development/i386/os/.treeinfo
        tif = self.instrepo.grab.urlopen('/.treeinfo')
        print "Fetched treeinfo from %s" % tif.url
        #cracked out ppc issues, etc
        archlist = [self.conf.yumvar['basearch']]
        if misc.isParaVirt():
            archlist = ['xen']
        if rpmUtils.arch.getCanonArch().startswith('ppc64'):
            archlist = [rpmUtils.arch.getCanonArch(), self.conf.yumvar['basearch']]
        
        cp = ConfigParser.ConfigParser()
        try:
            cp.readfp(tif)
        except ConfigParser.MissingSectionHeaderError:
            # Generally this means we failed to download the file
            raise PUError, "Failed to download installer metadata"
        except ConfigParser.Error:
            raise PUError, "Failed to parse installer metadata"
        
        ts = float(cp.get('general','timestamp','0'))
        if ts > 0:
            print "treeinfo timestamp: " + time.asctime(time.localtime(ts))
        else:
            print "treeinfo timestamp missing (possibly malformed)!"
        
        mysection = None
        for arch in archlist:
            if cp.has_section('images-%s' % arch):
                mysection = 'images-%s' % arch
                break
        
        if not mysection:
            raise PUError, "Cannot find any kernel/initrd section for this arch in treeinfo"
            
        try:
            self.kernel = cp.get(mysection, 'kernel')
            self.initrd = cp.get(mysection, 'initrd')
            if cp.has_section('stage2'):
                self.mainimage = cp.get('stage2', 'mainimage')
            else:
                self.mainimage = 'images/stage2.img'
        except ConfigParser.NoOptionError, e:
            raise PUError, "Invalid treeinfo: %s" % str(e)
        
        if cp.has_section('checksums'):
            for opt_fn in cp.options('checksums'):
                (csum_type, csum) = cp.get('checksums', opt_fn).split(':')
                self.image_csums[opt_fn] = (csum_type, csum)
 
        tif.close()
        
    
    def _retrieve_file(self, fileinfo, targetdir=''):
        """ for returning files from our important repo, but not package or metadata files"""
        if not targetdir:
            targetdir = self.bootpath
        
        if not os.path.exists(targetdir):
            os.makedirs(targetdir)

        # for each file:
        # check if we have the filesize and check it out
        # if it is okay, then continue
        # if it is not okay then unlink it 
        # check for space to download it
        
        # local name
        item_fname = os.path.basename(fileinfo)
        local = targetdir + '/' + item_fname
        # remote size
        tmp = self.instgrab.urlopen(fileinfo)
        mysize = tmp.hdr.dict.get('content-length','-1')
        tmp.close()

        # Check to see if we've already got the right file
        if os.path.exists(local):
            local_stat = os.stat(local)
            if int(local_stat.st_size) == int(mysize):
                # It's the right size - check checksum
                if self.image_csums.has_key(fileinfo):
                    csumtype, csum = self.image_csums[fileinfo]
                    if csum == checksum(csumtype, local):
                        print "%s checksum OK" % local
                        return True
                    else:
                        print "%s checksum failed - removing" % local
                        os.unlink(local)
            else:
                os.unlink(local)

        # if we are this far we're going to try to download it
        # check for available disk space
        dirstat = os.statvfs(targetdir)
        if (dirstat.f_bavail * dirstat.f_bsize) <= long(mysize):
            raise PUError,("Not enough space in %s to download %s." % \
                    (targetdir, item_fname))

        try:
            self.instrepo._getFile(relative=fileinfo, local=local)
        except URLGrabError, e:
            os.unlink(local)
            raise PUError, "failure downloading %s: %s" % (item_fname, e)
            
        if self.image_csums.has_key(fileinfo):
            csumtype, csum = self.image_csums[fileinfo]
            if csum != checksum(csumtype, local):
                os.unlink(local)
                raise PUError, "checksum failure downloading %s" % item_fname

        return True

    def remove_boot_files(self):
        '''Remove any previously-downloaded boot files.'''
        def rmtree_handler(func,path,exc_info):
            (e_type, e_value, e_traceback) = exc_info
            raise PUError, "Failed to remove %s: %s" % (path,e_value)
        if os.path.exists(self.bootpath):
            shutil.rmtree(self.bootpath,ignore_errors=False,\
                          onerror=rmtree_handler)

    def remove_boot_entry(self):
        kernel = os.path.join(self.bootpath,
                            os.path.basename(getattr(self,'kernel','vmlinuz')))
        os.system('/sbin/grubby --remove-kernel="%s"' % kernel)

    def remove_repo_cache(self):
        '''Remove the repo caches, if they exist.
        Will throw PUError if the repos are not yet set up.'''
        # Since we inherit from YumBase we've got the cleanXXX methods:
        #   self.clean{Headers,Packages,Sqlite,Metadata}
        # FIXME this is kind of a dumb way to see if we've set up the repos
        if not hasattr(self,'instrepo'):
            # we need to set up the repos before we can clean them!
            raise PUError, "Can't clear repo caches until repos are set up"
        self.cleanHeaders()
        self.cleanPackages()
        self.cleanSqlite()
        self.cleanMetadata()
        # For good measure, try to clear out any remaining files
        cachedir = "/var/cache/yum/preupgrade"
        for f in os.listdir(cachedir):
            if f.endswith(".rpm"):
                try:
                    os.unlink(os.path.join(cachedir,f))
                except OSError:
                    pass
        
    def resuming_run(self):
        '''If there is cached data from a previously-interrupted run of
        preupgrade, return the release name (self.pu_release)'''
        releasename = None
        try:
            f = open(runfile)
            releasename = f.readline().strip()
        except IOError:
            pass
        return releasename

    def fake_setup_repos(self):
        repoconfig = {}
        yc = yumConfigParser()
        yc.read(repofile)
        for s in yc.sections():
            id = yc.get(s,'name')
            repoconfig['id'] = id
            repoconfig['urltype'] = 'baseurl'
            repoconfig['url'] = 'http://fake.url/%s' % id
            # instrepo needs to be set for remove_repo_cache to work
            self.instrepo = self._setup_repo(**repoconfig)
            self.repos.enableRepo(id)

    def clear_incomplete_run(self):
        '''Clear all the saved data for an incomplete run.'''
        release = self.resuming_run()
        assert release != None
        self.fake_setup_repos()
        self.remove_repo_cache()
        self.remove_boot_files()
        self.remove_boot_entry()
        os.unlink(runfile)
        self.clear_all_repos()

    def retrieve_critical_boot_files(self):
        # relative tmp FIXME XXX
        
        for item in (self.kernel, self.initrd):
            self._retrieve_file(item)
        
    def retrieve_non_critical_files(self,targetdir=''):
        self._retrieve_file(self.mainimage,targetdir)
        return os.path.basename(self.mainimage)

    def generate_kickstart(self,targetfile=None,extra_cmds=""):
        '''Generate a simple kickstart to automate the upgrade.
        Returns a tuple of (bootdev,path) like bootpath_to_anacondapath()'''
        lang = misc.getlang()
        keymap = misc.getkeymap()
        if not targetfile:
            targetfile = os.path.join(self.bootpath,'ks.cfg')
        kernel = os.path.join(self.bootpath,
                            os.path.basename(getattr(self,'kernel','vmlinuz')))
        f = open(targetfile,'w')
        ks = ['# ks.cfg generated by preupgrade',
              'lang %s' % lang,
              'keyboard %s' % keymap,
              'bootloader --upgrade --location=none',
              'upgrade',
              'reboot',
              '%s' % extra_cmds,
              '\n%post',
              'grubby --remove-kernel=%s' % kernel,
              'rm -rf %s /var/cache/yum/preupgrade*' % self.bootpath,
              '%end\n']
        if float(self.target_version) > 10.0:
            ks[4] += ' --root-device=UUID=%s' % dev.device_to_uuid('/dev/root')
        f.write('\n'.join(ks))
        f.close()
        return dev.bootpath_to_anacondapath(targetfile,UUID=True)
        
    def get_download_pkgs(self):
       return map(lambda txmbr: txmbr.po,
               self.tsInfo.getMembersWithState(output_states=TS_INSTALL_STATES))

    def get_download_size(self,pkgs):
        localsize = 0
        totalsize = 0
        for pkg in pkgs:
            try:
                size = int(pkg.size)
                totalsize += size
                try:
                    if pkg.verifyLocalPkg():
                        localsize += size
                except:
                    pass
            except e:
                print "Error calculating download size for %s: %s" % \
                            (str(e),str(pkg))
        return(totalsize - localsize)

    def get_upgrade_size(self):
        '''Returns the (approximate) amount of space needed for anaconda to
        complete the upgrade, in bytes.'''
        # Why, yes, this is a made-up number. And yet it's still more accurate
        # than the results of the test transaction.
        return 500 * 1024 * 1024

    # XXX currently unused. It takes forever, the callback doesn't work, and
    # the results don't match anaconda's transaction test so it tends to give
    # us false positives and then the upgrade fails.
    def transaction_test_diskspace(self,callback=None):
        '''Run a test transaction to ensure that we'll have enough disk
        to complete the upgrade.
        
        Returns a dict of mountpoint:spaceneeded pairs.
        An empty dict indicates success.
        '''
        testcb = yum.RPMTransaction(self,test=True)
        if callback:
            testcb.display = callback
        else:
            testcb.display = yum.rpmtrans.SimpleCliCallBack
        # save our dsCallback
        dscb = self.dsCallback
        self.dsCallback = None # ssshhhhhh
        # Turn off all the problem filters *but* diskspace
        self.ts.setProbFilter(~rpm.RPMPROB_FILTER_DISKSPACE)
        # XXX do we need setColor (see anaconda yuminstall.py) here?
        self.populateTs(keepold=0)
        # Actually test the transaction
        self.ts.addTsFlag(rpm.RPMTRANS_FLAG_TEST)
        problems = self.ts.run(testcb.callback, '')
        # clean up, and restore the depcheck callback
        del self.ts
        self.dsCallback = dscb
        spaceneeded = {}
        if problems:
            for (desc, (problem_type, mount, need)) in problems:
                if problem_type == rpm.RPMPROB_DISKSPACE:
                    spaceneeded[mount] = need
        return spaceneeded

    def generate_repo(self, dir, comps=None, callback=None):
        self.gather_downloaded_packages(dir)
        misc.generate_repodata(dir,comps,callback)

    def gather_downloaded_packages(self, dir):
        '''Hardlink all downloaded packages to the given directory'''
        for r in self.repos.listEnabled():
            if r.pkgdir:
                fullpkgdir = os.path.join(r.basecachedir,r.pkgdir)
                for p in os.listdir(fullpkgdir):
                    newfile = os.path.join(dir,p)
                    try:
                        os.remove(newfile)
                    except:
                        pass
                    os.link(os.path.join(fullpkgdir,p),newfile)

    def return_available_versions(self, cur_ver, beta_ok=False):
        versions = []
        for name, releasedata in self.release_map.iteritems():
            # Skip releases that aren't newer than this one
            if cur_ver >= float(releasedata.get('version')):
                continue
            # Skip releases that aren't preupgrade-ok
            if not self.release_preupgrade_ok(releasedata):
                continue
            # If beta_ok get all releases, otherwise just the stable ones
            if beta_ok or self.release_is_stable(releasedata):
                versions.append(name)
        return versions
        
    def figure_out_what_we_need(self):
        # grab the comps file, we'll need it later for anaconda
        if self.comps.compscount == 0:
            raise PUError, "Cannot retrieve group lists from any repository"
        # also try to get the gzipped comps file if it's there so that
        # anaconda doesn't have to later.  this is kind of ugly (#439232)
        for r in self.repos.listGroupsEnabled():
            try:
                r.retrieveMD('group_gz')
            except:
                pass

        # update all
        self.update()
        # NOTE: anaconda sets installonly_limit to 1. Do we need that?
        # build a transaction
        self.buildTransaction()

    def add_boot_target(self,extra_args=''):
        mykernel  = self.bootpath  + '/' + os.path.basename(self.kernel)
        myinitrd = self.bootpath + '/' + os.path.basename(self.initrd)
        arch = self.conf.yumvar['basearch']
        argstr = "preupgrade"
        if extra_args: 
            argstr += " " + extra_args
        title="Upgrade to %s" % self.pu_release
        # yaboot (ppc) doesn't support spaces in titles
        if arch.startswith('ppc'):
            title = 'preupgrade' 
        #FIXME - do something with return codes below here :)

        misc.grubby(mykernel, myinitrd, title=title, args=argstr)
        if arch.startswith('ppc'):
            # yaboot doesn't like spaces in titles much
            # ppc machines need to run ybin to activate changes
            os.system('/sbin/ybin > /dev/null')

    def change_boot_target(self): 
        arch = self.conf.yumvar['basearch']
        if arch.startswith('ppc'):
            # ppc machines need to run ybin --bootonce
            os.system('/sbin/ybin --bootonce preupgrade > /dev/null')
        else:
            # everyone else uses grub --once
            misc.grubby_make_next_boot_only()
            
    def _main(self):
        '''Basically an example of how a run is supposed to proceed.'''
        self.setup('Rawhide')
        # In this example, we always clear old data rather than reusing it
        self.remove_repo_cache()
        self.remove_boot_files()
        self.figure_out_what_we_need()
        downloadpkgs = []
        downloadpkgs = map(lambda txmbr: txmbr.po,
                           self.tsInfo.getMembersWithState(output_states=TS_INSTALL_STATES))
        del(self.tsInfo)
        
       
        problems = self.downloadPkgs(downloadpkgs) 

        # find tree info data 
        self.retrieve_treeinfo()
        
        # grab kernel and initrd and stage2
        try:
            self.retrieve_critical_boot_files()
        except PUError, e:
            print 'Fatal: %s' % e
            sys.exit(1)

        try:            
            self.retrieve_non_critical_boot_files()        
        except PUError:
            pass
        
        # make it bootable
        self.add_boot_target()
        self.change_boot_target()