#!/usr/bin/python

# Execute operations on/in apport chroots.
#
# Copyright (c) 2007 Canonical Ltd.
# Author: Martin Pitt <martin.pitt@ubuntu.com>
#
# 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.  See http://www.gnu.org/copyleft/gpl.html for
# the full text of the license.

import optparse, os.path, sys, urllib, re, atexit, shutil, subprocess, tempfile
from glob import glob

import problem_report
from apport.chroot import Chroot, setup_fakeroot_env
from apport.crashdb import get_crashdb

#
# functions
#

def parse_options():
    '''Parse command line options and return (options, args) tuple.'''

    optparser = optparse.OptionParser('''%prog [options] create <release> <chroot path>
%prog [options] upgrade <chroot path>|<chroot release name>|all
%prog [options] installdeb <chroot path>|<chroot release name> <path to .deb> [...]
%prog [options] login <chroot path>|<chroot release name>
%prog [options] retrace <bugnumber>|<report file>''')

    optparser.add_option('--mirror', 
        help='Mirror for chroot creation',
        action='store', type='string', dest='mirror', metavar='URL', default=None)
    optparser.add_option('-a', '--apt-source', 
        help='Add an extra apt source',
        action='append', type='string', dest='extra_apt', metavar='SOURCE', default=[])
    optparser.add_option('-t', '--tar', 
        help='Create chroot tarball instead of directory',
        action='store_true', dest='tar', default=False)
    optparser.add_option('--save', 
        help='When logging in to a chroot tarball, update the tarball afterwards to save modifications if the shell exits with status 0.',
        action='store_true', dest='tar_save', default=False)
    optparser.add_option('-m', '--chroot-map', 
        help='Path to chroot map. This is a file that defines a Python dictionary, mapping DistroRelease: values to chroot paths',
        action='store', type='string', dest='chroot_map', metavar='FILE', default=None)
    optparser.add_option('-p', '--extra-package', 
        help='Install an extra package (can be specified multiple times)',
        action='append', type='string', dest='extra_packages', metavar='PACKAGE', default=[])
    optparser.add_option('-v', '--verbose', 
        help='Verbose operation (also passed to apport-retrace)',
        action='store_true', dest='verbose', default=False)
    optparser.add_option('--auth', 
        help='Passed on to apport-retrace in "retrace" mode',
        action='store', type='string', dest='auth_file', default=None)
    optparser.add_option('--duplicate-db', 
        help='Passed on to apport-retrace in "retrace" mode',
        action='store', type='string', dest='dup_db', default=None)
    optparser.add_option('--confirm-attach', 
        help='Display retraced stack traces and ask for confirmation before uploading them as bug attachments.',
        action='store_true', dest='confirm_attach', default=False)

    (opts, args) = optparser.parse_args()

    if len(args) < 1:
        optparser.error('no command specified (use --help for a short online help)')
        sys.exit(1)

    if opts.chroot_map:
        if not os.path.exists(opts.chroot_map):
            print >> sys.stderr, 'specified chroot map does not exist'
            sys.exit(1)

        # load chroot map and resolve relative paths
        map_file_dir = os.path.dirname(opts.chroot_map)
        opts.chroot_map = eval(open(opts.chroot_map).read(), {}, {})
        for n, p in opts.chroot_map.iteritems():
            if not p.startswith('/'):
                opts.chroot_map[n] = os.path.join(map_file_dir, p)

    return (opts, args)

def release_from_report(file):
    '''Return the distro release from the given Apport report.'''

    pr = problem_report.ProblemReport()
    pr.load(open(file), binary=False)
    return pr['DistroRelease']

def upgrade_chroot(chroot, verbose=False, extra_packages=[]):
    '''Update a chroot to the latest apt lists and packages.
    
    If run from a tarball and the dist-upgrade succeeds, then the tarball
    is updated as well. If the dist-upgrade fails, an assertion is raised.'''

    if verbose:
        assert chroot.run(['apt-get', 'update']) == 0
        assert chroot.run(['apt-get', '-y', '--allow-unauthenticated', 'dist-upgrade']) == 0
    else:
        assert chroot.run(['apt-get', '-qq', 'update']) == 0
        assert chroot.run(['apt-get', '-qqy', '--allow-unauthenticated', 'dist-upgrade']) == 0
    if extra_packages:
        assert chroot.run(['apt-get', 'install', '-y', '--allow-unauthenticated'] + extra_packages) == 0

    chroot.fix_symlinks()

    if chroot.root_tarball:
        assert chroot.run(['apt-get', 'clean']) == 0
        chroot.tar()

