#!/usr/bin/python

# Use the coredump in a crash report to regenerate the stack traces. This is
# helpful to get a trace with debug symbols.
#
# Copyright (c) 2006 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 sys, os, os.path, subprocess, optparse, shutil, tempfile, glob, gzip, re
import zlib
import tty, termios
import apport, apport.fileutils

import warnings
warnings.filterwarnings("ignore", "apt API not stable yet", FutureWarning)
import apt

from apport.crashdb import get_crashdb

#
# functions
#

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

    optparser = optparse.OptionParser('%prog [options] <apport problem report | crash ID>')
    optparser.add_option('-s', '--stdout', 
        help='Do not put the new traces into the report, but write them to stdout.',
        action='store_true', dest='stdout', default=False)
    optparser.add_option('-g', '--gdb', 
        help="Start an interactive gdb session with the report's core dump (-o ignored; does not rewrite report)",
        action='store_true', dest='gdb', default=False)
    optparser.add_option('-o', '--output', 
        help='Write modified report to given file instead of changing the original report',
        action='store', type='string', dest='output', metavar='FILE', default=None)
    optparser.add_option('-c', '--remove-core', 
        help='Remove the core dump from the report after stack trace regeneration',
        action='store_true', dest='removecore', default=False)
    optparser.add_option('-r', '--core-file', 
        help="Override report's CoreFile",
        action='store', type='string', dest='core_file', metavar='CORE', default=None)
    optparser.add_option('-x', '--executable', 
        help="Override report's ExecutablePath",
        action='store', type='string', dest='executable', metavar='EXE', default=None)
    optparser.add_option('-m', '--procmaps', 
        help="Override report's ProcMaps",
        action='store', type='string', dest='procmaps', metavar='MAPS', default=None)
    optparser.add_option('-R', '--rebuild-package-info', 
        help="Rebuild report's Package information",
        action='count', dest='rebuild', default=0)
    optparser.add_option('-v', '--verbose', 
        help='Report download/install progress when installing additional packages',
        action='count', dest='verbose', default=0)
    optparser.add_option('-u', '--unpack-only', 
        help='Only unpack the additionally required packages, do not configure them; purge packages again after retracing',
        action='store_true', dest='unpack_only', default=False)
    optparser.add_option('--no-dpkg', 
        help='Do not use dpkg when using -u and do not purge packages afterwards. This should only be used for temporarily unpacked chroot tarballs where it would just be a waste of time.',
        action='store_true', dest='no_dpkg', default=False)
    optparser.add_option('-p', '--extra-package', 
        help='Install an extra package (can be specified multiple times)',
        action='append', type='string', dest='extra_packages', default=[])
    optparser.add_option('--auth',
        help='Path to a file with the crash database authentication information. This is used when specifying a crash ID to upload the retraced stack traces (only if neither -g, -o, nor -s are specified)',
        action='store', type='string', dest='auth_file', default=None)
    optparser.add_option('--confirm',
        help='Display retraced stack traces and ask for confirmation before sending them to the crash database.',
        action='store_true', dest='confirm', default=False)
    optparser.add_option('--duplicate-db',
        help='Path to the duplicate sqlite database (default: no duplicate checking)',
        action='store', type='string', dest='dup_db', metavar='PATH', default=None)

    (opts, args) = optparser.parse_args()

    if len(args) != 1:
        optparser.error('incorrect number of arguments; use --help for a short online help')
        sys.exit(1)

    if not (opts.auth_file or opts.output or opts.stdout or opts.gdb):
        optparser.error('you either need to do a local operation (-s, -g, -o) or supply an authentication file (--auth); see --help for a short online help')
        sys.exit(1)

    return (args[0], opts)

