app/models.py

import base64
from datetime import datetime, timedelta
import os

from werkzeug.security import generate_password_hash, check_password_hash
from flask_login import UserMixin

from app import db, login

@login.user_loader
def load_user(id):
  return User.query.get(int(id))

class PaginatedAPIMixin(object):
  @staticmethod
  def to_collection_dict(query, page, per_page, endpoint, **kwargs):
    resources = query.paginate(page, per_page, False)
    data = {
      'items': [item.to_dict() for item in resources.items],
      '_meta': {
      'page': page,
      'per_page': per_page,
      'total_pages': resources.pages,
      'total_items': resources.total
      },
    }
    return data

class UserRole(db.Model):
  id = db.Column(db.Integer, primary_key=True)
  role_name = db.Column(db.String(32), index=True, unique=True, default='regular')
  role_description = db.Column(db.String(128), default='Regular user with access to only base stacks')

  users = db.relationship('User', backref='user_role', lazy='dynamic')

  def __repr__(self):
    return ('\tID: {}\n'
            '\tRole name: {}\n'
            '\tRole description: {}\n'
            .format(self.id, self.role_name, self.role_description))

class UserOrganization(db.Model):
  id = db.Column(db.Integer, primary_key=True)
  org_name = db.Column(db.String(32), index=True, unique=True, default='nordix')
  org_description = db.Column(db.String(128), default='Nordix Community users')

  users = db.relationship('User', backref='user_organization', lazy='dynamic')

  def __repr__(self):
    return ('\tID: {}\n'
            '\tOrganization name: {}\n'
            '\tOrganization description: {}\n'
            .format(self.id, self.org_name, self.org_description))

class User(PaginatedAPIMixin, UserMixin, db.Model):
  id = db.Column(db.Integer, primary_key=True)
  role_id = db.Column(db.Integer, db.ForeignKey('user_role.id'))
  org_id = db.Column(db.Integer, db.ForeignKey('user_organization.id'))
  username = db.Column(db.String(64), index=True, unique=True)
  email = db.Column(db.String(128), index=True, unique=True)
  fullname = db.Column(db.String(128))
  password_hash = db.Column(db.String(128))
  ssh_public_key = db.Column(db.String(1024))
  registered_on = db.Column(db.DateTime, index=True, default=datetime.utcnow)
  confirmed = db.Column(db.Integer, default=0)
  last_logged_in = db.Column(db.DateTime, index=True)
  confirmed_on = db.Column(db.DateTime, index=True)
  updated_on = db.Column(db.DateTime, index=True, default=datetime.utcnow)
  token = db.Column(db.String(32), index=True, unique=True)
  token_expiration = db.Column(db.DateTime)

  bookings = db.relationship('Booking', backref='user', lazy='dynamic')

  def __repr__(self):
    return ('\tID: {}\n'
            '\tUsername: {}\n'
            '\tFullname: {}\n'
            '\tEmail: {}\n'
            '\tPassword Hash: {}\n'
            '\tSSH Public Key: {}\n'
            '\tRegistered On: {}\n'
            '\tConfirmed: {}\n'
            '\tConfirmed On: {}\n'
            '\tUpdated On: {}\n'
            '\tRole ID: {}\n'
            .format(self.id, self.username, self.fullname, self.email,
              self.password_hash, self.ssh_public_key, self.registered_on,
              self.confirmed, self.confirmed_on, self.updated_on, self.role_id))

  def set_password(self, password):
    self.password_hash = generate_password_hash(password)

  def check_password(self, password):
    return check_password_hash(self.password_hash, password)

  def to_dict(self):
    data = {
      'id': self.id,
      'role_id': self.role_id,
      'role_name': UserRole.query.filter_by(id=self.role_id).first().role_name,
      'org_id': self.org_id,
      'org_name': UserOrganization.query.filter_by(id=self.org_id).first().org_name,
      'username': self.username,
      'email': self.email,
      'ssh_public_key': self.ssh_public_key,
      'registered_on': self.registered_on.isoformat(),
      'confirmed': self.confirmed
    }
    return data

  def from_dict(self, data, new_user=False):
    for field in ['confirmed', 'org_id', 'role_id']:
      if field in data:
        setattr(self, field, data[field])
      if new_user and 'password' in data:
        self.set_password(data['password'])

  def get_token(self, expires_in=3600):
    now = datetime.utcnow()
    if self.token and self.token_expiration > now + timedelta(seconds=60):
      return self.token
    self.token = base64.b64encode(os.urandom(24)).decode('utf-8')
    self.token_expiration = now + timedelta(seconds=expires_in)
    db.session.add(self)
    return self.token

  def revoke_token(self):
    self.token_expiration = datetime.utcnow() - timedelta(seconds=1)

  @staticmethod
  def check_token(token):
    user = User.query.filter_by(token=token).first()
    if user is None or user.token_expiration < datetime.utcnow():
      return None
    return user

  @staticmethod
  def check_none_regular_user(token):
    user = User.query.filter_by(token=token).first()
    current_user_role_id = user.role_id
    regular_user_role_id = UserRole.query.filter_by(role_name='regular').first().id
    if current_user_role_id == regular_user_role_id:
      return None
    return user

