models.py (1773B) download
1from email.policy import default
2from enum import Enum
3import json
4from flask_login import UserMixin
5
6from .app import db, login_manager
7
8""" function to load a user from database """
9@login_manager.user_loader
10def load_user(user_id):
11 return User.query.get(int(user_id))
12
13class UserType(Enum):
14 ADMIN = 0
15 CLIENT = 1
16
17class User(db.Model, UserMixin):
18 id = db.Column(db.Integer, primary_key=True)
19 type = db.Column(db.Enum(UserType), nullable=False, default=UserType.CLIENT)
20 email = db.Column(db.String(120), unique=True, nullable=False)
21 name = db.Column(db.String(20), unique=True, nullable=False)
22 password = db.Column(db.String(60), nullable=False)
23 image_file = db.Column(db.String(20), nullable=False, default='default.jpg')
24 phone = db.Column(db.Text, nullable=False)
25 address = db.Column(db.Text)
26
27 contact = db.Column(db.Integer, db.ForeignKey('user.id')) # set if user
28
29 def contact_class(self):
30 return User.query.filter_by(id=self.contact).first()
31
32
33class Trap(db.Model):
34 mac = db.Column(db.String(16), primary_key=True, nullable=False)
35 name = db.Column(db.Text)
36 last_heartbeat = db.Column(db.DateTime)
37 caught = db.Column(db.Boolean, nullable=False, default=False)
38 owner = db.Column(db.Integer, db.ForeignKey('user.id'))
39 connect_expired = db.Column(db.DateTime)
40 location_lat = db.Column(db.Float)
41 location_lon = db.Column(db.Float)
42
43 def pretty_mac(self):
44 upper = self.mac.upper()
45 return ':'.join([ upper[i] + upper[i+1] for i in range(0, len(upper), 2) ])
46
47 def owner_class(self):
48 return User.query.filter_by(id=self.owner).first()
49
50 def status_color(self):
51 if self.caught:
52 return '#f4a900'
53 return 'currentColor'