rlarch

#!/usr/bin/env python

# going to explore installed packaged.
import os
import pkg_resources
import subprocess
import tempfile
import shutil
import sys
import fileinput
import struct

# should have a pacman.conf reader.
url_download = 'http://mir.archlinux.fr'

class EnvDir:
    """
    Just a container for directory variables.
    """
    def __init__(self, prefix):
        """
        Setup variables
        """
        self.prefix = prefix
        self.db_path   = os.path.join(self.prefix, 'var/lib/pacman')
        self.local_db  = os.path.join(self.db_path, 'local')
        self.sync_db   = os.path.join(self.db_path, 'sync')
        self.cache_dir = os.path.join(self.prefix, 'var/cache/pacman/pkg')

    def create(self):
        for p in (self.prefix, self.db_path, self.local_db, 
                  self.sync_db, self.cache_dir):
            if not os.path.exists(p): os.makedirs(p)


def _parse_package(fd):
    """
    Parse a package file descriptor. 
    @param fd:   the file descriptor from where we read data.
    @return map of category. Each category contains a vector.
    """
    result = {}
    current_cat = ''
    for line in fd.readlines():
        if '%' in line:
            current_cat = line[:-1].replace('%', '')
            continue
        if line == '\n': continue
        if current_cat != '':
                result.setdefault(current_cat, []).append(line[:-1])
    return result


def _read_database(db_name, incoming_db_filename):
    """
    Read a database.

    @return a map<package name, (version, depends, db_name)>.
    """
    import tarfile, zlib

    tar = tarfile.open(incoming_db_filename, 'r:gz')
    name = ''; version = ''
    result = {}
    for tarinfo in tar:
        # fetch version from directory.
        if tarinfo.isdir():
            name, version, revision = tarinfo.name.rsplit('-', 2)
            version = pkg_resources.parse_version(version + "-" + revision)
            
        # extract depends information.
        elif tarinfo.isfile():
            if "/depends" in tarinfo.name:
                depends = _parse_package(tar.extractfile(tarinfo))
                result[name] = (version, depends, db_name)
        
    return result

def _read_descs_from_db(packages, db_name, incoming_db_filename):
    """
    Read description for all packages from the database.
    We open the data using tar and we browse it.
    
    Once we found a requested package, we read the desc file.

    @param packages: a set of packages.
    @param dn_name: the name of the database.
    @param incoming_db_filename: the database filepath.

    @return a map<pkgname, (version, desc, db_name)>
    """
    import tarfile, zlib

    tar = tarfile.open(incoming_db_filename, 'r:gz')
    name = ''; version = ''
    result = {}
    for tarinfo in tar:
        # fetch version from directory.
        if tarinfo.isdir():
            name, version, revision = tarinfo.name.rsplit('-', 2)
            version = pkg_resources.parse_version(version + "-" + revision)
            continue
        
        # extract depends information.
        elif tarinfo.isfile():
            if name  not in packages: continue
            
            if "/desc" in tarinfo.name:
                desc = _parse_package(tar.extractfile(tarinfo))
                result[name] = (version, desc, db_name)
        
    return result


def _file_replace_inplace(filename, replaces):
    """
    Make a replace inplace in a file.
    @param filename: the filename that will be updated.
    @param search: the search pattern.
    @param replace: the replace Value
    @return the number of line changed.
    """
    count = 0
    for line in fileinput.input(filename, inplace=1):
        
        for r in replaces:
            if r[0] in line:
                line = line.replace(r[0], r[1])
                count += 1
        sys.stdout.write(line)
            
    return count


def _is_elf_file(filename):
    """
    Check if the file is an ELF file.
    Read only the first 4 bytes.
    
    @param filename: the filename to test.
    @return True on success.
    """
    if os.path.getsize(filename) < 4: return False
    f = open(filename, 'rb')
    header = struct.unpack('4c', f.read(4))
    return ('\x7f', 'E', 'L', 'F') == header


def _list_elf_files(dirname):
    """
    List all elf file present in the hierarchy.
    @param dirname: the directory name to browse.
    @return a list of absolute elf file path.
    """
    def _over(arg, dirname, fnames):
        for f in fnames:
            p = os.path.join(dirname, f)
            if not os.path.isfile(p): continue
            if _is_elf_file(p):
                arg.append(p)
    tmp = []
    os.path.walk(dirname, _over, tmp)
    return tmp


