import share link token generation to models.py

app/corpora/routes.py

@@ -69,19 +69,9 @@ def disable_corpus_is_public(corpus_id):
 @login_required
 def follow_corpus(corpus_id, token):
     corpus = Corpus.query.get_or_404(corpus_id)
-    try:
-        payload = jwt.decode(
-                token,
-                current_app.config['SECRET_KEY'],
-                algorithms=['HS256'],
-                issuer=current_app.config['SERVER_NAME'],
-                # options={'require': ['exp', 'iat', 'iss', 'sub']}
-                options={'require': ['exp', 'iat', 'iss']}
-            )
-    except jwt.PyJWTError:
-        abort(410)
-    # permission = payload.get('sub')
-    if not current_user.is_following_corpus(corpus):
+    if not (current_user.is_authenticated and current_user.verify_follow_corpus_token(token)):
+        abort(403)
+    if not current_user.is_following_corpus(corpus) and current_user != corpus.user:
         current_user.follow_corpus(corpus)
         db.session.commit()
         flash(f'You are following {corpus.title} now', category='corpus')
@@ -174,9 +164,6 @@ def create_corpus():
 def corpus(corpus_id):
     corpus = Corpus.query.get_or_404(corpus_id)
     exp_date = (datetime.utcnow() + timedelta(days=7)).strftime('%b %d, %Y')
-    print(corpus.user)
-    print(current_user)
-    print(current_user.is_following_corpus(corpus))
     if corpus.user == current_user or current_user.is_administrator():
         return render_template(
             'corpora/corpus.html.j2',
@@ -201,18 +188,7 @@ def generate_corpus_share_link(corpus_id):
     # permission = data['permission']
     exp_data = data['expiration']
     expiration = datetime.strptime(exp_data, '%b %d, %Y')
-    now = datetime.utcnow()
-    payload = {
-            'exp': expiration,
-            'iat': now,
-            'iss': current_app.config['SERVER_NAME']
-            # 'sub': permission
-        }
-    token = jwt.encode(
-            payload,
-            current_app.config['SECRET_KEY'],
-            algorithm='HS256'
-        )
+    token = current_user.generate_follow_corpus_token(corpus_id, expiration)
     link = url_for('corpora.follow_corpus', corpus_id=corpus_id, token=token, _external=True)
     return link
 

app/models.py

@@ -1,6 +1,6 @@
 from datetime import datetime, timedelta
 from enum import Enum, IntEnum
-from flask import current_app, url_for
+from flask import abort, current_app, url_for
 from flask_hashids import HashidMixin
 from flask_login import UserMixin
 from sqlalchemy.ext.associationproxy import association_proxy
@@ -768,6 +768,37 @@ class User(HashidMixin, UserMixin, db.Model):
     def is_following_corpus(self, corpus):
         return corpus in self.followed_corpora
     
+    def generate_follow_corpus_token(self, corpus_id, expiration=7):
+        now = datetime.utcnow()
+        payload = {
+            'exp': expiration,
+            'iat': now,
+            'iss': current_app.config['SERVER_NAME'],
+            'sub': corpus_id
+        }
+        return jwt.encode(
+            payload,
+            current_app.config['SECRET_KEY'],
+            algorithm='HS256'
+        )
+    
+    def verify_follow_corpus_token(self, token):
+        try:
+            payload = jwt.decode(
+                token,
+                current_app.config['SECRET_KEY'],
+                algorithms=['HS256'],
+                issuer=current_app.config['SERVER_NAME'],
+                options={'require': ['exp', 'iat', 'iss', 'sub']}
+            )
+        except jwt.PyJWTError:
+            return False
+        corpus_id = payload.get('sub')
+        corpus = Corpus.query.get_or_404(corpus_id)
+        if corpus is None:
+            return False
+        return True
+
     def to_json_serializeable(self, backrefs=False, relationships=False, filter_by_privacy_settings=False):
         json_serializeable = {
             'id': self.hashid,