#!/usr/bin/python3
#
#    test-squid.py quality assurance test script
#    Copyright (C) 2008-2017 Canonical Ltd.
#    Author: Jamie Strandboge <jamie@canonical.com>
#    Author: Marc Deslauriers <marc.deslauriers@canonical.com>
#    Author: Christian Ehrhardt <christian.ehrhardt@canonical.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.
#
#    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, see <http://www.gnu.org/licenses/>.
#

import os
import os.path
import re
import shutil
import signal
import subprocess
import sys
import time
import unittest


class HttpdCommon(unittest.TestCase):
    '''Common functions'''
    def __init__(self, *args):
        '''This is called for each TestCase test instance, which isn't much better
           than SetUp.'''
        unittest.TestCase.__init__(self, *args)

    def _setUp(self):
        '''Setup'''
        if not hasattr(self, 'initscript'):
            self._set_initscript("/etc/init.d/apache2")

        self._stop()
        self._start()

    def _set_initscript(self, initscript):
        self.initscript = initscript

    def _tearDown(self):
        '''Clean up after each test_* function'''
        self._stop()

    def _start(self):
        '''Start httpd'''
        expected = 0
        ret, report = cmd([self.initscript, 'start'])
        result = 'Got exit code %d, expected %d\n' % (ret, expected)
        self.assertEqual(expected, ret, result + report)
        time.sleep(2)

    def _stop(self):
        '''Stop httpd'''
        expected = 0
        ret, report = cmd([self.initscript, 'stop'])
        result = 'Got exit code %d, expected %d\n' % (ret, expected)
        self.assertEqual(expected, ret, result + report)

    def _word_find(self, report, content, invert=False):
        '''Check for a specific string'''
        if invert:
            warning = 'Found "%s"\n' % content
            self.assertTrue(content not in report, warning + report)
        else:
            warning = 'Could not find "%s"\n' % content
            self.assertTrue(content in report, warning + report)

    def _test_url_proxy(self, url="http://localhost/", content="",
                        proxy="localhost:3128"):
        '''Test the given url'''
        ret, rep = cmd(['elinks', '-verbose', '2', '-no-home', '1',
                        '-eval',
                        'set protocol.ftp.proxy.host = "%s"' % (proxy),
                        '-eval',
                        'set protocol.http.proxy.host = "%s"' % (proxy),
                        '-eval',
                        'set protocol.https.proxy.host = "%s"' % (proxy),
                        '-dump', url])
        expected = 0
        result = 'Got exit code %d, expected %d\n' % (ret, expected)
        self.assertEqual(expected, ret, result + rep)

        if content != "":
            self._word_find(rep, content)


class BasicTest(HttpdCommon):
    '''Test basic functionality'''
    def setUp(self):
        '''Setup mechanisms'''

        self._set_initscript("/etc/init.d/squid")
        HttpdCommon._setUp(self)

        self.aa_profile = "usr.sbin.squid"
        self.aa_abs_profile = "/etc/apparmor.d/%s" % self.aa_profile

    def tearDown(self):
        '''Shutdown methods'''
        HttpdCommon._tearDown(self)

    def test_daemons(self):
        '''Test daemon'''

        pidfile = "/run/squid.pid"
        exe = "squid"

        self.assertTrue(check_pidfile(exe, pidfile))

    def test_http_proxy(self):
        '''Test http'''
        self._test_url_proxy("http://localhost/",
                             "It works!",
                             "http://localhost:3128/")

    def test_https_proxy(self):
        '''Test https'''
        self._test_url_proxy("https://localhost/",
                             "It works!",
                             "http://localhost:3128/")

    def test_ftp_proxy(self):
        '''Test ftp'''
        self._test_url_proxy("ftp://anonymous@localhost:21", "irectory",
                             "http://localhost:3128/")


    # Run this last so if we enable the profile then we don't unload it
    def test_zz_apparmor(self):
        '''Test apparmor'''

        # Currently while we have a profile, it is shipped disabled by default.
        # Verify that.
        ret, report = check_apparmor(self.aa_abs_profile, is_running=False)
        expected = 1
        result = 'Got exit code %d, expected %d\n' % (ret, expected)
        self.assertEqual(ret, expected, result + report)

        # Verify it is syntactically correct
        ret, report = cmd(['apparmor_parser', '-p', self.aa_abs_profile])
        expected = 0
        result = 'Got exit code %d, expected %d\n' % (ret, expected)
        self.assertEqual(ret, expected, result + report)

        # The remaining tests try to actually load a profile
        # skip them if securityfs isn't mounted (i.e., we are in a lxc container)
        ret, _ = _aa_status()
        # from the manpage:
        # 3   if the apparmor control files aren't available under /sys/kernel/security/.
        if ret == 3:
            return True

        # Verify it loads ok
        ret, report = cmd(['aa-enforce', self.aa_abs_profile])
        expected = 0
        result = 'Got exit code %d, expected %d\n' % (ret, expected)
        self.assertEqual(ret, expected, result + report)

        self._stop()
        self._start()

        ret, report = check_apparmor(self.aa_abs_profile, is_running=True)
        expected = 1
        result = 'Got exit code %d, expected %d\n' % (ret, expected)
        self.assertEqual(ret, expected, result + report)


