#!/usr/bin/python3
from passlib.apache import HtpasswdFile
from flask import Flask, request, make_response, jsonify, redirect
import python_jwt as jwt
from jwcrypto.jwk import JWK
import datetime
from json import dumps
import argparse

app = Flask(__name__)

def authorized_request():
    args = request.form

    user = args['user'] if 'user' in args else None
    pswd = args['pass'] if 'pass' in args else None
    remember = args['remember'] if 'remember' in args else None

    valid = user is not None and pswd is not None
    valid = valid and user in htpasswd.users()
    valid = valid and htpasswd.check_password(user, pswd)

    return valid, user, pswd, remember

@app.route('/auth', methods=['POST'])
def authorize():
    # on success, redirect to /. on failure, redirect to /login
    auth, user, _, remember  = authorized_request()
    if auth:
        resp = redirect('/')
        if remember:
            exp = None
        else:
            exp = datetime.timedelta(minutes=exptime)
        token = jwt.generate_jwt({}, privkey, "EdDSA", exp)
        resp.set_cookie('auth', token, max_age=exp)
        return resp
        # this stuff too
    else:
        resp = redirect('/login?err')
        return resp

@app.route('/logout', methods=['GET'])
def logout():
    resp = redirect('/login')
    resp.delete_cookie('auth')
    return resp

if  __name__ == '__main__':
    # argparse arguments
    parser = argparse.ArgumentParser(
            prog='login.py',
            description='A web server that handles htpasswd-file JWT auth logic')
    parser.add_argument('htpasswd')
    parser.add_argument('privkey')
    parser.add_argument('-e', '--expireminutes', default=30, type=int)

    args = parser.parse_args()
    htpasswd_filename = args.htpasswd
    privkey_filename = args.privkey
    exptime = args.expireminutes

    htpasswd = HtpasswdFile(htpasswd_filename)
    with open(privkey_filename, 'rb') as privkey_file:
        privkey = JWK()
        privkey.import_from_pem(privkey_file.read())
    app.run(debug=True)