#!/usr/bin/python3
from passlib.apache import HtpasswdFile
from flask import Flask, request, make_response, jsonify, redirect
import jwt
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
            payload = {}
        else:
            exp = datetime.datetime.utcnow() + datetime.timedelta(minutes=exptime)
            payload = {'exp': exp}
        token = jwt.encode(payload, privkey, "EdDSA")
        resp.set_cookie('auth', token, expires=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

# 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 = privkey_file.read()
if __name__ == '__main__':
    app.run(debug=True)