oneaccount/accounts.py

from flask import Flask, url_for
from flask_ldap3_login import LDAP3LoginManager
from flask_login import LoginManager, login_user, UserMixin, current_user, logout_user
from flask import render_template_string, redirect, render_template, request, send_file
from flask_ldap3_login.forms import LDAPLoginForm
from mfa import generateOTP, generateSecret, generateProvisioningUri
# from ad import updateMfaSecret
from db import setupMfaSecret, updateMfaSecret, getMfaSecret, removeMfa
import yamlcon
from wtforms import StringField
from wtforms.validators import DataRequired
import _forms
from ad import updatePassword
from vpn import genVPN, getVPN

app = Flask(__name__)
app.config['SECRET_KEY'] = 'secret'
app.config['DEBUG'] = True

# Setup LDAP Configuration Variables. Change these to your own settings.
# All configuration directives can be found in the documentation.

# Hostname of your LDAP Server
app.config['LDAP_HOST'] = 'corp.bbrunson.com'

app.config['LDAP_PORT'] = 389

# Base DN of your directory
app.config['LDAP_BASE_DN'] = 'dc=corp,dc=bbrunson,dc=com'

# Users DN to be prepended to the Base DN
app.config['LDAP_USER_DN'] = 'cn=users'

# Groups DN to be prepended to the Base DN
# app.config['LDAP_GROUP_DN'] = 'cn=groups'

# The RDN attribute for your user schema on LDAP
app.config['LDAP_USER_RDN_ATTR'] = 'cn'

# The Attribute you want users to authenticate to LDAP with.
app.config['LDAP_USER_LOGIN_ATTR'] = 'sAMAccountName'

adinfo = yamlcon.load("adinfo.yaml")

# The Username to bind to LDAP with
app.config['LDAP_BIND_USER_DN'] = adinfo.get('adbind_user', '')

# The Password to bind to LDAP with
app.config['LDAP_BIND_USER_PASSWORD'] = adinfo.get('adbind_pass', '')

login_manager = LoginManager(app)              # Setup a Flask-Login Manager
ldap_manager = LDAP3LoginManager(app)          # Setup a LDAP3 Login Manager.

# Create a dictionary to store the users in when they authenticate
# This example stores users in memory.
users = {}


# Declare an Object Model for the user, and make it comply with the
# flask-login UserMixin mixin.
class User(UserMixin):
    def __init__(self, dn, username, data):
        self.dn = dn
        self.username = username
        self.data = data

    def __repr__(self):
        return self.dn

    def get_id(self):
        return self.dn
    
class LDAPLoginForm0(LDAPLoginForm):
    mfacode = StringField('MFA')

# Declare a User Loader for Flask-Login.
# Simply returns the User if it exists in our 'database', otherwise
# returns None.
@login_manager.user_loader
def load_user(id):
    if id in users:
        return users[id]
    return None


# Declare The User Saver for Flask-Ldap3-Login
# This method is called whenever a LDAPLoginForm() successfully validates.
# Here you have to save the user, and return it so it can be used in the
# login controller.
@ldap_manager.save_user
def save_user(dn, username, data, memberships):
    user = User(dn, username, data)
    users[dn] = user
    return user


# Declare some routes for usage to show the authentication process.
@app.route('/')
def home():
    # Redirect users who are not logged in.
    if not current_user or current_user.is_anonymous:
        return redirect(url_for('login'))
    
    if getMfaSecret(user=current_user.data['sAMAccountName']):
        mfaStatus = "MFA is active"
    else:
        mfaStatus = "MFA is not active"

    # User is logged in, so show them a page with their cn and dn.
    # template = """
    # <h1>Welcome, {{ current_user.data.givenName }}</h1>

    # <h2> Email: {{ current_user.data.mail }}</h2>
    # <h2>{{ current_user.dn }}</h2> 
    # """
    # print(current_user)
    # print(current_user.data)
    # print(current_user.data['objectSid'])

    # return render_template_string(template)
    return render_template('home.html', current_user=current_user, mfaurl=url_for('two_factor'), \
                           mfastatus=mfaStatus, logouturl=url_for('logout'), changepwurl=url_for('changepw'), vpnurl=url_for('vpn'))

