#!/usr/bin/python
#
# Copyright 2005 Lars Wirzenius (liw@iki.fi)
# 
# 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


"""Debian package installation and uninstallation tester.

This program sets up a minimal Debian system in a chroot, and installs
and uninstalls packages and their dependencies therein, looking for
problems.

See the manual page (piuparts.1, generated from piuparts.docbook) for
more usage information.

Lars Wirzenius <liw@iki.fi>
"""


VERSION = "0.7"


import time
import getopt
import sys
import commands
import tempfile
import shutil
import os
import tarfile
import stat
import re
import pickle


class Settings:

    """Global settings for this program."""
    
    def __init__(self):
        self.max_command_output_size = 1024 * 1024
        self.args_are_package_files = True
        self.really_verbose = False
        self.verbose = False
        self.debian_mirrors = []
        self.debian_distros = []
        self.basetgz = None
        self.savetgz = None
        self.endmeta = None
        self.saveendmeta = None
        self.log_files = [sys.stdout]
        self.ignored_files = [
            "/etc/ld.so.cache",
            "/etc/ld.so.conf",
            "/usr/share/info/dir",
            "/usr/share/info/dir.old",
            "/var/backups/infodir.bak",
            "/var/cache/apt/archives/lock",
            "/var/cache/apt/pkgcache.bin", 
            "/var/cache/debconf",
            "/var/cache/debconf/config.dat",
            "/var/cache/debconf/config.dat-old",
            "/var/cache/debconf/passwords.dat",
            "/var/cache/debconf/templates.dat",
            "/var/cache/debconf/templates.dat-old",
            "/var/cache/man/index.db",
            "/var/lib/dpkg/available", 
            "/var/lib/dpkg/available-old", 
            "/var/lib/dpkg/lock", 
            "/var/lib/dpkg/status", 
            "/var/lib/dpkg/status-old", 
            "/var/log/dpkg.log",
            "/",
            ]
        self.ignored_patterns = [
            "/var/log/exim/.*",
            "/var/log/exim4/.*",
            "/var/mail/.*",
            "/var/spool/exim/.*",
            "/var/spool/exim4/.*",
            "/var/run/.*",
            ]


settings = Settings()
startup_time = time.time()


def log_msg(log_to_default, timestamp, prefix, msg, exit_code):
    """Write log message to stdout.
        'log_to_default' tells whether to log to stdout.
        'timestamp' is None or the the time (from time.time()).
       'prefix' is "DEBUG", or other prefix to add to message.
       'msg' is the actual message.
       If 'exit_code' is non-zero, exit process with that exit code."""
    if timestamp:
        t = timestamp - startup_time
        min = int(t) / 60
        sec = t % 60
    for f in settings.log_files:
        if f == sys.stdout and not log_to_default:
            continue
        if timestamp:
            f.write("%dmin%.1fs " % (min, sec))
        if prefix:
            f.write("%s: " % prefix)
        f.write(msg)
        f.write("\n")
        f.flush()
    if exit_code:
        sys.exit(exit_code)


def dump(msg):
    """Write really verbose debugging message."""
    log_msg(settings.really_verbose, None, "", msg, 0)


def debug(msg):
    """Write a debugging log message."""
    log_msg(settings.verbose, time.time(), "DEBUG", msg, 0)


def info(msg):
    """Write an informational log message."""
    log_msg(True, time.time(), "", msg, 0)


def error(msg, exit_code=1):
    """Write an error message and exit (unless exit_code is 0)."""
    log_msg(True, time.time(), "ERROR", msg, exit_code=exit_code)


def indent_string(str):
    """Indent all lines in a string with two spaces and return result."""
    return "\n".join(["  " + line for line in str.split("\n")])


