diff --git a/modules/frontend/__init__.py b/modules/frontend/__init__.py index 9306ac5..6ed0414 100644 --- a/modules/frontend/__init__.py +++ b/modules/frontend/__init__.py @@ -1,5 +1,7 @@ -from flask import Flask +from flask import Flask, render_template, flash import os +import sys +from functools import wraps from flask_sqlalchemy import SQLAlchemy import warnings warnings.filterwarnings("ignore") @@ -9,4 +11,24 @@ app.config['SECRET_KEY'] = 'secret-key-goes-here' app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///db.sqlite' -db = SQLAlchemy(app) \ No newline at end of file +db = SQLAlchemy(app) + +isAdmin = False +if(len(sys.argv) > 1 and sys.argv[1] == '--admin'): + isAdmin = True +else: + isAdmin = False + +def checkAdmin(func): + @wraps(func) + def decorated_function(*args, **kwargs): + isAdmin = False + if(len(sys.argv) > 1 and sys.argv[1] == '--admin'): + isAdmin = True + else: + isAdmin = False + if(isAdmin == False): + flash("Sorry, You do not have permission to access this page\nPlease contact admin") + return render_template('login.html') + return func(*args, **kwargs) + return decorated_function \ No newline at end of file diff --git a/modules/frontend/server.py b/modules/frontend/server.py index 616fa92..134b731 100644 --- a/modules/frontend/server.py +++ b/modules/frontend/server.py @@ -1,5 +1,6 @@ from flask import Flask, flash, request, redirect, url_for, render_template, send_file import os +import sys from flask_sqlalchemy import SQLAlchemy from flask_login import UserMixin from flask_login import LoginManager, login_user, login_required, current_user, logout_user @@ -7,7 +8,7 @@ warnings.filterwarnings("ignore") from werkzeug.security import generate_password_hash, check_password_hash -from __init__ import app, db +from __init__ import app, db, isAdmin, checkAdmin from models import User PEOPLE_FOLDER = os.path.join('static','styles') @@ -23,7 +24,7 @@ def load_user(user_id): @app.route("/", methods=['GET']) def index(): - return render_template('home.html') + return render_template('home.html', isAdmin = isAdmin) @app.route('/login', methods=['GET','POST']) def login(): @@ -43,9 +44,10 @@ def login(): # if the above check passes, then we know the user has the right credentials login_user(user, remember=remember) return render_template('home.html') - return render_template('login.html') + return render_template('login.html', isAdmin = isAdmin) @app.route('/signup', methods=['GET','POST']) +@checkAdmin def signup(): if request.method =='POST': email = request.form.get('email') @@ -64,15 +66,15 @@ def signup(): # add the new user to the database db.session.add(new_user) db.session.commit() - return render_template('login.html') + return render_template('login.html', isAdmin = isAdmin) - return render_template('signup.html') + return render_template('signup.html', isAdmin = isAdmin) @app.route('/logout') @login_required def logout(): logout_user() - return render_template('home.html') + return render_template('home.html', isAdmin = isAdmin) # @app.route("/png-extraction", methods = ['GET']) # @login_required diff --git a/modules/frontend/templates/base.html b/modules/frontend/templates/base.html index 0bb0fed..437b3c6 100644 --- a/modules/frontend/templates/base.html +++ b/modules/frontend/templates/base.html @@ -127,10 +127,12 @@