def _patch_elf_file(filename, new_root_path, replaces):
    """
    Call patchelft program to path the el file.
    We change the run path and the set interpreter.

    @param filename: the absolute path of the file to patch.
    @param new_root_path: the new root path used to set the run path.
    @param replaces: a vector containing what we have to replace in the rpath.
    """
    # TODO: we could put this stuff outside.
    mandatory_rpath = [os.path.join(new_root_path, p) for p in ['usr/lib', 'usr/local/lib', 'lib']]
    mandatory_rpath = ':'.join(mandatory_rpath)

    # get current rpath
    cmds = ["patchelf", "--print-rpath", filename]
    try:
        # get rpath and change it.
        rpath = subprocess.check_output(cmds)

        # we should add other rpath: /usr/lib:/usr/local/lib:/lib
        for r in replaces:
            rpath = rpath.replace(r[0], r[1])
        rpath = '%s:%s' % (rpath.replace('\n', ''), mandatory_rpath)

        # set rpath
        cmds = ["patchelf", "--set-rpath", rpath, filename]
        retcode = subprocess.call(cmds)

        # set interpreter
        cmds = ["patchelf", "--set-interpreter",
                os.path.join(new_root_path, "usr/lib/ld-linux-x86-64.so.2"), filename]
        retcode = subprocess.call(cmds)
                    
    except Exception, e:
        return False
    
    return True



def read_installed_packages(env_dir):
    """
    @param local_db_path: the local db path.
    @return a map indexed by package name. Value
    are a tuple (version, revision, description_map)
    """
    # detect installed package.
    installed_db = {} # indexed by program name.
    for soft in os.listdir(env_dir.local_db):
        name, version, revision = soft.rsplit('-', 2)
        version = pkg_resources.parse_version(version + "-" + revision)
        # read depends.
        descr = os.path.join(env_dir.local_db,
                             soft, 'desc')
        descr = _parse_package(open(descr))
        installed_db[name] = (version, descr)
    
    return installed_db



def read_databases(databases, env_dir):
    """
    Read databases passed by parameter.
    
    @param databases: array of databases
    @param env_dir: variable with all available path.
    
    @return a map<pkg name, (version, depends, db_name)>
    """
    # read all databases defined in databases variable.
    db_packages = {}
    for db_name in databases:
        db_file = os.path.join(env_dir.sync_db, db_name + '.db')
        pkgs = _read_database(db_name, db_file)
        db_packages.update(pkgs)
        
    return db_packages



def read_descriptions(packages, db_packages, env_dir):
    """
    Read description from databases for a set of packages.

    @param packages: a set of packages for which we want descriptions.
    @param databases: an array of database name.
    @param db_packages:

    @return a map<pkgname, (version, desc, db_name)>
    """
    # group packages per database.
    updates_per_db = {}
    for name in packages:
        pkg = db_packages[name]
        updates_per_db.setdefault(pkg[2], set()).add(name)

    # load descriptions.
    descriptions = {}
    for db_name, pkgs in updates_per_db.iteritems():

        db_file = os.path.join(env_dir.sync_db, db_name + '.db')
        pkgs = _read_descs_from_db(pkgs, db_name, db_file)
        descriptions.update(pkgs)

    return descriptions



def _compute_packages_order(update_packages,
                            installed_packages,
                            db_packages,
                            already_scheduled = {}
                            ):
    """
    Compute packages order for the install process. It
    will add required dependency to the list.
    """
    tmp = []
    for name in update_packages:

        # extract name.
        # ... TODO: do it better.
        if ">=" in name: name = name.split('>=')[0]
        if "==" in name: name = name.split('==')[0]
        
        # check if already scheduled
        if name in already_scheduled: continue
        
        # check if the package is available in the
        # database.
        db_pkg = db_packages.get(name)
        if db_pkg is None: continue

        # get dependencies of the package. If they are
        # not available only append the package to the
        # tmp list.
        deps = db_pkg[1].get('DEPENDS')
        if deps is None:
            # only append if a new version is available.
            if name not in installed_packages or db_packages[name][0] > installed_packages[name][0]:
                tmp.append(name)
                already_scheduled[name] = 1
            continue
        
        # process dependencies.
        tmp += _compute_packages_order(deps,
                                       installed_packages,
                                       db_packages,
                                       already_scheduled)

        # the current name
        if name in already_scheduled:
            print("strange: '%s' should not be already in the queue" % (name))
            continue
        
        # compare the version between installed 
        if name not in installed_packages or \
           db_packages[name][0] > installed_packages[name][0]:
            # append the current package.
            tmp.append(name)
            already_scheduled[name] = 1

    return tmp