def run(command, ignore_errors=False):
    """Run an external command and die with error message if it failes."""
    debug("Starting command: %s" % command)
    f = os.popen("{ export LC_ALL=C LANGUAGES=; %s ; } 2>&1" % 
                 command, "r")
    output = ""
    while 1:
        line = f.readline()
        if not line:
            break
        output += line
        output = output[-settings.max_command_output_size:]
        if line[-1:] == "\n":
            line = line[:-1]
        dump("  " + line)
    status = f.close()
    if status is None:
        status = 0
    failed = not os.WIFEXITED(status) or os.WEXITSTATUS(status) != 0
    if not failed:
        debug("Command ok: %s" % repr(command))
    elif ignore_errors:
        debug("Command failed (status=%d), but ignoring error: %s" % 
              (status, repr(command)))
    else:
        error("Command failed (status=%d): %s\n%s" % 
              (status, repr(command), indent_string(output)))
    return status, output


def run_in_chroot(root, command, ignore_errors=False):
    return run("chroot %s %s" % (root, command), ignore_errors=ignore_errors)


def create_temp_dir():
    """Create a temporary directory and return its full path."""
    path = tempfile.mkdtemp(dir=".")
    os.chmod(path, 0755)
    debug("Created temporary directory %s" % path)
    return path


def create_temp_file():
    """Create a temporary file and return its full path."""
    (fd, path) = tempfile.mkstemp(dir=".")
    debug("Created temporary file %s" % path)
    return (fd, path)


def remove_dir_tree(path):
    """Remove a directory tree, including all files and contents."""
    shutil.rmtree(path)
    debug("Removed directory tree at %s" % path)


def create_file(name, contents):
    """Create a new file with the desired name and contents."""
    try:
        f = file(name, "w")
        f.write(contents)
        f.close()
    except IOError, detail:
        error("Couldn't create file %s: %s" % (name, detail))


def copy_files(source_names, target_name):
    """Copy files in 'source_name' to file/dir 'target_name'."""
    debug("Copying %s to %s" % (", ".join(source_names), target_name))
    for source_name in source_names:
        try:
            shutil.copy(source_name, target_name)
        except IOError, detail:
            error("Error copying %s to %s: %s" % 
                  (source_name, target_name, detail))


def remove_files(filenames):
    """Remove some files."""
    for filename in filenames:
        debug("Removing %s" % filename)
        try:
            os.remove(filename)
        except OSError, detail:
            error("Couldn't remove %s: %s" % (filename, detail))


def create_apt_sources(root, distro):
    """Create an /etc/apt/sources.list in a chroot with a given distro."""
    lines = []
    for mirror, components in settings.debian_mirrors:
        lines.append("deb %s %s %s\n" % 
                     (mirror, distro, " ".join(components)))
    create_file(os.path.join(root, "etc/apt/sources.list"), "".join(lines))


def create_policy_rc_d(root):
    """Create a policy-rc.d that prevents daemons from running in chroot."""
    create_file(os.path.join(root, "usr/sbin/policy-rc.d"),
                "#!/bin/sh\nexit 101\n")


def setup_minimal_chroot(path):
    """Set up a minimal Debian system in a chroot at 'path'."""
    debug("Setting up minimal chroot for %s at %s." % 
          (settings.debian_distros[0], path))
    run("debootstrap %s '%s' '%s'" % 
        (settings.debian_distros[0], path, settings.debian_mirrors[0][0]))
    create_apt_sources(path, settings.debian_distros[0])
    create_policy_rc_d(path)


def create_chroot():
    """Create a chroot according to user's wishes. Return it's path."""
    root = create_temp_dir()
    if settings.basetgz:
        unpack_tarball_to_dir(settings.basetgz, root)
        run_in_chroot(root, "apt-get update")
        run_in_chroot(root, "apt-get clean")
    else:
        setup_minimal_chroot(root)
        run_in_chroot(root, "apt-get update")
        run_in_chroot(root, "apt-get clean")
        if settings.savetgz:
            save_dir_tree(root, settings.savetgz)
    return root
        

def save_dir_tree(path, result):
    """Tar and compress all files in the directory tree at 'path'."""
    debug("Saving %s to %s." % (path, result))
    try:
        tf = tarfile.open(result, "w:gz")
        tf.add(path, arcname=".")
        tf.close()
    except tarfile.TarError, detail:
        error("Couldn't create tar file %s: %s" % (result, detail))


