UnderCover/undercover/db.py

150 lines
4.1 KiB
Python

import os
from dataclasses import dataclass
from typing import Optional
import bcrypt
import psycopg
@dataclass
class Letter:
id: int
title: str
contents: str
@dataclass
class User:
id: int
email: str
@dataclass
class UserWithHash:
id: int
email: str
password_hash: str
def connect():
return psycopg.connect(
host=os.environ['UNDERCOVER_POSTGRES_HOST'],
dbname=os.environ['UNDERCOVER_POSTGRES_DBNAME'],
port=os.environ['UNDERCOVER_POSTGRES_PORT'],
user=os.environ['UNDERCOVER_POSTGRES_USER'],
password=os.environ['UNDERCOVER_POSTGRES_PASSWORD'])
def connected(action):
with connect() as con:
cur = con.cursor()
return action(cur, con)
def login(user_email: str, password: str):
pw_bytes: bytes = password.encode('utf-8')
user = __get_user(user_email)
if user:
return bcrypt.checkpw(pw_bytes, user.password_hash.encode('utf-8'))
return False
def add_user(username: str, password: str):
pw_bytes = password.encode('utf-8')
salt = bcrypt.gensalt()
pw_hash = bcrypt.hashpw(pw_bytes, salt)
decoded = pw_hash.decode('utf-8')
with connect() as con:
cur = con.cursor()
cur.execute("INSERT INTO users(email, password) VALUES (%s, %s)", (username, decoded))
con.commit()
def delete_user(username: str):
with connect() as con:
cur = con.cursor()
cur.execute("DELETE FROM users WHERE email = %s", (username,))
con.commit()
def add_user_lambda(username: str, password: str):
def f(cur, con):
cur.execute("INSERT INTO users(email, password) VALUES (%s, %s)", (username, password))
con.commit()
connected(f)
def add_letter(user_id: int, letter_title: str, letter_content: str):
with connect() as con:
cur = con.cursor()
cur.execute("INSERT INTO letter_data(user_id, letter_name, letter_data) VALUES (%s, %s, %s)",
(user_id, letter_title, letter_content))
con.commit()
def edit_letter(letter_id: int, letter_title: str, letter_content: str):
with connect() as con:
cur = con.cursor()
cur.execute("UPDATE letter_data SET letter_name = %s, letter_data = %s WHERE id = %s",
(letter_title, letter_content, letter_id))
con.commit()
def get_user_letters(user_id: int) -> [Letter]:
with connect() as con:
cur = con.cursor()
cur.execute("SELECT id, letter_name, letter_data FROM letter_data WHERE user_id = %s", (str(user_id),))
return list(map(lambda row: Letter(row[0], row[1], row[2]), cur.fetchall()))
def get_user(email: str) -> Optional[User]:
user = __get_user(email)
if user:
return User(user.id, user.email)
return None
def __get_user(email: str) -> Optional[UserWithHash]:
"""
:param email:
:return: User without their password_hash
"""
with connect() as con:
cur = con.cursor()
cur.execute("SELECT id, password FROM users WHERE users.email = %s", (email,))
row = cur.fetchone()
if row:
user_id, password = row
return UserWithHash(user_id, email, password)
return None
def get_users() -> [UserWithHash]:
with connect() as con:
cur = con.cursor()
cur.execute("SELECT id, email, password FROM users")
return map(lambda row: UserWithHash(row[0], row[1], row[2]), cur.fetchall())
if __name__ == "__main__":
add_user("hash_man", "hashword")
print("Can pull correctly: " + str(login("hash_man", "hashword")))
delete_user("hash_man")
# add_letter(1, "Dynamically-added", "This is a letter added from Python!")
# edit_letter(3, "Dynamically edited!", "This letter was dynamically edited from Python!")
# for letter in get_user_letters(1):
# print("\'" + letter.title + "\"" + ":")
# print(" id: " + str(letter.id))
# print(" letter-data: " + letter.contents)
# print()
# for user in get_users():
# print(user.email + ":")
# print(" id: " + str(user.id))
# print(" password: " + user.password_hash)
# print()