def determine_package_to_update(installed_packages,
                                db_packages):
    """
    Determine which package to update. We take the list
    of installed packages, and check if a new version is available.

    @param installed_packages: a map of installed packages.
    @param db_packages: the list of database packages.
    
    @return a tuple of (update or install pkgs, no present in db)
    """
    # first step is to check which packages we have to update. 
    packages_to_update = []
    packages_missing = []
    for name, oldpkg in installed_packages.iteritems():
        # check if the package continue to exists.
        npkg = db_packages.get(name)
        if npkg is None:
            packages_missing.append(name)

        # if a new version is available, append it to the
        # package to install.
        elif npkg[0] > oldpkg[0]:
            packages_to_update.append(name)

    # order packages.
    packages_to_update = _compute_packages_order(packages_to_update,
                                                 installed_packages,
                                                 db_packages)
    return packages_to_update, packages_missing


def determine_package_to_install(packages, installed_packages,
                                 db_packages):
    """
    Determine which package to install. We dont care about the
    presence of packages, we install it if the caller want it.
    """
    packages_to_install = []
    packages_missing    = []

    for name in packages:
        if name not in db_packages: packages_missing.append(name)
        else: packages_to_install.append(name)

    tmp = _compute_packages_order(packages_to_install,
                                  installed_packages,
                                  db_packages)
    # force the install of requested packages even if the version
    # has not changed.
    packages_to_install = set(tmp + packages_to_install)
    return packages_to_install, packages_missing
    

def download_package(source, target, md5sum=None):
    """
    Download a package
    If a checksum is passed in parameter, we compute
    a md5sum on the target file (if it exists). If the md5sum is not valid,
    we erase the target and we start the download.

    @param source: the uri of the source file.
    @param target: the target file.
    @param md5sum: optional checksum used to compare with the target md5sum.

    @return True on success, False otherwise.
    """
    #xfer_command = "/usr/bin/curl -C - -f %u > %o"
    # /usr/bin/wget --passive-ftp -c -O %o %u
    # download using curl.
    if os.path.exists(target) and md5sum:
        # on file exists, we check 
        output = subprocess.check_output(["md5sum", target])
        compute_md5sum = output.split(' ')[0]
    
        if md5sum == compute_md5sum:
            return True
        print("wrong checksum, deleting file: ", target)
        os.unlink(target)
    
    print("downloading: ", source, target)
    retcode = subprocess.call(["/usr/bin/wget", "--progress=bar", "-c", "-O",
                                target, source ])
    return retcode == 0


def download_databases(from_url, databases, env_dir):
    """
    Download a vector of databases

    @param from_url: the url used to download databases.
    @param databases: a set of databases.
    @param env_dir: contains directory where to put the file.

    @return True if all databases are downloaded correctly.
    """
    # url format: http://.../database/os/arch/database.db
    ret = True
    for db in databases:
        #TODO: python 3: os.uname().machine
        db_url = '%s/%s/os/%s/%s.db' % (from_url, db, os.uname()[4], db)
        db_target = '%s/%s.db' % (env_dir.sync_db, db)
        
        ret = True and download_package(db_url, db_target)

    return ret


def uncompress_archive(filename):
    """
    uncompress a xz archive into a directory.
    @param filename: the archive filename
    @output_dir: the output directory.
    @return the directory where the archive has been extracted.
    The caller has to remove the directory.
    """
    # create the temporary output directory.
    tmp_dir = tempfile.mkdtemp()

    shutil.copy(filename, tmp_dir)
    filename = os.path.join(tmp_dir, os.path.basename(filename))
    retcode  = subprocess.call(["xz", "-d", filename])
    filename = filename[:-3]
    retcode  = subprocess.call(["tar", "xf", filename, "-C", tmp_dir])
    os.unlink(filename)

    return tmp_dir




def rewrite_path(package_dir, new_root_path, text_only = False):
    """
    @param package_dir: the package directory.
    @param new_root_path: the root path used to rewrite
    """
    # list of file that maybe rewritten.
    cmds = ["find",  package_dir, 
            "!", "-name", "*.h", 
            "!", "-name", ".html", 
            "!", "-name", "*.hpp",
            "-type", "f", 
            "-exec", "grep", "-Il", ".", "{}", ";"]
    outputs = subprocess.check_output(cmds)
    
    to_replace = [
        ['/usr', os.path.join(new_root_path, 'usr')],
        ['/etc', os.path.join(new_root_path, 'etc')],
        ['/var', os.path.join(new_root_path, 'var')],   
    ]

    count_text_modified = 0
    for f in outputs[:-1].split('\n'):
        count = _file_replace_inplace(f, to_replace)
        if count: count_text_modified += 1
    #print "number of rewritten text file:", count_text_modified
    
    if text_only: return
    
    # now patch binary file, find all elf file inside
    # a directory. once found, execute the patchelf command
    elf_files = _list_elf_files(package_dir)
    for e in elf_files:
        _patch_elf_file(e, new_root_path, to_replace)

    return
    