def unpack_tarball_to_dir(tarball, path):
    """Unpack a tarball to a directory (cf. 'tar -C path -xzf tarball')."""
    debug("Unpacking %s into %s" % (tarball, path))
    run("tar -C '%s' -zxf '%s'" % (path, tarball))

    return
    try:
        tf = tarfile.open(tarball, "r|gz")
        for ti in tf:
            tf.extract(ti, path)
        tf.close()
    except tarfile.TarError, detail:
        error("Couldn't extract tar file %s into %s: %s" % 
              (tarball, path, detail))


def save_dir_tree_meta_data(root):
    """Return the meta data for all objects in a dir tree."""
    root = os.path.join(root, ".")
    dict = {}
    for dirpath, dirnames, filenames in os.walk(root):
        assert dirpath[:len(root)] == root
        for name in [dirpath] + [os.path.join(dirpath, f) for f in filenames]:
            st = os.lstat(name)
            if stat.S_ISLNK(st.st_mode):
                target = os.readlink(name)
            else:
                target = None
            dict[name[len(root):]] = (st, target)
    return dict


def objects_are_different(pair1, pair2):
    """Are filesystem objects different based on their meta data?"""
    (m1, target1) = pair1
    (m2, target2) = pair2
    if (m1.st_mode != m2.st_mode or 
        m1.st_uid != m2.st_uid or 
        m1.st_gid != m2.st_gid or
        target1 != target2):
        return True
    if stat.S_ISREG(m1.st_mode):
        return m1.st_size != m2.st_size # or m1.st_mtime != m2.st_mtime
    return False


def diff_meta_data(tree1, tree2):
    """Compare two dir trees and return list of new files (only in 'tree2'),
       removed files (only in 'tree1'), and modified files."""

    tree1 = tree1.copy()
    tree2 = tree2.copy()

    for name in settings.ignored_files:
        if name in tree1:
            del tree1[name]
        if name in tree2:
            del tree2[name]

    for pattern in settings.ignored_patterns:
        pat = re.compile(pattern)
        for name in tree1.keys():
            m = pat.match(name)
            if m:
                del tree1[name]
        for name in tree2.keys():
            m = pat.match(name)
            if m:
                del tree2[name]

    modified = []
    for name in tree1.keys()[:]:
        if name in tree2:
            if objects_are_different(tree1[name], tree2[name]):
                modified.append((name, tree1[name]))
            del tree1[name]
            del tree2[name]

    removed = [x for x in tree1.iteritems()]
    new = [x for x in tree2.iteritems()]

    return new, removed, modified


def file_list(meta_infos):
    """Return list of indented filenames."""
    meta_infos = meta_infos[:]
    meta_infos.sort()
    return "\n".join(["  " + name for (name, data) in meta_infos])


def get_chroot_selections(root):
    """Get current package selections in a chroot."""
    (status, output) = run_in_chroot(root, "dpkg --get-selections '*'")
    list = [line.split() for line in output.split("\n") if line.strip()]
    dict = {}
    for name, status in list:
        dict[name] = status
    return dict


def diff_selections(root, selections):
    """Compare original and current package selection.
       Return dict where dict[package_name] = original_status, that is,
       the value in the dict is the state that the package needs to be
       set to to restore original selections."""
    changes = {}
    current = get_chroot_selections(root)
    for name, value in current.iteritems():
        if name not in selections:
            changes[name] = "purge"
        elif selections[name] != current[name] and \
             selections[name] == "purge":
            changes[name] = selections[name]
    return changes


def remove_or_purge(root, operation, packages):
    """Remove or purge packages in a chroot."""
    for name in packages:
        run_in_chroot(root, "dpkg --%s %s" % (operation, name), 
                      ignore_errors=True)
    run_in_chroot(root, "dpkg --remove --pending", ignore_errors=True)