def deb_without_preinst(deb):
    '''If given .deb file has a preinst script, generate a <name>_noscript.deb
    file without it and return that name; otherwise, return deb.
    
    If the modified deb already exists, its name is returned without recreating
    it.'''

    ndeb = '/var/cache/apt/archives/%s_noscript%s' % os.path.splitext(os.path.basename(deb))

    if os.path.exists(ndeb):
        return ndeb

    # get control.tar.gz    
    ar = subprocess.Popen(['ar', 'p', deb, 'control.tar.gz'], stdout=subprocess.PIPE)
    control_tar = ar.communicate()[0]
    assert ar.returncode == 0

    # check if package has a preinst
    tar = subprocess.Popen(['tar', 'tz', './preinst'], stdin=subprocess.PIPE,
        stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    tar.communicate(control_tar)
    if tar.returncode != 0:
        return deb

    # unpack control.tar.gz and remove scripts
    d = tempfile.mkdtemp()
    d2 = tempfile.mkdtemp()
    try:
        tar = subprocess.Popen(['tar', '-C', d, '-xz'], stdin=subprocess.PIPE)
        tar.communicate(control_tar)
        assert tar.returncode == 0
        for s in ('preinst', 'postinst', 'prerm', 'postrm'):
            path = os.path.join(d, s)
            if os.path.exists(path):
                os.unlink(path)

        control_tar_new = os.path.join(d2, 'control.tar.gz')
        tar = subprocess.Popen(['tar', '-C', d, '-cz', '.'],
            stdin=subprocess.PIPE, stdout=open(control_tar_new, 'w'))
        assert tar.wait() == 0

        shutil.copy(deb, ndeb)
        r = subprocess.Popen(['ar', 'r', ndeb, control_tar_new])
        assert r.wait() == 0
    finally:
        shutil.rmtree(d)
        shutil.rmtree(d2)

    return ndeb

def fetch_unpack(cache, fetchProgress, no_dpkg=False, verbosity=0):
    '''Fetch the packages marked to be installed in the given apt.Cache()
    object and unpack them.
    
    fetchProgress must be a valid apt.progress.FetchProgress object.'''

    # fetch
    fetcher = apt.apt_pkg.GetAcquire(fetchProgress)
    pm = apt.apt_pkg.GetPackageManager(cache._depcache)
    try:
        res = cache._fetchArchives(fetcher, pm)
    except IOError, e:
        print >> sys.stderr, 'ERROR: could not fetch all archives:', e

    # extract
    if verbosity:
        so = sys.stderr
    else:
        so = subprocess.PIPE
    if no_dpkg:
        for i in fetcher.Items:
            if verbosity:
                print 'Extracting', i.DestFile
            if subprocess.call(['dpkg', '-x', i.DestFile, '/'], stdout=so,
                stderr=subprocess.STDOUT) != 0:
                print >> sys.stderr, 'WARNING: %s failed to extract' % i.DestFile
    else:
        res = subprocess.call(['dpkg', '--force-depends', '--force-overwrite', '--unpack'] + 
            [deb_without_preinst(i.DestFile) for i in fetcher.Items], stdout=so)
        if res != 0:
            raise IOError, 'dpkg failed to unpack archives'

    # remove other maintainer scripts
    for c in cache.getChanges():
        for script in ('postinst', 'prerm', 'postrm'):
            try:
                os.unlink('/var/lib/dpkg/info/%s.%s' % (c.name, script))
            except OSError:
                pass

def install_missing_packages(report, verbosity = 0, unpack_only=False,
    no_dpkg=False, extra_packages = []):
    '''Install packages which are required to retrace the given report.
    
    If package installation fails (e. g. because the user does not have root
    privileges), the list of required packages is printed out instead.
    
    Return a list of installed packages.'''

    c = apt.Cache()

    installed = []

    # create map of dependency package versions as specified in report
    dependency_versions = {}
    for l in (report['Package'] + '\n' + report.get('Dependencies', '')).splitlines():
        if not l.strip():
            continue
        (pkg, version) = l.split()[:2]
        dependency_versions[pkg] = version
        try:
            # this fails for packages which are still installed, but gone from
            # the archive; i. e. /var/lib/dpkg/status still knows about them
            if not c[pkg]._lookupRecord():
                raise KeyError
            if not 'Architecture: all' in c[pkg]._records.Record:
                dependency_versions[pkg+'-dbgsym'] = dependency_versions[pkg]
        except KeyError:
            print >> sys.stderr, 'WARNING: package %s not known to package cache' % pkg

    for pkg, ver in dependency_versions.iteritems():
	if not c.has_key(pkg):
	    print >> sys.stderr, 'WARNING: package %s not available' % pkg
	    continue

	# ignore packages which are already installed in the right version
	if (ver and c[pkg].installedVersion == ver) or \
	   (not ver and c[pkg].installedVersion):
	   continue

	if ver and c[pkg].candidateVersion != ver:
	    print >> sys.stderr, 'WARNING: %s version %s required, but %s is available' % (
		pkg, ver, c[pkg].candidateVersion)

	c[pkg].markInstall(False)

    # extra packages
    for p in extra_packages:
        c[p].markInstall(False)

    if verbosity:
        fetchProgress = apt.progress.TextFetchProgress()
        installProgress = apt.progress.InstallProgress()
    else:
        fetchProgress = apt.progress.FetchProgress()
        installProgress = apt.progress.DumbInstallProgress()

    try:
	if c.getChanges():
	    os.environ['DEBIAN_FRONTEND'] = 'noninteractive'
            if unpack_only:
                fetch_unpack(c, fetchProgress, no_dpkg, verbosity)
            else:
                try:
                    c.commit(fetchProgress, installProgress)
                except SystemError:
                    print >> sys.stderr, 'Error: Could not install all archives. If you use this tool on a production system, it is recommended to use the -u option. See --help for details.'
                    sys.exit(1)

	    # after commit(), the Cache object does not empty the pending
	    # changes, so we need to reinitialize it to avoid applying the same
	    # changes again below
            installed = [p.name for p in c.getChanges()]
	    c = apt.Cache()
    except IOError, e:
	pass # we will complain to the user later

    # check list of libraries that the crashed process referenced at
    # runtime and warn about those which are not available
    libs = set()
    if report.has_key('ProcMaps'):
        for l in report['ProcMaps'].splitlines():
            if not l.strip():
                continue
            cols = l.split()
            if 'x' in cols[1] and len(cols) == 6 and '.so' in cols[5]:
                lib = os.path.realpath(cols[5])
                libs.add(lib)

    # grab as much as we can
    for l in libs:
        if os.path.exists('/usr/lib/debug' + l):
            continue

        pkg = apport.packaging.get_file_package(l, True)
	if pkg:
            if not os.path.exists(l):
                if c.has_key(pkg):
                    c[pkg].markInstall(False)
                else:
                    print >> sys.stderr, 'WARNING: %s was loaded at runtime, but its package %s is not available' % (l, pkg)

	    if c.has_key(pkg+'-dbgsym'):
		c[pkg+'-dbgsym'].markInstall(False)
	    else:
		print >> sys.stderr, 'WARNING: %s-dbgsym is not available' % pkg
	else:
		print >> sys.stderr, 'WARNING: %s is needed, but cannot be mapped to a package' % l

    try:
	if c.getChanges():
	    os.environ['DEBIAN_FRONTEND'] = 'noninteractive'
            if unpack_only:
                fetch_unpack(c, fetchProgress, no_dpkg, verbosity)
            else:
                c.commit(fetchProgress, installProgress)
        installed += [p.name for p in c.getChanges()]
    except (SystemError, IOError), e:
	print >> sys.stderr, 'WARNING: could not install missing packages:', e
        if os.geteuid() != 0:
            print >> sys.stderr, 'You either need to call this program as root or install these packages manually:'
	for p in c.getChanges():
	    print >> sys.stderr, '  %s %s' % (p.name, p.candidateVersion)

    return installed

def getch():
    '''Read a single character from stdin.'''

    fd = sys.stdin.fileno()
    old_settings = termios.tcgetattr(fd)
    try:
        tty.setraw(sys.stdin.fileno())
        ch = sys.stdin.read(1)
    finally:
        termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)
    return ch

