path: root/presto-utils/deltarpmd.py

#!/usr/bin/python -t
# -*- mode: Python; indent-tabs-mode: nil; -*-
#
# 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
#
# Copyright (c) 2008 Jonathan Dieter
# Some portions copyright others

import errno, os, sys
import fnmatch, re
import rpmUtils.transaction, rpmUtils.miscutils
import commands
import string
import getopt
from dumpMetadata import byteranges
from packagelist import RpmItem, DrpmItem
from ConfigParser import SafeConfigParser
from SimpleXMLRPCServer import SimpleXMLRPCServer
import xmlrpclib
import threading, thread
import time
import deltasig

import gzip
from zlib import error as zlibError
from gzip import write32u, FNAME

DRPMCONF = "/etc/prestod.conf"
DEBUG = 0
__version__ = "0.1.0"
PORT = 8094

def log(cmds, stuff, mtype="ERROR"):
    """Simple logging function"""
    if mtype=="ERROR" or (mtype=="INFO" and cmds['verbose']) or (mtype=="DEBUG" and cmds['debug']):
        print "%s: %s" % (mtype, stuff)

def _(args):
    """Stub function for translation"""
    return args

def _getNevra(filename, cmds, ts):
    """Get nevra of an rpm or source rpm of a deltarpm"""
    try:
        hdr = rpmUtils.miscutils.hdrFromPackage(ts, filename)
    except:
        log(cmds, _("Unable to open %s") % filename, "ERROR")
        raise
    else:
        nm = hdr['name']
        e = hdr['epoch']
        if e is None:
            e = "0"

        nevra = (hdr['name'], e, hdr['version'], hdr['release'], hdr['arch'])
        return nevra

class BuildThread(threading.Thread):
    def __init__(self, parent):
        threading.Thread.__init__(self)
        self.parent = parent
    
    def run(self):
        """Thread that builds deltarpms"""
        while self.parent.building:
            self.parent.lock.acquire()
            
            # Exit thread if we don't have anything to build
            if len(self.parent.queue) == 0:
                self.parent.building = False
                self.parent.lock.release()
                break

            # Build top build
            build = self.parent.queue[0]
            self.parent.lock.release()
            if build[0] == 'build':
                self.parent.makeDeltas(build[2], build[3])
            elif build[0] == 'sign':
                self.parent.__attach_signature(build[3], build[2], build[4])
            self.parent.lock.acquire()
            del self.parent.queue[0]
            self.parent.lock.release()
        log(self.parent.cmds, _("No more deltarpms to build, exitting thread..."), "DEBUG")
            
class Builder():
    """Class that builds deltarpms"""

    def __init__(self, cmds, config, ts):
        self.lock = thread.allocate_lock()
        self.cmds = cmds
        self.config = config
        self.ts = ts
        self.queue = []
        self.building = False
        self.thread = None
        self.current_build = []
        
    def add_delta(self, source_rpm, target_rpm, tag=""):
        """Add deltarpm to list of deltarpms to build"""
        print threading.enumerate()
        self.lock.acquire()
        self.queue.append(('build', tag, source_rpm, target_rpm))
        if not self.building:
            if self.thread != None:
                self.thread.join()
                del self.thread
            self.thread = BuildThread(self)
            self.building = True
            self.thread.start()
            log(self.cmds, _("Starting new build thread"), "DEBUG")
        self.lock.release()
        
    def remove_tag(self, tag):
        """Remove all deltarpms with tag 'tag' from queue"""
        self.lock.acquire()
        x = 0
        while x < len(self.queue):
            if self.queue[x][1] == tag:
                del self.queue[x]
            else:
                x += 1
        self.lock.release()
        
    def attach_signature(self, signature, target_rpm, dest_dir, tag=""):
        """Attach signature to deltarpms that point to target_rpm"""
        
        # First make sure that we're not still waiting for deltarpms to be built
        self.lock.acquire()
        count = 0
        x = len(self.queue) - 1
        while x >= 0:
            if self.queue[x][0] == 'sign' and self.queue[x][2] == target_rpm:
                log(self.cmds, _("Deltarpms for %s already waiting to have signatures attached") % os.path.basename(target_rpm))
                return False
            if self.queue[x][0] == 'build' and self.queue[x][3] == target_rpm:
                val = self.queue.pop(x)
                self.queue.insert(0, val)
                count += 1
            x -= 1
            
        # Add 'attach signature to deltarpm' to queue after said deltarpms are built
        if count > 0:
            self.queue.insert(count, ('sign', tag, target_rpm, signature, dest_dir))
            self.lock.release()
        else:
            self.lock.release()
            self.__attach_signature(signature, target_rpm, dest_dir)
        return False
        
    def __attach_signature(self, signature, target_rpm_filename, dest_dir):
        # Attach signature to all deltarpms for target_rpm_filename, saving signed deltarpms in dest_dir
        try:
            target_nevra = _getNevra(target_rpm_filename, self.cmds, self.ts)
        except:
            log(self.cmds, _("Unable to read %s. Is it not an rpm?") % target_rpm_filename)
            return False
        
        target_rpm_dir = os.path.dirname(target_rpm_filename)
        for filename in os.listdir("%s/%s.drpm" % (target_rpm_dir, target_nevra[0])):
            if filename[-5:] == ".drpm":
                try:
                    deltasig.make_delta_from_sig(os.path.join(target_rpm_dir, filename), signature, os.path.join(dest_dir, filename))
                except:
                    log(self.cmds, _("Error attaching signature to %s") % filename)
        return True
        
    def list_current_build(self):
        return self.queue[0]
        
    def list_builds(self, tag=False):
        return self.queue
        