def _merge_directory(source_dir, target_dir):
    """
    used to move the contents of source directory
    to the target directoy.
    """

    skip_files = [os.path.join(source_dir, x) for x in [
        'var/lock'
    ]]

    installed_files = []
    len_target_dir = len(target_dir) + 1
    for src_dir, dirs, files in os.walk(source_dir):
        dst_dir = src_dir.replace(source_dir, target_dir)
        
        if not os.path.exists(dst_dir):
            os.mkdir(dst_dir)
        for file_ in files:
        
            
            src_file = os.path.join(src_dir, file_)
            dst_file = os.path.join(dst_dir, file_)
            
            installed_files.append(dst_file[len_target_dir:])

            if os.path.exists(dst_file):
                os.remove(dst_file)
            
            if src_file in skip_files: continue
            try:
                shutil.copy2(src_file, dst_dir)
            except Exception, e:
                print("+++ failed to copy file: ", src_file)
    
    return installed_files


def insert_to_local_db(package, version, descr,
                       depends, files, env_dir):
    # remove the directory if it does not exists.
    pkgver = "%s-%s" % (package, version)

    target_dir = os.path.join(env_dir.local_db, pkgver)
    if os.path.exists(target_dir):
        shutil.rmtree(target_dir)

    # create directory.
    os.makedirs(target_dir)

    descr.update(depends)

    # generate desc file.
    f = open(os.path.join(target_dir, 'desc'), 'w')
    for key, lines in descr.iteritems():
        f.write('%%%s%%\n' % (key))
        for line in lines:
            f.write(line + '\n')
    f.close()

    # write files
    f = open(os.path.join(target_dir, 'files'), 'w')
    for file in files:
        f.write(file + '\n')
    f.close()


def update_packages(packages, descriptions, 
                    installed_packages, db_packages, 
                    url, env_dir):


    for i, pkg in enumerate(to_update):
        desc     = descriptions[pkg][1]
        filename = desc['FILENAME'][0]
        db_version = desc['VERSION'][0]
        db_name  = db_packages[pkg][2]
        db_depends = db_packages[pkg][1]
        md5sum_expected = desc['MD5SUM'][0]

        # build source & target.
        source_url =  '%s/%s/os/x86_64/%s' % (url, db_name, filename)
        target_url = '%s/%s' % (env_dir.cache_dir, filename)

        # download
        print "downloading %s (%d/%d)" % (pkg, i + 1, len(to_update))
        download_package(source_url, target_url, md5sum = md5sum_expected)

        # note we should check if md5sum are valid.

    for pkg in to_update:

        desc     = descriptions[pkg][1]
        filename = desc['FILENAME'][0]
        db_version = desc['VERSION'][0]
        db_name  = db_packages[pkg][2]
        db_depends = db_packages[pkg][1]
        md5sum_expected = desc['MD5SUM'][0]

        # build source & target.
        source_url =  '%s/%s/os/x86_64/%s' % (url, db_name, filename)
        target_url = '%s/%s' % (env_dir.cache_dir, filename)
        
    
        # uncompress xz into temporary directory.
        tmp_dir = uncompress_archive(target_url)
        
        # rewrite rpath and other files.
        rewrite_path(tmp_dir, env_dir.prefix, pkg in ('glib',))

        # uninstall previous package.
        if pkg in installed_packages:
            print "+++ we have to uninstall", pkg

        # copy file.
        installed_files = _merge_directory(tmp_dir, env_dir.prefix)

        # should destroy the temporary directory.

        # generate local save.
        insert_to_local_db(pkg, db_version, desc, db_depends,
                           installed_files, env_dir)


def params():
    import optparse

    parser = optparse.OptionParser()

    # --> Verbosity
    parser.add_option('-v', '--verbose',
                    action = 'store_true',
                    dest = 'verbose',
                    default = False,
                    help = 'enable verbose messages [default: %default]')

    parser.add_option('-E', '--env',
                       action= 'store',
                       dest='env',
                       default=os.environ.get('RLARCH_PREFIX'))

    parser.add_option('--create-env',
                       action='store_true',
                       dest='create_env',
                       default=False,
                       help= 'use this option to bootstrap an environnement.')

    parser.add_option('-S',
                       action='store_true',
                       dest='synchronize',
                       default = False,
                       help = 'Synchronize package.')

    parser.add_option('-y',
                       action='store_true',
                       dest='fetch_db',
                       default = False,
                       help = 'Update the database from the Official ArchLinux repository.')

    parser.add_option('-u',
                       action='store_true',
                       dest='upgrade',
                       default = False,
                       help = 'Upgrade packages of the system')

    parser.add_option('--from-url',
                       action = 'store',
                       dest='from_url',
                       default = 'http://mir.archlinux.fr',
                       help= 'The url used to download database and packages.')

    parser.add_option('--databases',
                      action='store',
                      dest='databases',
                      default='core,extra,community',
                      help='List of comma separated databases.')

    config, args = parser.parse_args()
    if config.databases:
        config.databases = config.databases.split(',')

    return config, args