def restore_chroot_selections(root, changes, packages):
    """Restore package selections in a chroot by applying 'changes'.
       'changes' is a return value from diff_selections."""

    deps = {}
    nondeps = {}
    for name, state in changes.iteritems():
        if name in packages:
            nondeps[name] = state
        else:
            deps[name] = state

    deps_to_remove = [name for name, state in deps.iteritems()
                      if state == "remove"]
    deps_to_purge = [name for name, state in deps.iteritems()
                      if state == "purge"]
    nondeps_to_remove = [name for name, state in nondeps.iteritems()
                         if state == "remove"]
    nondeps_to_purge = [name for name, state in nondeps.iteritems()
                        if state == "purge"]

    # First remove all packages.
    remove_or_purge(root, "remove", deps_to_remove + deps_to_purge +
                                    nondeps_to_remove + nondeps_to_purge)

    # Then purge all packages being depended on.
    remove_or_purge(root, "purge", deps_to_purge)
    
    # Finally, purge actual packages.
    remove_or_purge(root, "purge", nondeps_to_purge)

    # Now do a final run to see that everything worked.
    run_in_chroot(root, "dpkg --purge --pending")
    run_in_chroot(root, "dpkg --remove --pending")


def get_package_names_from_package_files(filenames):
    """Return list of package names given list of package file names."""
    list = []
    for filename in filenames:
        (status, output) = run("dpkg --info " + filename)
        for line in [line.lstrip() for line in output.split("\n")]:
            if line[:len("Package:")] == "Package:":
                list.append(line.split(":", 1)[1].strip())
    return list


def apt_get_knows(root, package_names):
    """Does apt-get (or apt-cache) in chroot know about a set of packages?"""
    for name in package_names:
        (status, output) = run_in_chroot(root, "apt-cache show " + name,
                                             ignore_errors=True)
        if not os.WIFEXITED(status):
            error("Error occurred when running apt-cache in chroot:\n" +
                  output)
        if os.WEXITSTATUS(status) != 0 or not output.strip():
            return False
    return True
    

def install_package_files_into_chroot(root, filenames):
    """Install package files into ... er, read the function name, willya."""
    if filenames:
        root_tmp = os.path.join(root, "tmp")
        copy_files(filenames, root_tmp)
        tmp_files = [os.path.basename(a) for a in filenames]
        tmp_files = [os.path.join("tmp", name) for name in tmp_files]
        run_in_chroot(root, "dpkg -i " + 
                            " ".join(tmp_files), ignore_errors=True)
        run_in_chroot(root, "apt-get -yf install")
        run_in_chroot(root, "apt-get clean")
        remove_files([os.path.join(root, name) for name in tmp_files])


def check_results(root, root_info):
    """Check that current chroot state matches 'root_info'."""
    current_info = save_dir_tree_meta_data(root)
    (new, removed, modified) = diff_meta_data(root_info, current_info)
    ok = True
    if new:
        error("Package purging left files on system:\n" + file_list(new), 0)
        ok = False
    if removed:
        error("After purging files have disappeared:\n" +
              file_list(removed), 0)
        ok = False
    if modified:
        error("After purging files have been modified:\n" +
              file_list(modified), 0)
        ok = False

    return ok


def install_purge_test(root, root_info, selections, args, packages):
    """Do an install-purge test. Return True if successful, False if not.
       Assume 'root' is a directory already populated with a working
       chroot, with packages in states given by 'selections'."""

    # Install packages into the chroot.
    if args:
        install_package_files_into_chroot(root, args)
    else:
        if packages:
            run_in_chroot(root, "apt-get -y install " +
                                " ".join(["'%s'" % arg for arg in packages]))
        run_in_chroot(root, "apt-get clean")

    # Remove all packages from the chroot that weren't there initially.    
    changes = diff_selections(root, selections)
    restore_chroot_selections(root, changes, packages)
    
    return check_results(root, root_info)


