db.relationship #5

Author: eshaan7

src/FlaskRTBCTF/models.py

@@ -19,6 +19,7 @@ class User(db.Model, UserMixin):
     password = db.Column(db.String(60), nullable=False)
     confirmed_at = db.Column(db.DateTime(), default=datetime.utcnow)
     isAdmin = db.Column(db.Boolean, default=False)
+    score = db.relationship('Score', backref='user', lazy=True, uselist=False)
 
     def get_reset_token(self, expires_sec=1800):
         s = Serializer(current_app.config['SECRET_KEY'], expires_sec)
@@ -34,20 +35,21 @@ def verify_reset_token(token):
         return User.query.get(user_id)
 
     def __repr__(self):
-        return f"User('{self.username}', '{self.email}')"
+        return f"User('{self.username}', '{self.email}') | Score('{self.score}')"
 
 
 ''' Score Table '''
 
 class Score(db.Model):
-    userid = db.Column(db.Integer, primary_key=True)
+    user_id = db.Column(db.Integer, db.ForeignKey('user.id'),
+        nullable=False, primary_key=True)
     userHash = db.Column(db.Boolean, default=False)
     rootHash = db.Column(db.Boolean, default=False)
-    score = db.Column(db.Integer)
+    points = db.Column(db.Integer)
     timestamp = db.Column(db.DateTime(), default=datetime.utcnow)
 
     def __repr__(self):
-        return f"Score('{self.userid}', '{self.score}')"
+        return f"Score('{self.user_id}', '{self.points}')"
 
 
 ''' Notifications Table '''

src/FlaskRTBCTF/users/routes.py

@@ -22,7 +22,7 @@ def register():
             form.password.data).decode('utf-8')
         user = User(username=form.username.data,
                     email=form.email.data, password=hashed_password)
-        score = Score(userid=user.id, userHash=False, rootHash=False, score=0)
+        score = Score(user=user, userHash=False, rootHash=False, points=0)
         db.session.add(user)
         db.session.add(score)
         db.session.commit()

src/create_db.py

@@ -18,12 +18,25 @@
         confirmed_at=datetime.datetime.now(),
         isAdmin = True
     )
-    score = Score(userid=admin_user.id, userHash=False, rootHash=False, score=0)
+    admin_score = Score(user=admin_user, userHash=False, rootHash=False, points=0)
+    db.session.add(admin_user)
+    db.session.add(admin_score)
+
     notif = Notification(
         title=f"Welcome to {ctfname}",
         body = "The CTF is live now. Please read rules!"
     )
-    db.session.add(admin_user)
-    db.session.add(score)
     db.session.add(notif)
+
+    '''    
+    test = User(
+        username='test',
+        email='test@test.com',
+        password=bcrypt.generate_password_hash('test').decode('utf-8'),
+    )
+    testscore = Score(user=test, userHash=False, rootHash=False, points=0)
+    db.session.add(test)
+    db.session.add(testscore)
+    '''
+
     db.session.commit()