#!/usr/bin/python

# Copyright (C) 2007 - 2009 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 tempfile, shutil, os.path, os, subprocess, signal, time, stat, sys
import resource, errno
import apport.fileutils

test_executable = '/bin/cat'
test_package = 'coreutils'
test_source = 'coreutils'

required_fields = ['ProblemType', 'CoreDump', 'Date', 'ExecutablePath',
    'ProcCmdline', 'ProcEnviron', 'ProcMaps', 'Signal', 'UserGroups']

ifpath = os.path.expanduser(apport.report._ignore_file)
orig_ignore_file = None

def run(argv, msg=None, stdout=False):
    '''Execute the given program with arguments and verify that it exits
    successfully. 
    
    Set stdout to True to see the program's standard output. msg is the
    assertion error message in case of failure.'''

    if stdout:
        result = subprocess.call(argv)
    else:
        result = subprocess.call(argv, stdout=subprocess.PIPE)
    assert result == 0, msg + ' [cwd: %s]' % os.getcwd()


def create_test_process(check_running=True):
    assert os.access(test_executable, os.X_OK), test_executable + ' is not executable'
    if check_running:
        assert subprocess.call(['pidof', test_executable]) == 1, 'no running test executable processes'
    pid = os.fork()
    if pid == 0:
        sys.stdin.close()
        os.setsid()
        os.execv(test_executable, [test_executable])
        assert False, 'Could not execute ' + test_executable

    # wait until child process has execv()ed properly
    while 'test-apport' in open('/proc/%i/cmdline' % pid).read():
        time.sleep(0.1)
    time.sleep(0.3) # needs some more setup time
    return pid

def check_crash(expect_coredump=True, expect_corefile=False, sig=signal.SIGSEGV, check_running=True, sleep=0):
    assert not os.path.exists('core'), 'no core dump in current directory'
    pid = create_test_process(check_running)
    if sleep > 0:
        time.sleep(sleep)
    os.kill(pid, sig)
    result = os.waitpid(pid, 0)[1]
    assert not os.WIFEXITED(result), 'test process did not exit normally'
    assert os.WIFSIGNALED(result), 'test process died due to signal'
    if expect_coredump:
        assert os.WCOREDUMP(result), 'result: 0x%02X should include WCOREDUMP' % result
    else:
        assert not os.WCOREDUMP(result), 'result: 0x%02X should not have WCOREDUMP' % result
    assert os.WSTOPSIG(result) == 0, 'test process was not signaled to stop'
    assert os.WTERMSIG(result) == sig, 'test process died due to proper signal'
    # wait max 10 seconds for apport to finish
    timeout = 50
    while timeout >= 0:
	if subprocess.call(['pidof', '-x', 'apport'], stdout=subprocess.PIPE) != 0:
	    break
	time.sleep(0.2)
	timeout -= 1
    assert subprocess.call(['pidof', '-x', 'apport']) == 1, 'no running apport processes'
    if check_running:
        assert subprocess.call(['pidof', test_executable]) == 1, 'no running test executable processes'
    if expect_corefile:
        assert os.path.exists('core'), 'leaves wanted core file'
        try:
            # check that core file is valid
            gdb = subprocess.Popen(['gdb', '--batch', '--ex', 'bt',
                test_executable, 'core'], stdout=subprocess.PIPE,
                stderr=subprocess.PIPE)
            (out, err) = gdb.communicate()
            assert gdb.returncode == 0
            err = err.strip()
            assert err == '' or err.startswith('warning'), err
        finally:
            os.unlink('core')
    else:
        assert not os.path.exists('core'), 'leaves unexpected core file behind'

def check_safe_environment(env):
    allowed_vars = ['SHELL', 'PATH', 'LANGUAGE', 'LANG', 'LC_CTYPE',
        'LC_COLLATE', 'LC_TIME', 'LC_NUMERIC', 'LC_MONETARY', 'LC_MESSAGES',
        'LC_PAPER', 'LC_NAME', 'LC_ADDRESS', 'LC_TELEPHONE', 'LC_MEASUREMENT',
        'LC_IDENTIFICATION', 'LOCPATH']

    for l in env.splitlines():
        (k, v) = l.split('=', 1)
        assert k in allowed_vars, '%s is insensitive environment variable' % k