def install_upgrade_test(root, root_info, selections, args, package_names):
    """Install package via apt-get, then upgrade from package files.
       Return True if successful, False if not."""

    # First install via apt-get.
    if package_names:
        run_in_chroot(root, "apt-get -yf install " + " ".join(package_names))

    # Then from the package files.
    install_package_files_into_chroot(root, args)

    # Remove all packages from the chroot that weren't there
    # initially.
    changes = diff_selections(root, selections)
    restore_chroot_selections(root, changes, package_names)
    
    return check_results(root, root_info)


def upgrade_to_distros(root, distros):
    """Upgrade a chroot installation to each successive distro."""
    for distro in distros:
        debug("Upgrading to %s" % distro)
        create_apt_sources(root, distro)
        run_in_chroot(root, "apt-get update")
        run_in_chroot(root, "apt-get -yf dist-upgrade")


def save_meta_data(filename, root_info, selections):
    """Save directory tree meta data into a file for fast access later."""
    debug("Saving chroot meta data to %s" % filename)
    f = file(filename, "w")
    pickle.dump((root_info, selections), f)
    f.close()


def load_meta_data(filename):
    """Load meta data saved by 'save_meta_data'."""
    debug("Loading chroot meta data from %s" % filename)
    f = file(filename, "r")
    (root_info, selections) = pickle.load(f)
    f.close()
    return root_info, selections


def install_and_upgrade_between_distros(filenames, packages):
    """Install package and upgrade it between distributions, then remove.
       Return True if successful, False if not."""

    root = create_chroot()

    if settings.basetgz:
        root_tgz = settings.basetgz
    else:
        (fd, root_tgz) = create_temp_file()
        save_dir_tree(root, root_tgz)
        
    if settings.endmeta:
        root_info, selections = load_meta_data(settings.endmeta)
    else:
        upgrade_to_distros(root, settings.debian_distros[1:])
        run_in_chroot(root, "apt-get clean")
    
        root_info = save_dir_tree_meta_data(root)
        selections = get_chroot_selections(root)
        
        if settings.saveendmeta:
            save_meta_data(settings.saveendmeta, root_info, selections)
    
        remove_dir_tree(root)
        root = create_temp_dir()
        unpack_tarball_to_dir(root_tgz, root)
    
    run_in_chroot(root, "apt-get update")
    if packages:
        run_in_chroot(root, "apt-get -yf install " + " ".join(packages))

    upgrade_to_distros(root, settings.debian_distros[1:])

    install_package_files_into_chroot(root, filenames)
    run_in_chroot(root, "apt-get clean")

    changes = diff_selections(root, selections)
    restore_chroot_selections(root, changes, packages)
    result = check_results(root, root_info)

    if root_tgz != settings.basetgz:
        remove_files([root_tgz])
    remove_dir_tree(root)

    return result


def parse_mirror_spec(str, defaultcomponents=[]):
    """Parse a mirror specification from the --mirror option argument.
       Return (mirror, componentslist)."""
    parts = str.split()
    return parts[0], parts[1:] or defaultcomponents[:]


def find_default_debian_mirrors():
    """Find the default Debian mirrors."""
    mirrors = []
    try:
        f = file("/etc/apt/sources.list", "r")
        for line in f:
            parts = line.split()
            if len(parts) > 2 and parts[0] == "deb":
                mirrors.append((parts[1], parts[3:]))
        f.close()
    except IOError:
        return None
    return mirrors


