sysvitals/web/db.py

import hashlib
import hmac
import os

import tornado.gen
import psycopg2
import momoko

import config

def hash_pw(password, key=None):
	if key is None:
		key = os.urandom(16)
		digest = hmac.new(key, password, hashlib.sha256)
		key = key.encode("hex")
		hashed = digest.hexdigest()
		return key, hashed

class MomokoDB:
	db = momoko.Pool(dsn='dbname=%s user=%s' % (config.db.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, username, password):
		salt, hashed_password = hash_pw(password)
		query = 'INSERT INTO users (username, password, salt) VALUES (%s, %s, %s);'
		yield self.execute(query, username, hashed_password, salt)

	@tornado.gen.coroutine
	def get_user(self, username):
		query = 'SELECT * FROM users WHERE username=%s;'
		cursor = yield self.execute(query, username)
		return cursor.fetchone()

	@tornado.gen.coroutine
	def check_user(self, username, password):
		user = yield self.get_user(username)
		if not user:
			return
		_, hashed = hash_pw(password, user['salt'].decode("hex"))
		if hashed == user['password']:
			return user