@app.route('/2fa')
def two_factor():
    if not current_user or current_user.is_anonymous:
        return redirect(url_for('login'))
    
    if getMfaSecret(user=current_user.data['sAMAccountName']):
        currSecret = "MFA already setup"
        code = generateOTP(getMfaSecret(user=current_user.data['sAMAccountName']))
        uri = ''

    elif not getMfaSecret(user=current_user.data['sAMAccountName']):
        setupMfaSecret(user=current_user.data['sAMAccountName'], secret=generateSecret())
        # currSecret = current_user.data['mfaSecret']
        currSecret = getMfaSecret(user=current_user.data['sAMAccountName'])
        uri = generateProvisioningUri(currSecret=currSecret, cu=current_user.data['sAMAccountName'])
        code = ''
    
    print(currSecret)
    return render_template('2fa.html', currSecret=currSecret, uri=uri, code=code, homeurl=url_for('home'), delurl=url_for('delTwoFactor'))

@app.route('/del2fa')
def delTwoFactor():
    if not current_user or current_user.is_anonymous:
        return redirect(url_for('login'))
    
    removeMfa(user=current_user.data['sAMAccountName'])

    return redirect(url_for('home'))

@app.route('/manual_login')
def manual_login():
    # Instead of using the form, you can alternatively authenticate
    # using the authenticate method.
    # This WILL NOT fire the save_user() callback defined above.
    # You are responsible for saving your users.
    app.ldap3_login_manager.authenticate('username', 'password')


@app.route('/login', methods=['GET', 'POST'])
def login():
    template = """
    {{ get_flashed_messages() }}
    {{ form.errors }}
    <form method="POST">
        <label>Username{{ form.username() }}</label>
        <label>Password{{ form.password() }}</label>
        <label>MFA Code{{ form.mfacode() }}</label>
        {{ form.submit() }}
        {{ form.hidden_tag() }}
    </form>
    """

    # Instantiate a LDAPLoginForm which has a validator to check if the user
    # exists in LDAP.
    form = LDAPLoginForm0()

    if form.validate_on_submit():
        # Successfully logged in, We can now access the saved user object
        # via form.user.
        login_user(form.user) # Tell flask-login to log them in.
        if getMfaSecret(user=current_user.data['sAMAccountName']):
            if form.mfacode.data == generateOTP(getMfaSecret(user=current_user.data['sAMAccountName'])):
                return redirect('/')  # Send them home
            else:
                logout_user()
                form.mfacode.errors.append("Invalid MFA Code")
                return redirect('/')
        elif not getMfaSecret(user=current_user.data['sAMAccountName']):
            login_user(form.user) # Tell flask-login to log them in.
            return redirect('/')
    return render_template_string(template, form=form)

@app.route('/changepw', methods=['GET', 'POST'])
def changepw():
    if not current_user or current_user.is_anonymous:
        return redirect(url_for('login'))
    
    form = _forms.ChangePasswordForm(request.form)
    if request.method == 'POST' and form.validate():
        new_password = form.newpw.data
        if updatePassword(cu=current_user.data['sAMAccountName'], newpw=new_password, adinfo=adinfo):
            print('Password changed successfully!')
            # flash('Password changed successfully!', 'success')
            return redirect(url_for('home'))
        else:
            print('Failed to change password. Please try again.')
            # flash('Failed to change password. Please try again.', 'danger')
    return render_template('changepw.html', form=form)

@app.route('/vpn', methods=['GET', 'POST'])
def vpn():
    if not current_user or current_user.is_anonymous:
        return redirect(url_for('login'))
    
    if request.args.get('dev'):
        return send_file(genVPN(cu=current_user.data['sAMAccountName'], dev=request.args.get('dev')), as_attachment=True)

    return render_template('vpn.html')

@app.route('/logout')
def logout():
    logout_user()
    return redirect('/')


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=81)