#!/usr/bin/python
#
#    addrpm.py  Adds RPM(s) to a Red Hat Tree
#    Copyright (C) 2004 NC State University
#    Written by Jack Neely <jjneely@pams.ncsu.edu>
#               Seth Vidal <skvidal@duke.edu>
#
#    SDG
#
#    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 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
#
import os
import sys
import os.path
import shutil
import rpm
import getopt

debug = 0
testmode = 0
addnew = 0

#where to put the old rpms
trash = "/tmp/"

# What version of RPM are we dealing with?
userpm404 = 0

if sys.version[0] == "2" and not userpm404:
    ts = rpm.TransactionSet("", (rpm._RPMVSF_NOSIGNATURES or rpm.RPMVSF_NOHDRCHK or rpm._RPMVSF_NODIGESTS or rpm.RPMVSF_NEEDPAYLOAD))

def usage():
    print "Usage:" + sys.argv[0] + " -e dir -u dir -v ver -r arch [listofrpms]"
    print " Required arguments:"
    print "    -v, --ver [version of distro]"
    print "    -r, --arch [architecture of distro]"
    print "    -e, --tree [where the top of the tree is]"
    print "    -u, --updatedir [subdir for updates ex: other-pkgs/updates"
    print
    print " Optional arguments:"
    print "    -p, --product [Name of product, default='RedHat'"
    print "    -t, --testmode\tRun in test mode"
    print "    -a, --addnew\tAdd rpms to tree if they are not already there"
    print "    -d, --debug\t\tRun in debug mode"
    print "    -h, --help\t\tPrint this message"
    print
    sys.exit(1)

def log(mess):
    if debug:
        print mess


def getHeader(filename):
    """ Read a rpm header. """
    if debug:
        print "Reading " + filename
    fd = os.open(filename, os.O_RDONLY)
    if sys.version[0] == "2" and not userpm404:
        h = ts.hdrFromFdno(fd)
        issrc = h[rpm.RPMTAG_SOURCEPACKAGE]
    else:
        (h, issrc) = rpm.headerFromPackage(fd)
    os.close(fd)

    return (h, issrc)


# sanity checks on the list list.  Can preppend root to aid in
# absolute path resolving.
# Make damn sure this looks like an RPM....

def checkFiles(filelist, rt=""):
    files = []
    for file in filelist:
        path = os.path.abspath(os.path.join(rt,file))
        log(path)
        if not os.access(path, os.R_OK):
            print "Could not find %s" % file
            sys.exit(1)
        root, ext = os.path.splitext(path)
        if ext == ".rpm":
            files.append(path)
    return files

# Returns a dict indexed by filename as given in the rpmlist list
def getDictByFile(rpmlist):
    headers = {}
    for file in rpmlist:
        header, isSource = getHeader(file)
        if isSource:
            print "What are you thinking?  %s is a source RPM!" % file
            sys.exit(1)
        headers[file] = header
    return headers

# Returns a dict indexed by package name.
# The dict contains a tuple of (header, filename)
def getDictByName(rpmlist):
    headers = {}
    for file in rpmlist:
        header, isSource = getHeader(file)
        name = header[rpm.RPMTAG_NAME]
        arch = header[rpm.RPMTAG_ARCH]
        if isSource:
            print "What are you thinking?  %s is a source RPM!" % file
            sys.exit(1)
        headers[(name,arch)] = (header, file)

    return headers

# Add package to tree
# pkg_fn is a absolute path to the rpm to be added

def addPackage(pkg_fn, install_tree, updates_dir, updatereldir):
    basename = os.path.basename(pkg_fn)
    if not testmode:
        shutil.copy2(pkg_fn, updates_dir)
        os.symlink(os.path.join(updatereldir,basename),
                   os.path.join(install_tree,basename))
    else:
        print "Copying %s to %s" % (pkg_fn, updates_dir)
        print "Symlinking %s as %s" % (os.path.join(updatereldir,basename),
                                       os.path.join(install_tree,basename))

# Remove package from tree
# pkg_fn is an absolute path to the package in the install tree to remove

def removePackage(pkg_fn):
    basename = os.path.basename(pkg_fn)
    #where is the dir of the file we're checking
    rpmdirname = os.path.dirname(pkg_fn)
    #where are we now?
    startdir = os.getcwd()
    if not testmode:
        if os.path.islink(pkg_fn):
            #its a symlink, its probably relative. - lets go there :)
            os.chdir(rpmdirname)
            pkglinksource = os.path.abspath(os.readlink(pkg_fn))
            #and back.
            os.chdir(startdir)    
            shutil.copy2(pkglinksource, trash)
            os.unlink(pkglinksource)
            os.unlink(pkg_fn)
        else:
            shutil.copy2(pkg_fn, trash)
            os.unlink(pkg_fn)
    else:
        if os.path.islink(pkg_fn):
            #its a symlink, its probably relative. - lets go there :)
            os.chdir(rpmdirname)
            pkglinksource = os.path.abspath(os.readlink(pkg_fn))
            #and back.
            os.chdir(startdir)
            print 
            print "Copying %s to %s" % (pkglinksource, trash)
            print "Removing %s and %s" % (pkglinksource,pkg_fn)
        else:
            print
            print "Copying %s to %s" % (pkg_fn, trash)
            print "Removing %s" % pkg_fn        


