diff --git a/py/db.py b/py/db.py index c968561..fa8ab77 100644 --- a/py/db.py +++ b/py/db.py @@ -3,12 +3,13 @@ import json import os import pycouchdb - class DB: def __init__(self): + # Initialize the database dictionary to store user data self.database = {} def ensure_username(self, data): + # Ensure a username can be retrieved either from username or email if "username" in data: return data.get("username") elif "email" in data: @@ -18,15 +19,17 @@ class DB: @staticmethod def hash_password(password): - salt = "your_secret_salt" + # Hash the password with a salt for secure storage + salt = "your_secret_salt" # Consider using a secure random salt hashed_password = hashlib.sha256((password + salt).encode()).hexdigest() return hashed_password def add_user(self, data): + # Add a new user to the database if username is unique username = data.get("username") password = data.get("password") email = data.get("email") - hashed_password = self.hash_password(password) + hashed_password = self.hash_password(password) # Hash the password user_data = { "hashed_password": hashed_password, "email": email, @@ -35,112 +38,123 @@ class DB: } if username not in self.database: self.database[username] = user_data - self.save_database() + self.save_database() # Save changes to the database return True - return False + return False # User already exists def delete_user(self, data): + # Delete a user from the database if credentials are valid username = self.ensure_username(data) if not self.check_credentials(data): - return False + return False # Invalid credentials - del self.database[username] - self.save_database() + del self.database[username] # Remove user from database + self.save_database() # Save changes return True def update_password(self, data): + # Update the user's password if credentials are valid username = self.ensure_username(data) new_password = data.get("new_password") if not self.check_credentials(data): - return False + return False # Invalid credentials - hashed_new_password = self.hash_password(new_password) + hashed_new_password = self.hash_password(new_password) # Hash the new password self.database[username].update({"hashed_password": hashed_new_password}) - self.save_database() + self.save_database() # Save changes return True def check_credentials(self, data): + # Verify if provided credentials match stored data username = self.ensure_username(data) password = data.get("password") if username not in self.database: - return False + return False # User not found stored_hashed_password = self.database[username]["hashed_password"] entered_hashed_password = self.hash_password(password) - return stored_hashed_password == entered_hashed_password + return stored_hashed_password == entered_hashed_password # Check hashed password def change_settings(self, data): + # Change user settings if credentials are valid username = self.ensure_username(data) if not self.check_credentials(data): - return False + return False # Invalid credentials - self.database[username]["settings"] = data.get("data") - self.save_database() + self.database[username]["settings"] = data.get("data") # Update settings + self.save_database() # Save changes return True def get_settings(self, data): + # Retrieve user settings if credentials are valid username = self.ensure_username(data) if not self.check_credentials(data): - return None + return None # Invalid credentials - send_back = self.database[username].get("settings") + send_back = self.database[username].get("settings") # Get settings return send_back def change_history(self, data): + # Change user history if credentials are valid username = self.ensure_username(data) if not self.check_credentials(data): - return False + return False # Invalid credentials - self.database[username]["history"] = data.get("data") - self.save_database() + self.database[username]["history"] = data.get("data") # Update history + self.save_database() # Save changes return True def get_history(self, data): + # Retrieve user history if credentials are valid username = self.ensure_username(data) if not self.check_credentials(data): - return None + return None # Invalid credentials - send_back = self.database[username].get("history") + send_back = self.database[username].get("history") # Get history return send_back def get_email(self, data): + # Retrieve user email if credentials are valid username = self.ensure_username(data) if not self.check_credentials(data): - return None + return None # Invalid credentials - send_back = self.database[username].get("email") + send_back = self.database[username].get("email") # Get email return send_back def get_name(self, data): + # Retrieve username if credentials are valid if not self.check_credentials(data): - return None + return None # Invalid credentials - send_back = self.ensure_username(data) + send_back = self.ensure_username(data) # Get username return send_back def save_database(self): + # Save the database to the specified storage (CouchDB or JSON file) if os.environ.get("PRODUCTION") == "YES": server = pycouchdb.Server("http://admin:admin@localhost:5984/") db = server.database("interstellar_ai") - db.save(self.database) + db.save(self.database) # Save to CouchDB else: with open("database.json", "w") as file: - json.dump(self.database, file) + json.dump(self.database, file) # Save to JSON file def load_database(self): + # Load the database from the specified storage (CouchDB or JSON file) if os.environ.get("PRODUCTION") == "YES": server = pycouchdb.Server("http://admin:admin@localhost:5984/") db = server.database("interstellar_ai") if db: - self.database = db + self.database = db # Load from CouchDB else: - server.create("interstellar_ai") + server.create("interstellar_ai") # Create database if it doesn't exist db = server.database("interstellar_ai") - db.save(self.database) + db.save(self.database) # Save initial empty database else: try: with open("database.json", "r") as file: - self.database = json.load(file) + self.database = json.load(file) # Load from JSON file except FileNotFoundError: - pass + pass # File not found, do nothing