#    def makeDeltas(self, source_rpm_filename, target_rpm_filename):
#        # Function to test threading without actually building deltarpms
#        log(self.cmds, _("Simulating building %s -> %s") % (source_rpm_filename, target_rpm_filename), "DEBUG")
#        time.sleep(60)
#        return True
        
    def makeDeltas(self, source_rpm_filename, target_rpm_filename):
        """Generate deltarpms from <source_rpm_filename> to <target_rpm_filename>"""
        
        target_dir = os.path.dirname(target_rpm_filename)
        source_dir = os.path.dirname(source_rpm_filename)

        build_count = self.config.getint(self.config.get('DEFAULT', 'default-section'), 'build-count')
        do_first = self.config.getboolean(self.config.get('DEFAULT', 'default-section'), 'do-first')
        
        if build_count == 0: # If we're not supposed to build deltarpms, just return True
            return True
            
        try:
            source_nevra = _getNevra(source_rpm_filename, self.cmds, self.ts)
        except:
            log(self.cmds, _("Unable to read %s. Is it not an rpm?") % source_rpm_filename)
            return False
        try:
            target_nevra = _getNevra(target_rpm_filename, self.cmds, self.ts)
        except:
            log(self.cmds, _("Unable to read %s. Is it not an rpm?") % target_rpm_filename)
            return False

        source_rpm = RpmItem(source_nevra, source_rpm_filename)
        target_rpm = RpmItem(target_nevra, target_rpm_filename)
        target_drpm = self.__buildDeltaRpm(source_rpm, target_rpm, target_dir)
        if not target_drpm:
            return False
            
        # Check for deltarpms in source_rpm's path.  If there are some,
        # we can combinedeltarpm them with the deltarpm we've just created
        # to make them rebuild to the target rpm rather than the source rpm.
        # This is how we create more than one deltarpm.
        source_delta_dir = "%s/%s.drpm" % (source_dir, source_nevra[0])
        if os.path.isdir(source_delta_dir) and build_count > 1:
            log(self.cmds, _("Deltarpms exist for source rpm %s") % os.path.basename(source_rpm_filename), "INFO")
            source_drpm_list = []
            for filename in os.listdir(source_delta_dir):
                if filename[-5:] == ".drpm":
                    filename = os.path.join(source_delta_dir, filename)
                    try:
                        (n1,e1,v1,r1) = self.__getDeltaNevr(os.path.join(source_delta_dir, filename))
                    except:
                        continue
                    
                    (n2,e2,v2,r2,a2) = _getNevra(filename, self.cmds, self.ts)
                    source_drpm = DrpmItem((n1,e1,v1,r1,a2),(n2,e2,v2,r2,a2), filename)
                    source_drpm_list.append(source_drpm)

            def sortByEVR(drpm1, drpm2, src=False):
                """Sort function for deltarpms.  Used to sort in reverse order (newest deltarpm first)"""
                if src:
                    (n1,e1,v1,r1,a1) = drpm1.getSrcNevra()
                    (n2,e2,v2,r2,a2) = drpm2.getSrcNevra()
                else:                        
                    (n1,e1,v1,r1,a1) = drpm1.getDstNevra()
                    (n2,e2,v2,r2,a2) = drpm2.getDstNevra()
                rc = rpmUtils.miscutils.compareEVR((e1,v1,r1),(e2,v2,r2))
                if rc == 0:
                    if src:
                        return 0
                    else:
                        return sortByEVR(drpm1, drpm2, True)
                if rc > 0:
                    return -1
                if rc < 0:
                    return 1   
            source_drpm_list.sort(sortByEVR)  # highest first in list

            # We've already created one deltarpm
            done = 1
            
            # If we're going to create a deltarpm against the oldest rpm, let's
            # just count it now
            if do_first:
                done += 1
                
            # Cycle through list of deltarpms making them rebuild to the target
            # rpm rather than the source rpm.  Once we reach build_count, break
            # out of the loop.
            for source_drpm in source_drpm_list:
                if source_drpm.getDstNevra != target_drpm.getDstNevra:
                    self.__combineDeltaRpm(source_drpm, target_drpm, target_rpm, target_dir)
                    done += 1
                    if done == build_count:
                        break
                        
            # Build deltarpm against oldest rpm, if that's what we're supposed to do
            if do_first:
                if source_drpm_list[-1].getDstNevra != target_drpm.getDstNevra:
                    self.__combineDeltaRpm(source_drpm_list[-1], target_drpm, target_rpm, target_dir)
                
        return True
                            
    def __getDeltaNevr(self, filename):
        # Get nevr for source rpm of a deltarpm.  (The deltarpm format doesn't store source arch.)
        # Return a tuple of (n, e, v, r)
        def _getLength(in_data):
            length = 0
            for val in in_data:
                length = length * 256
                length += ord(val)
            return length
            
        def _stringToVersion(strng):
            """Convert string to (e, v, r)"""
            i = strng.find(':')
            if i != -1:
                epoch = strng[:i]
            else:
                epoch = '0'
            j = strng.find('-')
            if j != -1:
                if strng[i + 1:j] == '':
                    version = None
                else:
                    version = strng[i + 1:j]
                release = strng[j + 1:]
            else:
                if strng[i + 1:] == '':
                    version = None
                else:
                    version = strng[i + 1:]
                release = None
            return (epoch, version, release)
            
        def _stringToNEVR(string):
            """Convert string to (n, e, v, r)"""
            i = string.rfind("-", 0, string.rfind("-")-1)
            name = string[:i]
            (epoch, ver, rel) = _stringToVersion(string[i+1:])
            return (name, epoch, ver, rel)
            
        (range_start, range_end) = byteranges(filename)
        fd = open(filename, "r")
        fd.seek(range_end)
        try:
            compobj = gzip.GzipFile("", "rb", 9, fd)
        except:
            raise zlibError("Data not stored in gzip format")
            
        if compobj.read(4)[:3] != "DLT":
            raise Exception("Not a deltarpm")
        
        nevr_length = _getLength(compobj.read(4))
        nevr = compobj.read(nevr_length).strip("\x00")
        oldnevr = _stringToNEVR(nevr)
        compobj.close()
        return oldnevr
    
    def __doWork(self, target_nevra, source_nevra, target_filename, dest_dir, deltaCommand):
        # Do the actual work of building a deltarpm.
        (n1,e1,v1,r1,a1) = target_nevra
        (n2,e2,v2,r2,a2) = source_nevra
        
        # Create directory that will store deltarpm
        drpm_drop_dir = os.path.join(dest_dir, '%s.drpm' % n1)
        try:
            os.mkdir(drpm_drop_dir)
        except OSError:
            pass
            
        deltaRPMName  = os.path.join(dest_dir, '%s.drpm/%s-%s-%s_%s-%s.%s.drpm' % (n1, n1, v2, r2, v1, r1, a1))
        deltaCommand += deltaRPMName
        log(self.cmds, deltaCommand, "DEBUG")
                    
        (code, out) = commands.getstatusoutput(deltaCommand)
        if code:
            log(self.cmds, _("Error generating deltarpm for %s: exitcode was %s - Reported Error: %s") % (n1, code, out))
            return False
        else:
            # Check whether or not we should keep the drpm
            drpmsize = os.path.getsize(deltaRPMName)
            if not self.__drpmIsWorthKeeping(drpmsize, target_filename):
                log(self.cmds, 'Deleting %s - not enough savings' % os.path.basename(deltaRPMName), "INFO")
                try:
                    os.unlink(deltaRPMName)
                except Exception, e:
                    log(self.cmds, _("Error deleting deltarpm %s: %s") % (os.path.basename(deltaRPMName), str(e)))
                try:
                    os.rmdir(drpm_drop_dir)
                except Exception, e:
                    log(self.cmds, _("Error deleting directory %s: %s") % (drpm_drop_dir, str(e)))
                return False
            else:
                if e1 == e2:
                    log(self.cmds, 'Generated delta rpm for %s.%s - %s-%s => %s-%s' % (n1, a1, v2, r2, v1, r1), "INFO")
                else:
                    log(self.cmds, 'Generated delta rpm for %s.%s - %s:%s-%s => %s:%s-%s' % (n1, a1, e2, v2, r2, e1, v1, r1), "INFO")
                drpm = DrpmItem(source_nevra, target_nevra, deltaRPMName)
                return drpm
    
    def __combineDeltaRpm(self, drpm1, drpm2, target_rpm, dest_dir):
        # Combine two deltarpms
        
        deltaCommand = 'combinedeltarpm %s %s ' % (drpm1.filename, drpm2.filename)
        return self.__doWork(drpm2.getDstNevra(), drpm1.getSrcNevra(), target_rpm.filename, dest_dir, deltaCommand)


    def __buildDeltaRpm(self, source_rpm, target_rpm, dest_dir):
        # Build a deltarpm
        
        deltaCommand = 'makedeltarpm %s %s ' % (source_rpm.filename, target_rpm.filename)
        return self.__doWork(target_rpm.getNevra(), source_rpm.getNevra(), target_rpm.filename, dest_dir, deltaCommand)

    def __drpmIsWorthKeeping(self, deltarpm_size, target_rpm):
        # Check to see whether the deltarpm is work keeping
        
        target_size = os.path.getsize(target_rpm)
        log(self.cmds, _("RPM Size: %i, DRPM Size: %i") % (target_size, deltarpm_size), "DEBUG")
        if deltarpm_size > config.getfloat(config.get('DEFAULT', 'default-section'), 'threshold') * target_size:
            return False
        return True
        
