#!/usr/bin/env python3
import sys
import requests
import argparse
from termcolor import colored
from bs4 import BeautifulSoup

parser = argparse.ArgumentParser(epilog='The password can either be supplied on the command line or in the file "/etc/bps/testuser.passwd".')
parser.add_argument('-u', '--username', required=True, help='A valid Uvt username')
parser.add_argument('-p', '--password', required=False, help='A valid password')
parser.add_argument('-H', '--hostname', required=False, default='bps.uvt.nl', help='The hostname of the BPS server')
parser.add_argument('-v', '--verbose', action='store_true', help='Prints more verbose output on errors')
parser.add_argument('-q', '--quiet', action='store_true', help='Suppresses all output')
args = parser.parse_args()

if not args.password:
    try:
        with open('/etc/bps/testuser.passwd') as f:
            args.password = f.read().rstrip()
    except IOError:
        parser.error('Password is missing. Use --help for more information about how to use this command.')

BPS_VERSION = '1.0.1'
STATE_OK = 0
STATE_WARNING = 1
STATE_CRITICAL = 2
STATE_UNKNOWN = 3
EXIT_STATUS = STATE_OK
HTTP_URL = 'http://{}/'.format(args.hostname)
HTTPS_URL = 'https://{}/'.format(args.hostname)
SSO_URL = 'https://sso.uvt.nl/login?service=https%3A%2F%2F' + args.hostname + '%2Flogin%2Fsso%2F%3Fnext%3D%252F'
SSO_POST_URL = 'https://sso.uvt.nl/login'

s = requests.Session()

def check_bps():
    check_version()
    check_login_redirect()
    check_sso_redirect()
    check_sso_login()

def print_docstring(func):
    'A simple decorator that prints a function’s docstring'
    def wrapper(*a, **k):
        if not args.quiet:
            print('', func.__doc__, '=' * len(func.__doc__), sep='\n')
        return func(*a, **k)
    return wrapper

@print_docstring
def check_version():
    'Check the version number of the BPS installation to be tested'

    r = requests.get(HTTPS_URL)
    soup = BeautifulSoup(r.text, 'lxml')
    tag = soup.find(id='version')
    version = tag.string if tag else None

    check('The installation’s version number ({}) equals the target version number ({})'.format(version, BPS_VERSION), version == BPS_VERSION)

@print_docstring
def check_login_redirect():
    'Check if anonymous users are redirected to the login page'

    r = requests.get(HTTP_URL)

    check('Requests over plain HTTP are redirected',
          r.history and r.history[0].is_redirect)

    check('Plain HTTP requests are redirected to ' + HTTPS_URL,
          r.history and r.history[1].url == HTTPS_URL)

    check('Cookieless requests are redirected',
          r.history and r.history[1].is_redirect)

    check('The login page responds with 200 OK',
          r.status_code == requests.codes.OK)

@print_docstring
def check_sso_redirect():
    'Check if the SSO login page redirects to the SSO server'

    r = requests.get(HTTPS_URL + 'login/sso/')

    check('The SSO login page response with a redirect',
          r.history and r.history[0].is_redirect)

    check('Requests to the SSO login page are redirected to the SSO server',
          r.url == SSO_URL)

    check('The SSO server responds with 200 OK',
          r.status_code == requests.codes.OK)

@print_docstring
def check_sso_login():
    'Check the login procedure via the SSO server'

    r = requests.get(SSO_URL)
    sso_cookie = r.cookies['session']

    check('The SSO server provides a session cookie',
          sso_cookie)

    soup = BeautifulSoup(r.content, 'lxml')
    sso_login_form = {
        'username': args.username,
        'password': args.password,
    }
    for tag in soup.find_all("input", type="hidden"):
        sso_login_form[tag['name']] = tag['value']
    cookies = r.cookies
    r = s.post(SSO_POST_URL, data=sso_login_form, cookies=cookies)

    check('After logging in, the SSO server returns a redirect',
          r.history and r.history[0].is_redirect)

    check('After logging in, the SSO server redirects to {}'.format(HTTPS_URL),
          r.history and r.history[1].url.startswith(HTTPS_URL))

    check('The login view returns (yet another!) redirect',
          r.history and r.history[1].is_redirect)

    check('The login view redirects to the homepage',
          r.url == HTTPS_URL)

    check('The login view has set the sessionid token',
          'sessionid' in s.cookies)

    r = s.get(HTTPS_URL)

    check('Subsequent requests have their session persisted (i.e. do not require logging in)',
          not r.history and r.url == HTTPS_URL)

def check(message, succeeded):
    if succeeded:
        show('  OK  ', 'green', message)
    else:
        show('FAILED', 'red', message)
        global EXIT_STATUS
        EXIT_STATUS = STATE_CRITICAL

def show(status, color, message):
    if sys.stdout.isatty():
        status = colored(status, color, attrs=['bold'])
    if not args.quiet:
        print('[{}] {}'.format(status, message))

if __name__ == '__main__':
    try:
        check_bps()
        exit(EXIT_STATUS)
    except Exception as e:
        if args.verbose:
            import traceback
            traceback.print_exc()
        show('ERROR:', 'red', e)
        exit(STATE_UNKNOWN)