#
# commands
#

def command_create(opts, args):
    '''Create a chroot.'''

    if len(args) != 2:
        print >> sys.stderr, 'create needs exactly two arguments (use --help for a short online help)'
        sys.exit(1)
    (release, destpath) = args

    # create chroot directory
    if opts.tar:
        root = tempfile.mkdtemp()
        atexit.register(shutil.rmtree, root)
        if os.path.isfile(destpath):
            print >> sys.stderr, 'target file', destpath, 'exists already, aborting'
            sys.exit(1)
    else:
        root = destpath
        if os.path.isdir(root):
            print >> sys.stderr, 'target directory', root, 'exists already, aborting'
            sys.exit(1)
        os.makedirs(root)

    # call debootstrap
    setup_fakeroot_env()
    debootstrap_argv = ['debootstrap',
        '--variant=fakechroot', release, root]
    if opts.mirror:
        debootstrap_argv.append(opts.mirror)

    assert subprocess.call(debootstrap_argv) == 0

    # if we have a file:// mirror, create a symlink
    if opts.mirror and opts.mirror.startswith('file://'):
        mirrordir = os.path.abspath(opts.mirror[7:])
        targetdir = os.path.normpath(root + '/' + os.path.dirname(mirrordir))
        if not os.path.isdir(targetdir):
            os.makedirs(targetdir)
        os.symlink(mirrordir, os.path.join(targetdir, os.path.basename(mirrordir)))

    # set up apt sources
    if opts.mirror:
        f = open(os.path.join(root, 'etc', 'apt', 'sources.list'), 'w')
        print >> f, 'deb %s %s main' % (opts.mirror, release)
    else:
        # debootstrap puts default mirror there
        f = open(os.path.join(root, 'etc', 'apt', 'sources.list'), 'a')

    for s in opts.extra_apt:
        print >> f, s
    f.close()

    # disable invoke-rc.d
    policyrc = os.path.join(root, 'usr', 'sbin', 'policy-rc.d')
    open(policyrc, 'w').write('#!/bin/sh\nexit 101')
    os.chmod(policyrc, 0755)

    # set up apt-get and required packages
    chroot = Chroot(root)
    assert chroot.run(['apt-get', 'update']) == 0
    chroot.run(['apt-get', 'install', '-y', '--allow-unauthenticated', 'gpgv', 'apport-retrace'] + opts.extra_packages)

    chroot.fix_symlinks()

    # clean up cruft
    for path, dirs, files in os.walk(os.path.join(root, 'var', 'cache', 'apt', 'archives')):
        for f in files:
            try:
                os.unlink(os.path.join(path, f))
            except OSError:
                pass

    # tar it up
    if opts.tar:
        chroot.tar(destpath)

def command_upgrade(opts, args):
    '''Upgrade one or all chroots.'''

    if len(args) != 1:
        print >> sys.stderr, 'upgrade needs exactly one argument (use --help for a short online help)'
        sys.exit(1)
    if not opts.chroot_map and not os.path.exists(args[0]):
        print >> sys.stderr, 'you must specify a chroot map with -m (use --help for a short online help)'
        sys.exit(1)

    if args[0] == 'all':
        for c in opts.chroot_map.itervalues():
            print 'Upgrading %s...' % c
            upgrade_chroot(Chroot(c), opts.verbose, opts.extra_packages)
    elif os.path.exists(args[0]):
        upgrade_chroot(Chroot(args[0]), opts.verbose, opts.extra_packages)
    elif opts.chroot_map.has_key(args[0]):
        c = opts.chroot_map[args[0]]
        print 'Upgrading %s...' % c
        upgrade_chroot(Chroot(c), opts.verbose, opts.extra_packages)
    else:
        print >> sys.stderr, 'invalid chroot'