class XMLRPCFunctions:
    """Class that contains all XML-RPC functions"""
    def __init__(self, cmds, config):
        self.cmds = cmds
        self.config = config
        self.ts = rpmUtils.transaction.initReadOnlyTransaction()
        self.builder = Builder(cmds, config, self.ts)
        
    def _listMethods(self):
        """List available XML-RPC functions"""
        return ['makeDeltas', 'listWaitingBuilds', 'listBuilding']
        
    def _methodHelp(self, method):
        """Get help for XML-RPC functions"""
        if method == 'makeDeltas':
            return "makeDeltas(source_rpm, target_rpm, tag)\n" \
                   "\n" \
                   "Generate deltarpms from <source_rpm> (plus deltarpms in \n" \
                   "    <source_rpm's directory) to <target_rpm>\n" \
                   "tag is optional (Default: '')\n" \
                   "Return value: True if build was successfully queued, False otherwise\n"
        if method == 'listWaitingBuilds':
            return "listWaitingBuilds(tag)\n" \
                   "\n" \
                   "List builds that are waiting to be built\n" \
                   "If tag is not false, it will return *only* builds that match the tag\n" \
                   "Return value: List of (tag, source_rpm, target_rpm) tuples\n"
        else:
            return ""

    def makeDeltas(self, source_rpm_filename, target_rpm_filename):
        self.builder.add_delta(source_rpm_filename, target_rpm_filename)
        return True
    
    def _makeDeltas(self, source_rpm_filename, target_rpm_filename):
        """Generate deltarpms from <source_rpm_filename> to <target_rpm_filename>"""

        build_count = self.config.getint(self.config.get('DEFAULT', 'default-section'), 'build-count')
        
        if build_count == 0: # If we're not supposed to build deltarpms, just return True
            return True
            
        try:
            source_nevra = _getNevra(source_rpm_filename, self.cmds, self.ts)
        except:
            log(self.cmds, _("Unable to read %s. Is it not an rpm?") % source_rpm_filename)
            return False
        try:
            target_nevra = _getNevra(target_rpm_filename, self.cmds, self.ts)
        except:
            log(self.cmds, _("Unable to read %s. Is it not an rpm?") % target_rpm_filename)
            return False
            
        self.builder.add_delta(source_rpm_filename, target_rpm_filename)
        return True
        
    def listBuilding(self):
        return self.builder.list_current_build()
        
    def listWaitingBuilds(self, tag=False):
        return self.builder.list_builds(tag)
        