def main():
    """This utility adds RPMs to a Red Hat Install Tree.
       Used to maintain the NCSU Realm Kit for Red Hat
       Linux."""

    global addnew
    global testmode
    global debug
    args = sys.argv[1:]

    try:
        optlist, rpmlist = getopt.getopt(args, 'hdtv:r:au:e:p:', ['help', 'debug', 'testmode', 'ver=', 'arch=', 'add', 'updatedir=', 'tree=', 'product='])
    except getopt.error:
        usage()

    archarg=0
    verarg=0
    treearg=0
    updatedirarg=0
    product = "RedHat"
    trees = "/tmp/dumb"
    updatesubdir = "/other/path"
    
    
    for o, a in optlist:
        if o in ("--product", "-p"):
            product = a
        if o in ("--help", "-h"):
            usage()
        if o in ("--testmode", "-t"):
            testmode=1
        if o in ("--ver", "-v"):
            version=a
            verarg=1
        if o in ("--arch", "-r"):
            arch=a
            archarg=1
        if o in ("--tree", "-e"):
            trees=a
            treearg=1
        if o in ("--updatedir", "-u"):
            import string
            updatesubdir = a
            slashindex = string.find(updatesubdir,"/")
            log(slashindex)
            if slashindex == 0:
                updatedirarg=0
                print "Error: Updates Dir must not start with a /"
            else:
                updatedirarg=1
            
        if o in ("--add", "-a"):
            addnew=1
        if o in ("--debug", "-d"):
            debug=1
            
    if archarg+verarg+treearg+updatedirarg != 4:
        print "Not enough or bad args"
        usage()

    # Define the location of install tree and updated dir
    install_tree = os.path.join(trees, version, arch, product, "RPMS/")
    updatereldir = os.path.join("../../", updatesubdir)
    updates_dir = os.path.join(trees, version, arch, updatesubdir)
    install_tree = os.path.normpath(install_tree)
    updates_dir = os.path.normpath(updates_dir)
    
    if not os.access(install_tree, os.R_OK | os.X_OK) or not os.path.isdir(install_tree):
        print "You do not have access to %s" % install_tree
        sys.exit(1)

    if not os.access(updates_dir, os.R_OK | os.X_OK) or not os.path.isdir(updates_dir):
        print "You do not have access to %s" % updates_dir
        sys.exit(1)

    log("tree path = %s" % install_tree)
    log("updates path = %s" % updates_dir)
        
    # Resolve absolute path names
    log("Checking new RPMS")
    newrpms = rpmlist
    log("Getting headers for new RPMS")
    newrpms_headers = getDictByFile(newrpms)
    log("Getting headers for RPMS in install tree.")
    filelist = checkFiles(os.listdir(install_tree), install_tree)
    current_rpms = getDictByName(filelist)

    # Do the packages upgrades
    for pkg in newrpms:
        log("Working on " + pkg)
        basename = os.path.basename(pkg)
        newheader = newrpms_headers[pkg]
        pkg_name = newheader[rpm.RPMTAG_NAME]
        pkg_arch = newheader[rpm.RPMTAG_ARCH]
        if current_rpms.has_key((pkg_name,pkg_arch)):
            header = current_rpms[(pkg_name,pkg_arch)][0]
            fn = current_rpms[(pkg_name,pkg_arch)][1]
            compare = rpm.versionCompare(header, newheader)
            if compare == 1:
                log("Old %s.%s package is newer.  Not adding." % (pkg_name, pkg_arch))
            elif compare == 0:
                log("Both %s.%s packages are equal. Not adding." % (pkg_name, pkg_arch))
            elif compare == -1:
                log("Adding %s.%s package." % (pkg_name, pkg_arch))
                addPackage(pkg, install_tree,
                           os.path.join(updates_dir, pkg_arch), 
                           os.path.join(updatereldir, pkg_arch))
                log("Removing old %s.%s package." % (pkg_name, pkg_arch))
                removePackage(fn)
                print "%s.%s updated." % (pkg_name, pkg_arch)
            else:
                print "Unexpected result from rpm.versionCompare()"
                print "Self destructing in 5 seconds."
                print "Have a nice day."
                sys.exit(1)
        else:
            if addnew:
                print "Warning: %s.%s is not present in the install tree. Adding anyway." % (pkg_name, pkg_arch)
                addPackage(pkg, install_tree,
                           os.path.join(updates_dir, pkg_arch), 
                           os.path.join(updatereldir, pkg_arch))
                print "%s.%s updated." % (pkg_name, pkg_arch)
            else:
                print "Warning: %s.%s is not present in the install tree. Ignoring." % (pkg_name, pkg_arch)

    print "Done adding RPMs!"


if __name__ == "__main__":
    main()