1 files changed, 24 insertions, 3 deletions
diff --git a/app/database.py b/app/database.py
index 532f971..c66f603 100644
--- a/app/database.py
+++ b/app/database.py
@@ -5,15 +5,24 @@ from . import db
 from . import login
 from flask_login import UserMixin
 
+from flask import redirect, url_for, flash
+
 class User(db.Model, UserMixin):
     id = Column(String, primary_key=True)
     email = Column(String, unique=True, nullable=False)
     password = Column(String, nullable=False)
     pref_name = Column(String, nullable=False)
     last_login = Column(DateTime, nullable=False)
-    fs_uniquifier = Column(String, unique=True, nullable=False)
     active = Column(Boolean, nullable=False)
+    is_admin = Column(Boolean, nullable=False, default=False)
     roles = relationship('Role', secondary='roles_users',backref=backref('users', lazy='dynamic'))
+    peers = relationship('Peer', backref='owner', lazy=True)
+    networks = relationship('Network', backref="manager", lazy=True)
+
+    def __repr__(self):
+        return f"<User {email}>"
+    def __str__(self):
+        return self.pref_name
 
 class Role(db.Model):
     id = Column(String, primary_key=True)
@@ -28,15 +37,27 @@ class RolesUsers(db.Model):
 
 @login.user_loader
 def load_user(user_id):
-    return User.query.filer_by(id=user_id)
+    return User.query.filter_by(id=user_id).first()
 
+@login.unauthorized_handler
+def unauth():
+    flash("Please log in first")
+    return redirect(url_for("auth.login"))
 
 class Peer(db.Model):
     id = Column(String, primary_key=True)
     addr = Column(CIDR, nullable=False)
     public_key = Column(String, nullable=False)
+    description = Column(String, nullable=False)
+    owner_id = Column(String, ForeignKey('user.id'), nullable=False)
 
 class Network(db.Model):
     id = Column(String, primary_key=True)
-    subnet = Column(CIDR, nullable=False)
+    subnet = Column(CIDR, nullable=False, unique=True)
     description = Column(String, nullable=True)
+    manager_id = Column(String, ForeignKey('user.id'), nullable=False)
+
+    def __repr__(self):
+        return f"{self.description} ({self.subnet})"
+    def __str__(self):
+        return f"{self.description} ({self.subnet})"