1 files changed, 19 insertions, 2 deletions
diff --git a/app/database.py b/app/database.py
index 24e1930..532f971 100644
--- a/app/database.py
+++ b/app/database.py
@@ -1,7 +1,9 @@
 from sqlalchemy import Boolean, Column, DateTime, ForeignKey, Integer, String
+from sqlalchemy.dialects.postgresql import CIDR
 from sqlalchemy.orm import relationship, backref
 from . import db
-from flask_security.core import RoleMixin, UserMixin
+from . import login
+from flask_login import UserMixin
 
 class User(db.Model, UserMixin):
     id = Column(String, primary_key=True)
@@ -13,7 +15,7 @@ class User(db.Model, UserMixin):
     active = Column(Boolean, nullable=False)
     roles = relationship('Role', secondary='roles_users',backref=backref('users', lazy='dynamic'))
 
-class Role(db.Model, RoleMixin):
+class Role(db.Model):
     id = Column(String, primary_key=True)
     name = Column(String, unique=True)
     description = Column(String)
@@ -23,3 +25,18 @@ class RolesUsers(db.Model):
     id = Column(Integer(), primary_key=True)
     user_id = Column('user_id', String(), ForeignKey('user.id'))
     role_id = Column('role_id', String(), ForeignKey('role.id'))
+
+@login.user_loader
+def load_user(user_id):
+    return User.query.filer_by(id=user_id)
+
+
+class Peer(db.Model):
+    id = Column(String, primary_key=True)
+    addr = Column(CIDR, nullable=False)
+    public_key = Column(String, nullable=False)
+
+class Network(db.Model):
+    id = Column(String, primary_key=True)
+    subnet = Column(CIDR, nullable=False)
+    description = Column(String, nullable=True)