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()