def usage(retval=1):
    print _("""
    prestod [options]

    Options:
     -c, --conf <file> = location of configuration file
     -n, --no-daemon = don't run as daemon
     -q, --quiet = run quietly
     -v, --verbose = run verbosely
     -d, --debug = turn debugging messages on
     -h, --help = show this help
     -V, --version = output version
    """)

    sys.exit(retval)
                
def parseArgs(args):
    """Parse the command line args return a commands dict and directory.
       Sanity check all the things being passed in."""
    cmds = {}
    cmds['quiet'] = False
    cmds['verbose'] = False
    cmds['conf'] = False
    cmds['no-daemon'] = False
    cmds['debug'] = False
    
    config = SafeConfigParser({'build-count': '5', 'do-first': 'True', 'default-section': 'DEFAULT', 'threshold': '0.5'})

    try:
        gopts, argsleft = getopt.getopt(args, 'hqvVc:nd', ['help', 'quiet', 'verbose', 'version', 'conf=', 'no-daemon', 'debug'])
    except getopt.error, e:
        log(cmds, _('Options Error: %s.') % e)
        usage()

    try:
        for arg,a in gopts:
            if arg in ['-h','--help']:
                usage(retval=0)
            elif arg in ['-V', '--version']:
                print '%s' % __version__
                sys.exit(0)
    except ValueError, e:
        log(cmds, _('Options Error: %s') % e)
        usage()

    if len(argsleft) != 0:
        usage()

    try:
        for arg,a in gopts:
            if arg in ['-v', '--verbose']:
                cmds['verbose'] = True
            elif arg in ['-q', '--quiet']:
                cmds['quiet'] = True
            elif arg in ['-n', '--no-daemon']:
                cmds['no-daemon'] = True
            elif arg in ['-d', '--debug']:
                cmds['debug'] = True
            elif arg in ['-c', '--conf']:
                cmds['conf'] = a
                
    except ValueError, e:
        log(cmds, _('Options Error: %s') % e)
        usage()
            
    if cmds['conf']:
        try:
            config.readfp(open(cmds['conf']))
        except:
            log(cmds, _("Configuration Error: Unable to read %s") % cmds['conf'])
            usage()
        else:
            log(cmds, _("Using %s as configuration file") % cmds['conf'], "INFO")
    else:
        try:
            config.readfp(open(DRPMCONF))
        except:
            log(cmds, _("Configuration Error: Unable to read %s") % DRPMCONF)
            log(cmds, _("Using defaults for configuration"), "INFO")
        else:
            log(cmds, _("Using %s as configuration file") % DRPMCONF, "INFO")
    
    # Set default-section for DEFAULT to 'default' if there is a default section in the .conf file
    # Basically, this is a hack so that we don't have to check whether to use 'DEFAULT' or 'default'
    # each time we want to access the default configuration.
    if config.has_section('default'):
        config.set("DEFAULT", "default-section", "default")
        config.remove_option("default", "default-section")
                
    log(cmds, _("Default section is \"%s\"") % config.get('DEFAULT', 'default-section'), "DEBUG")
    if cmds['debug']:
        for section in config.sections():
            log(cmds, "[%s]" % section, "DEBUG")
            for item in config.items(section):
                if item[0] != "default-section": # This item should never be used or set in the conf file
                    log(cmds, "%s: %s" % item, "DEBUG")
    return cmds, config

if __name__ == '__main__':
    cmds, config = parseArgs(sys.argv[1:])
        
    server = SimpleXMLRPCServer(('', PORT))
    server.register_introspection_functions()
    server.register_instance(XMLRPCFunctions(cmds, config))
    server.serve_forever()