def confirm_traces(report):
    '''Display the retraced stack traces and ask the user whether or not to
    upload them to the crash database.
    
    Return True if the user agrees.'''

    print '--- stack trace ---'
    print report['Stacktrace']
    print '--- thread stack trace ---'
    print report['ThreadStacktrace']
    print '---'

    ch = None
    while ch not in ['y', 'n']:
        print 'OK to send these as attachments? [y/n]'
        ch = getch().lower()

    return ch == 'y'

def find_file_dir(name, dir, limit=None):
    '''Return a path list of all files with given name which are in or below
    dir.
    
    If limit is not None, the search will be stopped after finding the given
    number of hits.'''

    result = []
    for root, dirs, files in os.walk(dir):
        if name in files:
            result.append(os.path.join(root, name))
            if limit and len(result) >= limit:
                break
    return result

def get_code(srcdir, filename, line, context=5):
    '''Find the given filename in the srcdir directory and return the code
    lines around the given line number.'''

    files = find_file_dir(filename, srcdir, 1)
    if not files:
        return '  [Error: %s was not found in source tree]\n' % filename

    result = ''
    lineno = 0
    # make enough room for all line numbers
    format = '  %%%ii: %%s' % len(str(line+context))

    for l in open(files[0]):
        lineno += 1
        if lineno >= line-context and lineno <= line+context:
            result += format % (lineno, l)

    return result

