From 0c8437edd8d5813fd357341d22f0de1c5d6dd66d Mon Sep 17 00:00:00 2001
From: Xevion <ryanchwalters@gmail.com>
Date: Wed, 3 Jul 2019 13:46:12 -0500
Subject: [PATCH] Updated roles, db migrate, @require_role v1

---
 app/models.py                     | 28 ++++++++++++++++++++++++++++
 app/routes.py                     | 22 ++++++++++++++++++----
 app/templates/hidden_history.html |  4 ++--
 config.py                         |  7 ++++++-
 4 files changed, 54 insertions(+), 7 deletions(-)

diff --git a/app/models.py b/app/models.py
index 8c7be6d..cc19a25 100644
--- a/app/models.py
+++ b/app/models.py
@@ -11,6 +11,7 @@ class User(UserMixin, db.Model):
     password_hash = db.Column(db.String(64))
     posts = db.relationship('Post', backref='author', lazy='dynamic')
     search_history = db.relationship('Search', backref='user', lazy='dynamic')
+    uroles = db.Column(db.String(80), default='')
 
     def set_password(self, password):
         self.password_hash = generate_password_hash(password)
@@ -19,6 +20,33 @@ class User(UserMixin, db.Model):
         if self.password_hash is None:
             raise "{} has no password_hash set!".format(self.__repr__())
         return check_password_hash(self.password_hash, password)
+    
+    def get_roles(self):
+        return self.uroles.split(' ')
+
+    def add_roles(self, roles):
+        user_roles = self.uroles.split(' ')
+        if type(roles) == str:
+            user_roles.append(roles)
+        elif type(roles) == list:
+            user_roles.extend(roles)
+        user_roles = ' '.join(user_roles)
+        self.uroles = user_roles
+
+    # Input: ['Insane', ['Fortunate', 'Blessed']]
+    # Meaning: Must have 'Insane' role, as well as 'Fortunate' or 'Blessed' roles.
+    def has_roles(self, roles):
+        user_roles = self.uroles.split(' ')
+        for reqrole in roles:
+            # If we have this role
+            if type(reqrole) == str:
+                if reqrole not in user_roles:
+                    return False
+            # If we have any of  these roles
+            elif type(reqrole) == list:
+                if not any([subreqrole in user_roles for subreqrole in reqrole]):
+                    return False
+        return True
 
     def __repr__(self):
         return '<User {}>'.format(self.username)
diff --git a/app/routes.py b/app/routes.py
index be60ece..a9c5f77 100644
--- a/app/routes.py
+++ b/app/routes.py
@@ -1,8 +1,8 @@
-from app import app, db
+from app import app, db, login
 from app.models import User, Search
 from app.forms import LoginForm, RegistrationForm
 from werkzeug.urls import url_parse
-from flask import render_template, redirect, url_for, flash, request, jsonify
+from flask import render_template, redirect, url_for, flash, request, jsonify, abort
 from flask_login import current_user, login_user, logout_user, login_required
 import requests
 import xmltodict
@@ -16,16 +16,29 @@ fake = faker.Faker()
 
 def strgen(length): return ''.join(random.choices(list(string.ascii_letters), k=length))
 
+def require_role(func, roles=[]):
+    def auth(*args, **kwargs):
+        if current_user.is_authenticated:
+            if current_user.has_roles(roles):
+                return func(*args, **kwargs)
+        return abort(401)
+    return auth
+
 @app.route('/api')
+@login_required
+@require_role
 def api():
     return 'fuckoff'
 
+@app.errorhandler(401)
+def unauthorized(e):
+    return redirect(url_for('login'))
+
 @app.route('/dashboard')
 def dashboard():
     return ''
 
 @app.route('/userinfo')
-@login_required
 def user_info():
     prepare = {
         'id' : current_user.get_id(),
@@ -35,7 +48,8 @@ def user_info():
         'is_active' : current_user.is_active,
         'is_anonymous' : current_user.is_anonymous,
         'is_authenticated' : current_user.is_authenticated,
-        'metadata' : current_user.metadata.info
+        'metadata' : current_user.metadata.info,
+        'uroles' : current_user.get_roles()
     }
     return jsonify(prepare)
 
diff --git a/app/templates/hidden_history.html b/app/templates/hidden_history.html
index e231621..cd3f727 100644
--- a/app/templates/hidden_history.html
+++ b/app/templates/hidden_history.html
@@ -20,7 +20,7 @@
                                     <th><abbr title="Search ID">ID</abbr></th>
                                     <th>Exact URL</th>
                                     <th>Query Arguments</th>
-                                    <th><abbr title="in UTC time">Timestamp</abbr></th>
+                                    <th>Timestamp</th>
                                 </tr>
                             </thead>
                             <tbody>
@@ -29,7 +29,7 @@
                                     <td>{{ post.id }}</td>
                                     <td>{{ post.exact_url }}</td>
                                     <td>{{ post.query_args }}</td>
-                                    <td>{{ post.timestamp }} </td>
+                                    <td>{{ post.timestamp }} UTC</td>
                                 </tr>
                                 {% endfor %}
                             </tbody>
diff --git a/config.py b/config.py
index 0631dd6..c677ea2 100644
--- a/config.py
+++ b/config.py
@@ -13,4 +13,9 @@ class Config(object):
     TEMPLATES_AUTO_RELOAD=True
     SQLALCHEMY_DATABASE_URI = os.environ.get('DATABASE_URL') or \
         'sqlite:///' + os.path.join(basedir, 'app.db')
-    SQLALCHEMY_TRACK_MODIFICATIONS = False
\ No newline at end of file
+    SQLALCHEMY_TRACK_MODIFICATIONS = False
+    # Flask-User settings
+    USER_APP_NAME = "Flask-User QuickStart App"      # Shown in and email templates and page footers
+    USER_ENABLE_EMAIL = False      # Disable email authentication
+    USER_ENABLE_USERNAME = True    # Enable username authentication
+    USER_REQUIRE_RETYPE_PASSWORD = True
\ No newline at end of file