sysvitals/web/db.py

import hashlib
import hmac
import binascii
import os

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 = h.hexdigest()
	salt_hex = binascii.hexlify(salt).decode()
	return hashed, salt_hex

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):
		hashed_password, salt = 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
		salt = binascii.unhexlify(bytes(user['salt'], 'ascii'))
		hashed, _ = hash_pw(password, salt)
		print(hashed)
		print(user['password'])
		if hashed == user['password']:
			return user