def gen_source_stacktrace(report):
    '''Generate StacktraceSource, a version of Stacktrace with the surrounding
    code lines (where available) and with local variables removed.'''

    if not report.has_key('Stacktrace') or not report.has_key('SourcePackage'):
        return

    workdir = tempfile.mkdtemp()
    try:
        try:
            version = report['Package'].split()[1]
        except (IndexError, KeyError):
            version = None
        srcdir = apport.packaging.get_source_tree(report['SourcePackage'], workdir, version)
        if not srcdir:
            return

        src_frame = re.compile('^#\d+\s.* at (.*):(\d+)$')
        other_frame = re.compile('^#\d+')
        result = ''
        for frame in report['Stacktrace'].splitlines():
            m = src_frame.match(frame)
            if m:
                result += frame + '\n' + get_code(srcdir, m.group(1), int(m.group(2)))
            else:
                m = other_frame.search(frame)
                if m:
                    result += frame + '\n'

        report['StacktraceSource'] = result
    finally:
        shutil.rmtree(workdir)
        pass

#
# main
#

(reportfile, options) = parse_options()

crashdb = get_crashdb(options.auth_file)

# load the report
if os.path.exists(reportfile):
    try:
        report = apport.Report()
        report.load(open(reportfile))
    except (MemoryError, TypeError, ValueError, IOError, zlib.error), e:
        print >> sys.strerr, 'Cannot open report file:', e
        sys.exit(1)
elif reportfile.isdigit():
    # crash ID
    report = crashdb.download(int(reportfile))
    crashid = reportfile
    reportfile = None
else:
    print >> sys.stderr, 'ERROR: "%s" is neither an existing report file nor a \
crash ID' % reportfile
    sys.exit(1)

if options.core_file:
    report['CoreDump'] = file(options.core_file).read()
if options.executable:
    report['ExecutablePath'] = options.executable
if options.procmaps:
    report['ProcMaps'] = file(options.procmaps).read()
if options.rebuild:
    report.add_package_info()

# sanity checks
required_fields = set(['CoreDump', 'ExecutablePath', 'Package'])
if not required_fields.issubset(set(report.keys())):
    print >> sys.stderr, 'report file does not contain required fields: ' + \
	' '.join(required_fields)
    sys.exit(1)

installed = install_missing_packages(report, options.verbose,
    options.unpack_only, options.no_dpkg, options.extra_packages)

# interactive gdb session
if options.gdb:
    core = tempfile.NamedTemporaryFile()
    core.write(report['CoreDump'])
    core.flush()
    subprocess.call(['gdb',
        '--ex', 'file ' + report.get('InterpreterPath', report['ExecutablePath']),
        '--ex', 'core-file ' + core.name])
    core.close()
else:
    # regenerate gdb info
    report.add_gdb_info()
    gen_source_stacktrace(report)

modified = False

if options.removecore:
    del report['CoreDump']
    modified = True

if options.stdout:
    print '--- stack trace ---'
    print report['Stacktrace']
    print '--- thread stack trace ---'
    print report['ThreadStacktrace']
    if report.has_key('StacktraceSource'):
        print '--- source code stack trace ---'
        print report['StacktraceSource']
else:
    if not options.gdb:
        modified = True

if modified:
    if not reportfile and not options.output:
        if not options.auth_file:
            print >> sys.stderr, 'You need to specify --auth for uploading retraced results back to the crash database.'
            sys.exit(1)
        if not options.confirm or confirm_traces(report):
            if not report.crash_signature():
                print 'Report has no crash signature, so retrace is flawed'
                crashdb.mark_retrace_failed(crashid)

            # check for duplicates
            update_bug = True
            if options.dup_db:
                crashdb.init_duplicate_db(options.dup_db)
                res = crashdb.check_duplicate(int(crashid), report)
                if res:
                    if res[1] == None:
                        print 'Report is a duplicate of #%i (not fixed yet)' % res[0]
                    elif res[1] == '':
                        print 'Report is a duplicate of #%i (fixed in latest version)' % res[0]
                    else:
                        print 'Report is a duplicate of #%i (fixed in version %s)' % res
                    update_bug = False
                else:
                    print 'Duplicate check negative'

            if update_bug:
                crashdb.update(crashid, report)
                print 'New attachments uploaded to crash database #' + crashid
    else:
        if options.output == None:
            out = open(reportfile, 'w')
        elif options.output == '-':
            out = sys.stdout
        else:
            out = open(options.output, 'w')

        report.write(out)

if options.unpack_only and installed and not options.no_dpkg:
    if options.verbose:
        so = sys.stderr
    else:
        so = subprocess.PIPE
    subprocess.call(['dpkg', '-P'] + installed, stdout=so)