class Booking(PaginatedAPIMixin, db.Model):
  id = db.Column(db.Integer, primary_key=True)
  user_id = db.Column(db.Integer, db.ForeignKey('user.id'))
  status_id = db.Column(db.Integer, db.ForeignKey('booking_status.id'))
  booking_length = db.Column(db.Integer)
  created_on = db.Column(db.DateTime, index=True, default=datetime.utcnow)
  updated_on = db.Column(db.DateTime, index=True, default=datetime.utcnow)
  expires_on = db.Column(db.DateTime, index=True, default=datetime.utcnow)
  heat_stack_name = db.Column(db.String(64))
  floating_ip = db.Column(db.String(64))
  engine_version = db.Column(db.String(64))
  stack = db.Column(db.String(64))
  stack_version = db.Column(db.String(64))
  scenario = db.Column(db.String(64))
  scenario_version = db.Column(db.String(64))
  scenario_deploy_log_url = db.Column(db.String(256))

  def __repr__(self):
    return ('\tID: {}\n'
            '\tUser ID: {}\n'
            '\tStatus: {}\n'
            '\tCreated: {}\n'
            '\tUpdated: {}\n'
            '\tExpires: {}\n'
            '\tHeat Stack Name: {}\n'
            '\tFloating IP: {}\n'
            '\tStack Name: {}\n'
            '\tStack Version: {}\n'
            '\tScenario: {}\n'
            '\tScenario Version: {}\n'
            .format(self.id, self.user_id, self.status, self.created,
              self.updated, self.expires, self.heat_stack_name, self.floating_ip,
              self.engine_version, self.stack, self.stack_version, self.scenario, self.scenario_version))

  def to_dict(self):
    data = {
      'id': self.id,
      'user_id': self.user_id,
      'booking_length': self.booking_length,
      'heat_stack_name': self.heat_stack_name,
      'scenario_deploy_log_url': self.scenario_deploy_log_url,
      'engine_version': self.engine_version,
      'scenario_version': self.scenario_version,
      'floating_ip': self.floating_ip,
      'created_on': self.created_on.isoformat(),
      'stack': self.stack,
      'status_text': BookingStatus.query.filter_by(id=self.status_id).first().status_text,
      'status_id': self.status_id,
      'scenario': self.scenario
    }
    return data

  def from_dict(self, data, new_booking=False):
    for field in ['status_text', 'floating_ip', 'heat_stack_name', \
                  'scenario_deploy_log_url', 'engine_version', \
                  'scenario_version']:
      if field in data:
        setattr(self, field, data[field])

class BookingStatus(db.Model):
  id = db.Column(db.Integer, primary_key=True)
  status_text = db.Column(db.String(64), default='new')
  status_description = db.Column(db.String(128), default='New booking request')

  bookings = db.relationship('Booking', backref='booking_status', lazy='dynamic')

  def __repr__(self):
    return ('\tID: {}\n'
            '\tStatus: {}\n'
            .format(self.id, self.status))

class Stack(db.Model):
  id = db.Column(db.Integer, primary_key=True)
  stack_name = db.Column(db.String(64), index=True, unique=True)
  stack_description = db.Column(db.String(128), unique=True)
  is_sandbox_enabled = db.Column(db.Boolean(8), default=False)

  scenarios = db.relationship('Scenario', backref='stack', lazy='dynamic')

  def __repr__(self):
    return ('\tID: {}\n'
            '\tStack name: {}\n'
            '\tStack description: {}\n'
            '\tIs sandbox enabled: {}\n'
            .format(self.id, self.stack_name, self.stack_description, self.is_sandbox_enabled))

class Scenario(db.Model):
  id = db.Column(db.Integer, primary_key=True)
  stack_id = db.Column(db.Integer, db.ForeignKey('stack.id'))
  scenario_name = db.Column(db.String(64), index=True, unique=True)
  scenario_description = db.Column(db.String(128), unique=True)
  is_sandbox_enabled = db.Column(db.Boolean(8), default=False)

  def __repr__(self):
    return ('\tID: {}\n'
            '\tScenario name: {}\n'
            '\tScenario description: {}\n'
            '\tIs sandbox enabled: {}\n'
            .format(self.id, self.scenario_name, self.scenario_description, self.is_sandbox_enabled))