def get_tempdir_reports(dir):
    '''Call apport.fileutils.get_all_reports() for a temporary report
    directory.'''

    old_dir = apport.fileutils.report_dir
    apport.fileutils.report_dir = dir
    reports = apport.fileutils.get_all_reports()
    apport.fileutils.report_dir = old_dir
    return reports

#
# main
#

all_passed = False

# we sometimes need a local report dir because we cannot access root's .lock file
temp_report_dir = tempfile.mkdtemp() 

try:
    core_pattern = open('/proc/sys/kernel/core_pattern').read().strip()
    assert core_pattern[0] == '|', \
            'kernel crash dump helper is not active; please enable before running this test.'
    apport_path = core_pattern[1:].split()[0]

    assert apport.fileutils.get_all_reports() == [], 'please remove all crash reports from /var/crash/ for this test suite'

    # do not write core files for the first tests
    resource.setrlimit(resource.RLIMIT_CORE, (0, -1))

    # move aside current ignore file
    if os.path.exists(ifpath):
	orig_ignore_file = ifpath + '.apporttest'
	os.rename(ifpath, orig_ignore_file)

    print '* empty core dumps do not generate a report'
    test_proc = create_test_process()
    try:
        app = subprocess.Popen([apport_path, str(test_proc), '42', '0'], 
            env={'APPORT_REPORT_DIR': temp_report_dir}, close_fds=True,
            stdin=subprocess.PIPE, stderr=subprocess.PIPE)
        app.stdin.close()
        assert app.wait() == 0, app.stderr.read()
    finally:
        os.kill(test_proc, 9)
        os.waitpid(test_proc, 0)

    assert get_tempdir_reports(temp_report_dir) == [], 'no report created for empty core dump'

    print '* check test process creation/killing with apport'
    check_crash()

    # check crash report
    report = os.path.join(apport.fileutils.report_dir, '%s.%i.crash' %
        (test_executable.replace('/', '_'), os.getuid()))
    assert apport.fileutils.get_all_reports() == [report], 'report was created'
    st = os.stat(report)
    assert stat.S_IMODE(st.st_mode) == 0600, 'report has correct permissions'

    print '* a subsequent crash does not alter unseen report'
    check_crash()
    st2 = os.stat(report)
    assert st == st2, 'original unread report did not change'

    print '* a subsequent crash alters seen report'
    apport.fileutils.mark_report_seen(report)
    check_crash()
    st2 = os.stat(report)
    assert st != st2, 'original read report changed'

    print '* report has required fields'
    pr = apport.Report()
    pr.load(open(report))
    assert set(required_fields).issubset(set(pr.keys())), 'report has required fields'
    assert pr['ExecutablePath'] == test_executable
    assert pr['ProcCmdline'] == test_executable
    assert pr['Signal'] == '%i' % signal.SIGSEGV

    print '* dumped environment only has insensitive variables'
    check_safe_environment(pr['ProcEnviron'])

    print '* collected system groups has nonempty user groups information'
    assert pr['UserGroups']
    print '* collected system groups are not those from root'
    assert 'root' not in pr['UserGroups']

    apport.fileutils.delete_report(report)
    assert apport.fileutils.get_all_reports() == [], \
        'no reports present after deleting (present: %s)' % str(apport.fileutils.get_all_reports())

    print '* only one apport instance is ran at a time'
    test_proc = create_test_process()
    test_executable = '/bin/dd'
    test_proc2 = create_test_process(False)
    test_executable = '/bin/cat'
    try:
        app = subprocess.Popen([apport_path, str(test_proc), '42', '0'], 
            env={'APPORT_REPORT_DIR': temp_report_dir}, close_fds=True,
            stdin=subprocess.PIPE, stderr=subprocess.PIPE)

        app2 = subprocess.Popen([apport_path, str(test_proc2), '42', '0'], 
            env={'APPORT_REPORT_DIR': temp_report_dir}, close_fds=True,
            stdin=subprocess.PIPE, stderr=subprocess.PIPE)

        # app should wait indefinitely for stdin, while app2 should terminate
        # immediately (give it 5 seconds)
        timeout = 50
        while timeout >= 0:
            if app2.poll():
                break

            time.sleep(0.1)
            timeout -= 1

        assert timeout > 0, 'second apport instance terminates immediately'
        assert not app.poll(), 'first apport instance is still running'

        # properly terminate app and app2
        app2.stdin.close()
        app.stdin.write('boo')
        app.stdin.close()

        assert app.wait() == 0, app.stderr.read()
    finally:
        os.kill(test_proc, 9)
        os.waitpid(test_proc, 0)
        os.kill(test_proc2, 9)
        os.waitpid(test_proc2, 0)

    for r in get_tempdir_reports(temp_report_dir):
        apport.fileutils.delete_report(r)

    print '* existing .lock file as dangling symlink does not create the file'
    # prepare a symlink trap
    lockpath = os.path.join(temp_report_dir, '.lock')
    trappath = os.path.join(temp_report_dir, '0wned')
    os.unlink(lockpath)
    assert not os.path.exists(trappath)
    os.symlink(trappath, lockpath)

    # now call apport
    test_proc = create_test_process()

    try:
        app = subprocess.Popen([apport_path, str(test_proc), '42', '0'], 
            env={'APPORT_REPORT_DIR': temp_report_dir}, close_fds=True,
            stdin=subprocess.PIPE, stderr=subprocess.PIPE)
        app.stdin.write('boo')
        app.stdin.close()

        assert app.wait() != 0, app.stderr.read()
    finally:
        os.kill(test_proc, 9)
        os.waitpid(test_proc, 0)

    assert get_tempdir_reports(temp_report_dir) == [], 'no report created'
    assert not os.path.exists(trappath), 'dangling .lock symlink created ' + trappath
    # remove our trap again
    os.unlink(lockpath)
    # remove crash reports about apport itself (when the packaged version is
    # ran from the test suite)
    for r in apport.fileutils.get_all_reports():
        apport.fileutils.delete_report(r)

    print '* non-packaged executables do not create a report'
    orig_test_executable = test_executable
    (fd, test_executable) = tempfile.mkstemp()
    os.write(fd, open(orig_test_executable).read())
    os.close(fd)
    os.chmod(test_executable, 0755)
    check_crash()
    os.unlink(test_executable)
    test_executable = orig_test_executable
    assert apport.fileutils.get_all_reports() == [], \
        'no reports present after a non-packaged crashed process (present: %s)' % \
        str(apport.fileutils.get_all_reports())

    print '* apport ignores SIGQUIT'
    check_crash(sig=signal.SIGQUIT)
    assert apport.fileutils.get_all_reports() == [], \
        'no reports present after SIGQUIT (present: %s)' % \
        str(apport.fileutils.get_all_reports())

    print '* existence of user-inaccessible files does not leak'
    orig_test_executable = test_executable
    dir = tempfile.mkdtemp()
    test_executable = os.path.join(dir, 'python') # set name to existing tool != self
    exec_file = open(test_executable, 'w')
    exec_file.write('#!/usr/bin/perl\nsystem("mv $0 $0.exe");\nsystem("ln -sf /etc/shadow $0");\n$0="..$0";\nsleep(10);\n')
    exec_file.close()
    os.chmod(test_executable, 0755)
    check_crash(check_running=False, sleep=2)
    os.unlink(test_executable)
    os.unlink(test_executable+".exe")
    os.rmdir(dir)
    leak = os.path.join(apport.fileutils.report_dir, '_usr_bin_perl.%i.crash' %
        (os.getuid()))
    pr = apport.Report()
    pr.load(open(leak))
    # On a leak, no report is created since the executable path will be replaced
    # by the symlink path, and it doesn't belong to any package.
    assert pr['ExecutablePath'] == '/usr/bin/perl'
    assert not pr.has_key('InterpreterPath')
    apport.fileutils.delete_report(leak)
    test_executable = orig_test_executable

    print '* non-packaged scripts do not create a report'
    orig_test_executable = test_executable
    (fd, test_executable) = tempfile.mkstemp()
    os.write(fd, '#!/bin/sh\nkill -SEGV $$')
    os.close(fd)
    os.chmod(test_executable, 0755)
    check_crash()
    assert apport.fileutils.get_all_reports() == [], \
        'no reports present after a non-packaged crashed script (present: %s)' % \
        str(apport.fileutils.get_all_reports())

    # relative path case
    old_cwd = os.getcwd()
    try:
        os.chdir(os.path.dirname(test_executable))
        test_executable = './' + os.path.basename(test_executable)
        check_crash()
        os.unlink(test_executable)
        test_executable = orig_test_executable
    finally:
        os.chdir(old_cwd)
    assert apport.fileutils.get_all_reports() == [], \
        'no reports present after a non-packaged crashed script with relative path(present: %s)' % \
        str(apport.fileutils.get_all_reports())

    print '* limitation of flooding: iteration',
    count = 0
    while count < 7:
        print count,
        sys.stdout.flush()
        check_crash()
        if not apport.fileutils.get_new_reports():
            break
        apport.fileutils.mark_report_seen(apport.fileutils.get_new_reports()[0])
        count += 1
    print
    assert count >= 1, 'gets at least 2 repeated crashes'
    assert count < 7, 'stops flooding after less than 7 repeated crashes'

    apport.fileutils.delete_report(report)
    assert apport.fileutils.get_all_reports() == [], \
        'no reports present after deleting (present: %s)' % str(apport.fileutils.get_all_reports())

    print '* core dump works for non-writable cwds'
    old_cwd = os.getcwd()
    try:
        os.chdir('/')
        check_crash()
    finally:
        os.chdir(old_cwd)
    pr = apport.Report()
    assert os.path.exists(report)
    pr.load(open(report))
    assert set(required_fields).issubset(set(pr.keys()))
    apport.fileutils.delete_report(report)
    assert apport.fileutils.get_all_reports() == [], \
        'no reports present after deleting (present: %s)' % str(apport.fileutils.get_all_reports())

    print '* non-packaged executables create core dumps on proper ulimits'
    orig_test_executable = test_executable
    (fd, test_executable) = tempfile.mkstemp()
    os.write(fd, open(orig_test_executable).read())
    os.close(fd)

    os.chmod(test_executable, 0755)
    resource.setrlimit(resource.RLIMIT_CORE, (1, -1))
    check_crash(expect_coredump=False, expect_corefile=False)
    apport.fileutils.delete_report(report)
    resource.setrlimit(resource.RLIMIT_CORE, (10, -1))
    check_crash(expect_corefile=False)
    apport.fileutils.delete_report(report)
    resource.setrlimit(resource.RLIMIT_CORE, (10000, -1))
    check_crash(expect_corefile=True)
    apport.fileutils.delete_report(report)
    resource.setrlimit(resource.RLIMIT_CORE, (-1, -1))
    check_crash(expect_corefile=True)
    apport.fileutils.delete_report(report)

    print '* non-packaged executables create core dumps on proper ulimits for SIGABRT'
    os.chmod(test_executable, 0755)
    resource.setrlimit(resource.RLIMIT_CORE, (1, -1))
    check_crash(expect_coredump=False, expect_corefile=False, sig=signal.SIGABRT)
    apport.fileutils.delete_report(report)
    resource.setrlimit(resource.RLIMIT_CORE, (10, -1))
    check_crash(expect_corefile=False, sig=signal.SIGABRT)
    apport.fileutils.delete_report(report)
    resource.setrlimit(resource.RLIMIT_CORE, (10000, -1))
    check_crash(expect_corefile=True, sig=signal.SIGABRT)
    apport.fileutils.delete_report(report)
    resource.setrlimit(resource.RLIMIT_CORE, (-1, -1))
    check_crash(expect_corefile=True, sig=signal.SIGABRT)
    apport.fileutils.delete_report(report)

    os.unlink(test_executable)
    test_executable = orig_test_executable

    print '* packaged executables create core dumps on proper ulimits'
    resource.setrlimit(resource.RLIMIT_CORE, (1, -1))
    check_crash(expect_coredump=False, expect_corefile=False)
    apport.fileutils.delete_report(report)
    resource.setrlimit(resource.RLIMIT_CORE, (10, -1))
    check_crash(expect_corefile=False)
    apport.fileutils.delete_report(report)
    resource.setrlimit(resource.RLIMIT_CORE, (10000, -1))
    check_crash(expect_corefile=True)
    apport.fileutils.delete_report(report)
    resource.setrlimit(resource.RLIMIT_CORE, (-1, -1))
    check_crash(expect_corefile=True)

    print '* crashes create core dumps with an existing crash report'
    check_crash(expect_corefile=True)
    apport.fileutils.delete_report(report)

    print '* packaged executables create core dumps on proper ulimits for SIGABRT'
    resource.setrlimit(resource.RLIMIT_CORE, (1, -1))
    check_crash(expect_coredump=False, expect_corefile=False, sig=signal.SIGABRT)
    resource.setrlimit(resource.RLIMIT_CORE, (10, -1))
    check_crash(expect_corefile=False, sig=signal.SIGABRT)
    resource.setrlimit(resource.RLIMIT_CORE, (10000, -1))
    check_crash(expect_corefile=True, sig=signal.SIGABRT)
    apport.fileutils.delete_report(report)
    resource.setrlimit(resource.RLIMIT_CORE, (-1, -1))
    check_crash(expect_corefile=True, sig=signal.SIGABRT)
    apport.fileutils.delete_report(report)

    print '* core dumps are capped on available memory size',
    # determine how much data we have to pump into apport in order to make sure
    # that it will refuse the core dump
    r = apport.Report()
    r.load(open('/proc/meminfo'))
    totalmb = long(r['MemFree'].split()[0]) + long(r['Cached'].split()[0])
    totalmb /= 1024
    r = None

    test_proc = create_test_process()
    try:
        app = subprocess.Popen([apport_path, str(test_proc), '42', '0'], 
            env={'APPORT_REPORT_DIR': temp_report_dir}, close_fds=True,
            stdin=subprocess.PIPE, stderr=subprocess.PIPE)
        # pipe an entire total memory size worth of spaces into it, which must be
        # bigger than the 'usable' memory size. apport should digest that and the
        # report should not have a core dump; NB that this should error out
        # with a SIGPIPE when apport aborts reading from stdin
        onemb = ' ' * 1048576
        while totalmb > 0:
            if totalmb & 31 == 0:
                sys.stdout.write('.')
                sys.stdout.flush()
            try:
                app.stdin.write(onemb)
            except IOError, e:
                if e.errno == errno.EPIPE:
                    break
                else:
                    raise
            totalmb -= 1
        print
        app.stdin.close()
        assert app.wait() == 0, app.stderr.read()
        onemb = None
    finally:
        os.kill(test_proc, 9)
        os.waitpid(test_proc, 0)

    reports = get_tempdir_reports(temp_report_dir)
    assert len(reports) == 1, 'report was created'

    pr = apport.Report()
    pr.load(open(reports[0]))
    os.unlink(reports[0])

    assert pr['Signal'] == '42'
    assert pr['ExecutablePath'] == test_executable
    assert not pr.has_key('CoreDump'), 'report does not have core dump'


    print '* binary blacklisting works'
    pr.mark_ignore()
    check_crash()
    assert apport.fileutils.get_all_reports() == [], \
        'no reports present for ignored binary'

    all_passed = True

finally:
    # clean up our ignore file
    if os.path.exists(ifpath):
	os.unlink(ifpath)
    if orig_ignore_file:
	os.rename(orig_ignore_file, ifpath)

    # clean up temporary work directory
    try:
        apport.fileutils.delete_report(report)
    except:
        pass
    shutil.rmtree(temp_report_dir)