def display_update(packages, installed_packages, descriptions):
    print "-" * 62
    print "%-40s %10s %10s" % ("Package", "Cur Vers", "New Vers")
    for pkg in to_update:
        inst_pkg = installed_packages.get(pkg)
        db_pkg   = descriptions.get(pkg)
        inst_version = None
        if inst_pkg: inst_version = inst_pkg[1].get('VERSION')[0]
        db_version = db_pkg[1].get('VERSION')[0]
        print "%-40s %10s %10s" % (pkg, inst_version, db_version)
    print "-" * 62
    

def getch():
  import termios
  import fcntl
  import sys
  fd = sys.stdin.fileno()

  oldterm = termios.tcgetattr(fd)
  newattr = termios.tcgetattr(fd)
  newattr[3] = newattr[3] & ~termios.ICANON & ~termios.ECHO
  termios.tcsetattr(fd, termios.TCSANOW, newattr)

  oldflags = fcntl.fcntl(fd, fcntl.F_GETFL)
  fcntl.fcntl(fd, fcntl.F_SETFL, oldflags | os.O_NONBLOCK)

  try:        
    while 1:            
      try:
        c = sys.stdin.read(1)
        break
      except IOError: pass
  finally:
    termios.tcsetattr(fd, termios.TCSAFLUSH, oldterm)
    fcntl.fcntl(fd, fcntl.F_SETFL, oldflags)
  return c



if __name__ == '__main__':

    # create parser
    config, packages = params()

    # make sure that environnement is defined.
    if config.env is None:
        print("you have to define the prefix to use through the RLARCH_PREFIX variable or --env parameter.")
        sys.exit(1)

    # make sure that the directory exists.
    # TODO: create a check env function that will ensure that basic requirements are present.
    if not os.path.exists(config.env):
        print("the environnement '%s' does not exist, you have to bootstrap it." % (config.env))
        sys.exit(1)

    env_dir = EnvDir(config.env)

    if config.create_env:
        env_dir.create()


    # synchronize stuff
    if config.synchronize:
        if config.fetch_db:
            # make sure that the target dictionary exists.
            print("going to synchronize database.")
            download_databases(config.from_url, config.databases, env_dir)

        if config.upgrade:
            # list installed packages
            installed_packages = read_installed_packages(env_dir)
            print("nb installed packages: ", len(installed_packages))

            # read databases
            db_packages = read_databases(config.databases, env_dir)

            # determine which package to update.
            to_update, _ = determine_package_to_update(installed_packages,
                                                       db_packages)

            print("reading description for: ", to_update)
            descriptions = read_descriptions(to_update,
                                             db_packages, env_dir)

            if len(to_update) == 0:
                print "nothing to do"
                sys.exit(0)
                
            display_update(to_update, installed_packages, descriptions)
            print "do you want to make the update ? (y/n): ",
            answer = getch()
            if answer != 'y':
                print "leaving ..."
                sys.exit(1)
            
            # make the update.
            update_packages(to_update, descriptions, 
                            installed_packages, db_packages,
                            config.from_url, env_dir)            
        
        if len(packages):   
            
            # read installed packages.
            installed_packages = read_installed_packages(env_dir)
            print("nb installed packages: ", len(installed_packages))

            # read databases
            db_packages = read_databases(config.databases, env_dir)

            # determine which package to update.
            to_update, missing = determine_package_to_install(packages,
                                                              installed_packages,
                                                              db_packages)
            if len(missing):
                print "missing packages are: ", missing
            
            print("reading description for: ", to_update)
            descriptions = read_descriptions(to_update,
                                             db_packages, env_dir)

            display_update(to_update, installed_packages, descriptions)
            print "do you want to continue ? (y/n): ",
            answer = getch()
            if answer != 'y':
                print "leaving ..."
                sys.exit(1)
                
            # update packages.
            update_packages(to_update, descriptions, 
                            installed_packages, db_packages,
                            config.from_url, env_dir)