# http://www.chiark.greenend.org.uk/ucgi/~cjwatson/blosxom/2009-07-02-python-sigpipe.html
# This is needed so that the subprocesses that produce endless output
# actually quit when the reader goes away.
def subprocess_setup():
    # Python installs a SIGPIPE handler by default. This is usually not what
    # non-Python subprocesses expect.
    signal.signal(signal.SIGPIPE, signal.SIG_DFL)


def _restore_backup(path):
    pathbackup = path + '.autotest'
    if os.path.exists(pathbackup):
        shutil.move(pathbackup, path)


def _save_backup(path):
    pathbackup = path + '.autotest'
    if os.path.exists(path) and not os.path.exists(pathbackup):
        shutil.copy2(path, pathbackup)
        # copy2 does not copy ownership, so do it here.
        # Reference: http://docs.python.org/library/shutil.html
        stat = os.stat(path)
        os.chown(pathbackup, stat[4], stat[5])


def config_replace(path, contents, append=False):
    '''Replace (or append) to a config file'''
    _restore_backup(path)
    if os.path.exists(path):
        _save_backup(path)
        if append:
            with open(path) as fileh:
                contents = fileh.read() + contents
    with open(path, 'w') as fileh:
        fileh.write(contents)


def config_restore(path):
    '''Rename a replaced config file back to its initial state'''
    _restore_backup(path)


def recursive_rm(dir_path, contents_only=False):
    '''recursively remove directory'''
    names = os.listdir(dir_path)
    for name in names:
        path = os.path.join(dir_path, name)
        if os.path.islink(path) or not os.path.isdir(path):
            os.unlink(path)
        else:
            recursive_rm(path)
    if not contents_only:
        os.rmdir(dir_path)


def check_pidfile(exe, pidfile):
    '''Checks if pid in pidfile is running'''
    if not os.path.exists(pidfile):
        return False

    # get the pid
    try:
        with open(pidfile, 'r') as filed:
            pid = filed.readline().rstrip('\n')
    except:
        return False

    return check_pid(exe, pid)


def check_pid(exe, pid):
    '''Checks if pid is running'''
    cmdline = "/proc/%s/cmdline" % (str(pid))
    if not os.path.exists(cmdline):
        return False

    # get the command line
    try:
        with open(cmdline, 'r') as filed:
            tmp = filed.readline().split('\0')
    except:
        return False

    # this allows us to match absolute paths or just the executable name
    if re.match(r'^' + exe + r'$', tmp[0]) or \
       re.match(r'.*/' + exe + r'$', tmp[0]) or \
       re.match(r'^' + exe + r': ', tmp[0]) or \
       re.match(r'^\(' + exe + r'\)', tmp[0]):
        return True

    return False


def cmd(command, stderr=subprocess.STDOUT, stdout=subprocess.PIPE,
        stdin=None):
    '''Try to execute given command (array) and return its stdout, or return
    a textual error if it failed.'''

    try:
        subp = subprocess.Popen(command, stdin=stdin, stdout=stdout,
                                stderr=stderr, close_fds=True,
                                preexec_fn=subprocess_setup,
                                universal_newlines=True)
    except OSError as err:
        return [127, str(err)]

    out, outerr = subp.communicate(None)
    # Handle redirection of stdout
    if out is None:
        out = ''
    # Handle redirection of stderr
    if outerr is None:
        outerr = ''
    return [subp.returncode, out + outerr]


def _aa_status():
    '''Get aa-status output'''
    exe = "/usr/sbin/aa-status"
    assert os.path.exists(exe)
    if os.geteuid() == 0:
        return cmd([exe])
    return cmd(['sudo', exe])


def is_apparmor_loaded(path):
    '''Check if profile is loaded'''
    ret, report = _aa_status()
    if ret != 0:
        return False

    for line in report.splitlines():
        if line.endswith(path):
            return True
    return False


def is_apparmor_confined(path):
    '''Check if application is confined'''
    ret, report = _aa_status()
    if ret != 0:
        return False

    for line in report.splitlines():
        if re.search(r'%s \(' % path, line):
            return True
    return False


def check_apparmor(path, is_running=True):
    '''Check if path is loaded and confined'''
    ret = -1

    if not os.path.exists('/sbin/apparmor_parser'):
        return (ret, "Skipped (couldn't find apparmor_parser)")

    ret = 0
    msg = ""
    if not is_apparmor_loaded(path):
        ret = 1
        msg = "Profile not loaded for '%s'" % path

    # this check only makes sense it the 'path' is currently executing
    if is_running and ret == 0 and not is_apparmor_confined(path):
        ret = 1
        msg = "'%s' is not running in enforce mode" % path

    return (ret, msg)


if __name__ == '__main__':
    SUITE = unittest.TestSuite()
    SUITE.addTest(unittest.TestLoader().loadTestsFromTestCase(BasicTest))

    RC = unittest.TextTestRunner(verbosity=2).run(SUITE)
    if not RC.wasSuccessful():
        sys.exit(1)