def parse_command_line():
    """Parse the command line, change global settings, return non-options."""
    (opts, args) = getopt.getopt(sys.argv[1:], "ab:B:d:hi:I:l:m:nps:S:vV",
                                 ["distribution=", "mirror=", "verbose",
                                  "log-file=", "ignore=", "no-ignores",
                                  "ignore-regex=", "help", "basetgz=",
                                  "save=", "save-end-meta=", "version", 
                                  "apt", "pbuilder", "end-meta"])
    exit = None
    for opt, optarg in opts:
        if opt in ["-a", "--apt"]:
            settings.args_are_package_files = False
        elif opt in ["-b", "--basetgz"]:
            settings.basetgz = optarg
        elif opt in ["-B", "--end-meta"]:
            settings.endmeta = optarg
        elif opt in ["-d", "--distribution"]:
            settings.debian_distros.append(optarg)
        elif opt in ["-h", "--help"]:
            sys.stdout.write("Usage: piuparts [options] package-file ...\n")
            sys.stdout.write("See manual page piuparts(1) for more help.\n")
            exit = 0
        elif opt in ["-i", "--ignore"]:
            settings.ignored_files.append(optarg)
        elif opt in ["-I", "--ignore-regex"]:
            settings.ignored_patterns.append(optarg)
        elif opt in ["-l", "--log-file"]:
            try:
                f = file(optarg, "a")
            except IOError, detail:
                error("Can't open log file %s: %s" % (optarg, detail))
                exit = 1
            settings.log_files.append(f)
        elif opt in ["-m", "--mirror"]:
            settings.debian_mirrors.append(parse_mirror_spec(optarg,
                                                             ["main",
                                                              "contrib",
                                                              "non-free"]))
        elif opt in ["-n", "--no-ignores"]:
            settings.ignored_files = []
            settings.ignored_patterns = []
        elif opt in ["-p", "--pbuilder"]:
            settings.basetgz = "/var/cache/pbuilder/base.tgz"
        elif opt in ["-s", "--save"]:
            settings.savetgz = optarg
        elif opt in ["-S", "--save-end-meta"]:
            settings.saveendmeta = optarg
        elif opt in ["-v", "--verbose"]:
            if not settings.verbose:
                settings.verbose = True
                settings.really_verbose = False
            else:
                settings.verbose = True
                settings.really_verbose = True
        elif opt in ["-V", "--version"]:
            sys.stdout.write("Version: %s\n" % VERSION)
            exit = 0

    if not settings.debian_distros:
        settings.debian_distros = ["sid"]

    if not settings.debian_mirrors:
        settings.debian_mirrors = find_default_debian_mirrors()
        if not settings.debian_mirrors:
            settings.debian_mirrors = [("http://ftp.debian.org/",
                                        ["main", "contrib", "non-free"])]

    if exit is not None:
        sys.exit(exit)

    return args
    

def main():
    """Main program. But you knew that."""

    args = parse_command_line()

    info("---------------------------------------------------------------")
    info("piuparts version %s starting up." % VERSION)
    info("Command line arguments: %s" % " ".join(sys.argv))

    # Make sure debconf does not ask questions and stop everything.
    # Packages that don't use debconf will lose.
    os.environ["DEBIAN_FRONTEND"] = "noninteractive"

    # Find the names of packages.
    if settings.args_are_package_files:
        packages = get_package_names_from_package_files(args)
    else:
        packages = args
        args = []

    if len(settings.debian_distros) == 1:
        root = create_chroot()

        root_info = save_dir_tree_meta_data(root)
        selections = get_chroot_selections(root)
    
        if not install_purge_test(root, root_info, selections,
				  args, packages):
            remove_dir_tree(root)
            error("FAIL: Installation and purging test.")
        info("PASS: Installation and purging test.")
    
        if not settings.args_are_package_files:
            info("Can't test upgrades: -a or --apt option used.")
        elif not apt_get_knows(root, packages):
            info("Can't test upgrade: packages not known by apt-get.")
        elif install_upgrade_test(root, root_info, selections, args, 
                                  packages):
            info("PASS: Installation, upgrade and purging tests.")
        else:
            remove_dir_tree(root)
            error("FAIL: Installation, upgrade and purging tests.")
    
        remove_dir_tree(root)
    else:
        if install_and_upgrade_between_distros(args, packages):
            info("PASS: Upgrading between Debian distributions.")
        else:
            error("FAIL: Upgrading between Debian distributions.")

    info("PASS: All tests.")
    info("piuparts run ends.")


if __name__ == "__main__":
    main()
