#!/usr/bin/python

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 clean_preload_lib():
    print '* Cleaning preloadlib directory'
    run(['make', '-C', 'preloadlib', 'clean'], 'Cleaning preloadlib directory')

def build_preload_lib():
    print '* Building preload library with local parameters'
    run(['make', '-C', 'preloadlib', 'AGENTPATH=%s' % os.path.join(os.getcwd(),
        'bin', 'apport')], 'Building preloadlib')

def enable_preloadlib():
    assert os.access('preloadlib/libapport.so', os.X_OK)
    os.environ['LD_PRELOAD'] = os.path.join(os.getcwd(), 'preloadlib', 'libapport.so')

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
    return pid

def check_crash(expect_coredump=True, expect_corefile=False, sig=signal.SIGSEGV, check_running=True, sleep=0):
    global mode # hack to avoid passing the mode every time

    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 and mode == 'kernel':
        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') and len(err.splitlines()) == 1) , 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
#

if len(sys.argv) != 2 or sys.argv[1] not in ('lib', 'kernel'):
    print 'Usage: %s lib|kernel' % sys.argv[0]
    sys.exit(1)
mode = sys.argv[1]

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()
    # check if kernel crashdump helper is already active
    if mode == 'lib':
        assert core_pattern[0] != '|', \
                'kernel crash dump helper is still active; please disable before running this test.'

        # set up temporary report directory
        os.environ['APPORT_REPORT_DIR'] = temp_report_dir
        apport.fileutils.report_dir = temp_report_dir
        (fd, os.environ['APPORT_LOG_FILE']) = tempfile.mkstemp()
        os.close(fd)

        # setup preload library
        clean_preload_lib()
        build_preload_lib()

        apport_path = 'bin/apport'
    else:
        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 '* Check that 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'

    if mode == 'lib':
        print '* Check test process creation/killing without apport'
        check_crash(False)
        assert apport.fileutils.get_all_reports() == [], 'no report created without apport'

    print '* Check test process creation/killing with apport'
    if mode == 'lib':
        enable_preloadlib()
    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 '* Check that a subsequent crash does not alter unseen report'
    check_crash()
    st2 = os.stat(report)
    assert st == st2, 'original unread report did not change'

    print '* Check that 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 '* Check that 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 '* Check that dumped environment only has insensitive variables'
    check_safe_environment(pr['ProcEnviron'])

    print '* Check that collected system groups has nonempty user groups information'
    assert pr['UserGroups']
    print '* Check that 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 '* Check that 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 '* Check that 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 '* Check that 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 '* Check that 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 '* Check that 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 '* Check that 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 '* Test limitation of flooding: iteration',
    count = 0
    while count < 7:
        print count,
        sys.stdout.flush()
        check_crash()
        if mode == 'lib':
            time.sleep(1)
        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 '* Check that 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 '* Check that 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_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 '* Check that 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_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 '* Check that packaged executables create core dumps on proper ulimits'
    resource.setrlimit(resource.RLIMIT_CORE, (1, -1))
    check_crash(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 '* Check that crashes create core dumps with an existing crash report'
    check_crash(expect_corefile=True)
    apport.fileutils.delete_report(report)

    print '* Check that packaged executables create core dumps on proper ulimits for SIGABRT'
    resource.setrlimit(resource.RLIMIT_CORE, (1, -1))
    check_crash(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 '* Check that 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 '* Check that 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)
    if mode == 'lib':
        clean_preload_lib()
    if os.environ.has_key('APPORT_LOG_FILE'):
        if not all_passed:
            print '----- apport log file -----'
            print open(os.environ['APPORT_LOG_FILE']).read()
        os.unlink(os.environ['APPORT_LOG_FILE'])
