Geoff's Project

from flask import Flask, redirect, url_for, request
from flask_sqlalchemy import SQLAlchemy
from flask_login import LoginManager, current_user, login_user, logout_user  # UserMixin
from flask_admin import Admin, expose, AdminIndexView
from flask_admin.helpers import validate_form_on_submit
from flask_admin.contrib import sqla
from flask_security import Security, SQLAlchemyUserDatastore, RoleMixin, UserMixin
from flask_security.utils import hash_password
from wtforms import Form, StringField, PasswordField, ValidationError
from wtforms.validators import DataRequired
from werkzeug.security import generate_password_hash, check_password_hash
import os

from sqlalchemy import Column, Integer, String, Boolean, DateTime, ForeignKey

# Create Flask application
app = Flask(__name__)
FILE = 'sample_db.sqlite'
config = {
    'DATABASE_FILE': FILE,
    'SQLALCHEMY_DATABASE_URI': f'sqlite:///{FILE}',
    'SECRET_KEY': '123456790',
    'SQLALCHEMY_ECHO': True,
    'SECURITY_PASSWORD_SALT': 'ATGUOHAELKiubahiughaerGOJAEGj',
}
app.config.update(config)

# Instantiate extensions
db = SQLAlchemy(app)

# Define models
if hasattr(db, 'Table'):
    roles_users = db.Table(
        'roles_users',
        Column('user_id', Integer(), ForeignKey('user.id')),
        Column('role_id', Integer(), ForeignKey('role.id'))
    )


class Role(db.Model, RoleMixin):
    id = Column(Integer(), primary_key=True)
    name = Column(String(80), unique=True)
    description = Column(String(255))

    def __str__(self):
        return self.name


class User(db.Model, UserMixin):
    id = Column(Integer, primary_key=True)
    username = Column(String(20), unique=True, nullable=False)
    email = Column(String(120), unique=True, nullable=False)
    password = Column(String(60), nullable=False)
    image_file = Column(String(20), nullable=False, default='default.jpg')
    # posts = db.relationship('Post', backref='author', lazy=True)
    active = Column(Boolean())
    confirmed_at = Column(DateTime())

    if hasattr(db, 'relationship') and hasattr(db, 'backref'):
        roles = db.relationship('Role', secondary=roles_users, backref=db.backref('users', lazy='dynamic'))

    def __repr__(self):
        return f"User('{self.username}', '{self.email}'')"


# Setup Flask-Security
user_datastore = SQLAlchemyUserDatastore(db, User, Role)
security = Security(app, user_datastore)


# Define login and registration forms (for flask-login)
class LoginForm(Form):
    username = StringField(validators=[DataRequired()])
    password = PasswordField(validators=[DataRequired()])

    def validate_login(self):
        user = self.get_user()

        if user is None:
            raise ValidationError('Invalid user')

        if not check_password_hash(user.password, self.password.data):
            raise ValidationError('Invalid password')

    def get_user(self):
        return db.session.query(User).filter_by(username=self.username.data).first()


class RegistrationForm(Form):
    username = StringField(validators=[DataRequired()])
    email = StringField()
    password = PasswordField(validators=[DataRequired()])

    def validate_login(self):
        if db.session.query(User).filter_by(username=self.username.data).count() > 0:
            raise ValidationError('Duplicate username')


# Create customized model view class
class MyModelView(sqla.ModelView):

    def is_accessible(self):
        return current_user.is_authenticated


# Create customized index view class that handles login & registration
class MyAdminIndexView(AdminIndexView):

    @expose('/')
    def index(self):
        if not current_user.is_authenticated:
            return redirect(url_for('.login_view'))
        return super(MyAdminIndexView, self).index()

    @expose('/login/', methods=('GET', 'POST'))
    def login_view(self):
        if current_user.is_authenticated:
            return redirect(url_for('.index'))
        form = LoginForm(request.form)
        if validate_form_on_submit(form):
            user = form.get_user()
            if user is None:
                print('User not found!')
                return redirect(url_for('.index'))
            elif not check_password_hash(user.password, form.password.data):
                print('Invalid password!')
                return redirect(url_for('.index'))
            else:
                login_user(user)
                print('logged user in...')

        if current_user.is_authenticated:
            return redirect(url_for('.index'))
        link = '<p>Don\'t have an account? <a href="' + url_for('.register_view') + '">Click here to register.</a></p>'
        self._template_args['form'] = form
        self._template_args['link'] = link
        return super(MyAdminIndexView, self).index()

    @expose('/register/', methods=('GET', 'POST'))
    def register_view(self):
        form = RegistrationForm(request.form)
        if validate_form_on_submit(form):
            user = User()

            form.populate_obj(user)
            user.password = generate_password_hash(form.password.data)

            db.session.add(user)
            db.session.commit()

            login_user(user)
            return redirect(url_for('.index'))
        link = '<p>Already have an account? <a href="' + url_for('.login_view') + '">Click here to log in.</a></p>'
        self._template_args['form'] = form
        self._template_args['link'] = link
        return super(MyAdminIndexView, self).index()

    @expose('/logout/')
    def logout_view(self):
        logout_user()
        return redirect(url_for('.index'))


# Initialize flask-login
def init_login():
    login_manager = LoginManager()
    login_manager.init_app(app)

    # Create user loader function
    @login_manager.user_loader
    def load_user(user_id):
        return db.session.query(User).get(user_id)


# Flask views
@app.route('/')
def index():
    return '<a href="/admin/">Click me to get to Admin!</a>'


# Initialize flask-login
init_login()

# Create admin
admin = Admin(app, 'Example: Auth', index_view=MyAdminIndexView(), base_template='my_master.html')

# Add view
admin.add_view(MyModelView(User, db.session))
admin.add_view(MyModelView(Role, db.session))


def build_sample_db():

    db.drop_all()
    db.create_all()
    # test_user = User(username="test", email="[email protected]", password=generate_password_hash("test"))
    # db.session.add(test_user)

    with app.app_context():
        user_role = Role()
        super_user_role = Role()
        db.session.add(user_role)
        db.session.add(super_user_role)
        db.session.commit()

        user_datastore.create_user(
            username='admin',
            email='admin',
            password=hash_password('admin')  # ,
            # roles=[user_role, super_user_role]
        )
        db.session.commit()
    return


if __name__ == '__main__':
    # Build a sample db on the fly, if one does not exist yet.
    app_dir = os.path.realpath(os.path.dirname(__file__))
    database_path = os.path.join(app_dir, app.config['DATABASE_FILE'])
    if not os.path.exists(database_path):
        build_sample_db()
    app.run(debug=True)