sysvitals/web/db.py

import binascii
from collections import defaultdict
import hashlib
import hmac
import os
import time

import tornado.gen
import psycopg2
import momoko

import config

def hash_pw(password, salt=None):
	if salt is None:
		salt = os.urandom(16)
	h = hmac.new(salt, password.encode('utf-8'), hashlib.sha256)
	hashed = hashlib.pbkdf2_hmac('sha512', password.encode('utf-8'), salt, 100000)
	hashed_hex = binascii.hexlify(hashed).decode()
	salt_hex = binascii.hexlify(salt).decode()
	return hashed_hex, salt_hex

class MomokoDB:
	db = momoko.Pool(dsn='dbname=%s user=%s' % (config.database, config.db_user), size=2)

	@tornado.gen.coroutine
	def execute(self, query, *args):
		result = yield momoko.Op(self.db.execute, query, args, cursor_factory=psycopg2.extras.DictCursor)
		return result

	@tornado.gen.coroutine
	def create_user(self, email, password):
		hashed, salt = hash_pw(password)
		sql = 'INSERT INTO users (email, password, salt) VALUES (%s, %s, %s) RETURNING id;'
		cursor = yield self.execute(sql, email, hashed, salt)
		return cursor.fetchone()['id']

	@tornado.gen.coroutine
	def check_user(self, email, password):
		sql = 'SELECT id, email, password, salt FROM users WHERE email=%s;'
		cursor = yield self.execute(sql, email)
		user = cursor.fetchone()
		if not user:
			return
		salt = binascii.unhexlify(user['salt'].encode())
		hashed, _ = hash_pw(password, salt)
		if hashed == user['password']:
			return user

	@tornado.gen.coroutine
	def create_group(self, user_id, group_name):
		hmac_msg = ('%d%s%d' % (user_id, group_name, time.time())).encode('utf-8')
		h = hmac.new(config.cookie_secret.encode('utf-8'), hmac_msg, hashlib.sha1)
		sql = 'INSERT INTO groups (name, api_key) VALUES(%s, %s) RETURNING id;'
		cursor = yield self.execute(sql, group_name, h.hexdigest())
		group_id = cursor.fetchone()['id']
		yield self.execute('INSERT INTO user_groups (user_id, group_id) VALUES(%s, %s);', user_id, group_id)
		return group_id

	@tornado.gen.coroutine
	def invite_user_group(self, email, group_id):
		cursor = yield self.execute('SELECT id FROM users WHERE email = %s;', email)
		user_id = cursor.fetchone()['id']
		yield self.execute('INSERT INTO user_groups (user_id, group_id) VALUES(%s, %s);', user_id, group_id)

	@tornado.gen.coroutine
	def get_groups(self, user_id):
		cursor = yield self.execute('''
			SELECT groups.id, groups.name, api_key FROM user_groups
			JOIN groups ON user_groups.group_id = groups.id
			WHERE user_id = %s;
		''', user_id)
		return cursor.fetchall()

	@tornado.gen.coroutine
	def get_servers(self, user_id):
		cursor = yield self.execute('''
			SELECT servers.id, servers.group_id, servers.hostname FROM user_groups
			JOIN servers ON user_groups.group_id = servers.group_id
			WHERE user_id = %s;
		''', user_id)
		servers = defaultdict(list)
		for row in cursor.fetchall():
			servers[row['group_id']].append(row)
		return servers

	@tornado.gen.coroutine
	def get_api_key(self, group_id):
		cursor = yield self.execute('SELECT api_key FROM groups WHERE id = %s', group_id)
		return cursor.fetchone()['api_key']