def command_installdeb(opts, args):
    '''Install a bunch of .debs into the chroot.'''

    if len(args) < 2:
        print >> sys.stderr, 'installdeb needs more arguments (use --help for a short online help)'
        sys.exit(1)
    if not opts.chroot_map and not os.path.exists(args[0]):
        print >> sys.stderr, 'you must specify a chroot map with -m (use --help for a short online help)'
        sys.exit(1)

    if os.path.exists(args[0]):
        chroot = Chroot(args[0])
    elif opts.chroot_map.has_key(args[0]):
        c = opts.chroot_map[args[0]]
        chroot = Chroot(c)
    else:
        print >> sys.stderr, 'invalid chroot'

    # symlink the debs into the chroot
    chroot_deb_dir = os.path.join(chroot.root, 'installdebs')
    os.makedirs(chroot_deb_dir)
    for deb in args[1:]:
        os.symlink(deb, os.path.join(chroot_deb_dir, os.path.basename(deb)))

    result = chroot.run(['dpkg', '-i'] + ['/installdebs/' + os.path.basename(d) for d in args[1:]])
    shutil.rmtree(chroot_deb_dir)

    chroot.fix_symlinks()

    if chroot.root_tarball and result == 0:
        chroot.tar()

def command_login(opts, args):
    '''Start a shell in a chroot.'''

    if len(args) != 1:
        print >> sys.stderr, 'login needs exactly one argument (use --help for a short online help)'
        sys.exit(1)
    if not opts.chroot_map and not os.path.exists(args[0]):
        print >> sys.stderr, 'you must specify a chroot map with -m (use --help for a short online help)'
        sys.exit(1)

    if os.path.exists(args[0]):
        chroot = Chroot(args[0])
    elif opts.chroot_map.has_key(args[0]):
        c = opts.chroot_map[args[0]]
        chroot = Chroot(c)
    else:
        print >> sys.stderr, 'invalid chroot'
        sys.exit(1)

    del os.environ['APPORT_CRASHDB_CONF']
    ret = chroot.run([os.environ.get('SHELL', 'bash')])

    if chroot.root_tarball and opts.tar_save and ret == 0:
        chroot.fix_symlinks()
        chroot.tar()

def command_retrace(opts, args):
    '''Retrace a bug or report file.'''

    if len(args) != 1:
        print >> sys.stderr, 'retrace needs exactly one argument (use --help for a short online help)'
        sys.exit(1)
    if not opts.chroot_map:
        print >> sys.stderr, 'you must specify a chroot map with -m (use --help for a short online help)'
        sys.exit(1)

    apport_retrace_argv = ['apport-retrace', '-u']
    if opts.verbose:
        apport_retrace_argv.append('-v')
    for p in opts.extra_packages:
	apport_retrace_argv += ['-p', p]

    if args[0].isdigit():
        crashdb = get_crashdb(opts.auth_file)
        release = crashdb.get_distro_release(args[0])
        chroot_path = opts.chroot_map[release]
        if os.path.isfile(chroot_path):
            apport_retrace_argv.append('--no-dpkg')
        c = Chroot(chroot_path)
        if opts.auth_file:
            # symlink the file into the chroot
            chroot_auth = os.path.join(c.root, 'tmp', 'auth')
            os.symlink(opts.auth_file, chroot_auth)
            apport_retrace_argv += ['--auth', chroot_auth]
        else:
            apport_retrace_argv += ['-s']
            chroot_auth = None
        if opts.dup_db:
            # symlink the file into the chroot
            chroot_dupdb = os.path.join(c.root, 'tmp', 'dup.db')
            os.symlink(opts.dup_db, chroot_dupdb)
            apport_retrace_argv += ['--duplicate-db', chroot_dupdb]
        if opts.confirm_attach:
            apport_retrace_argv.append('--confirm')
        apport_retrace_argv.append(args[0])

        del os.environ['APPORT_CRASHDB_CONF']
        c.run(apport_retrace_argv)
        if chroot_auth:
            os.unlink(chroot_auth)
    else:
        # symlink the report into the chroot
        release = release_from_report(args[0])
        chroot_path = opts.chroot_map[release]
        if os.path.isfile(chroot_path):
            apport_retrace_argv.append('--no-dpkg')
        c = Chroot(chroot_path)
        chroot_report = os.path.join(c.root, 'tmp', os.path.basename(args[0]))
        os.symlink(os.path.realpath(args[0]), chroot_report)
        apport_retrace_argv.append('-s')
        apport_retrace_argv.append(chroot_report)
        c.run(apport_retrace_argv)
        os.unlink(chroot_report)

#
# main
#

opts, args = parse_options()
try:
    command = globals()['command_' + args[0]]
except KeyError:
    print >> sys.stderr, 'unknown command (use --help for a short online help)'
    sys.exit(1)
command(opts, args[1:])
