interstellar_ai/py/db.py

import hashlib
import json
import os
import pycouchdb


class DB:
    def __init__(self):
        self.database = {}

    def ensure_username(self, data):
        if 'username' in data:
            return data.get('username')
        elif 'email' in data:
            for index, entry in self.database:
                if entry.get('email') == data.get('email'):
                    return index

    @staticmethod
    def hash_password(password):
        salt = "your_secret_salt"
        hashed_password = hashlib.sha256((password + salt).encode()).hexdigest()
        return hashed_password

    def add_user(self, data):
        username = data.get('username')
        password = data.get('password')
        email = data.get('email')
        hashed_password = self.hash_password(password)
        user_data = {"hashed_password": hashed_password, "email": email, "settings": None, "history": None}
        if username not in self.database:
            self.database[username] = user_data
            self.save_database()
            return True
        return False

    def delete_user(self, data):
        username = self.ensure_username(data)
        if not self.check_credentials(data):
            return False

        del self.database[username]
        self.save_database()
        return True

    def update_password(self, data):
        username = self.ensure_username(data)
        new_password = data.get('new_password')
        if not self.check_credentials(data):
            return False

        hashed_new_password = self.hash_password(new_password)
        self.database[username].update({"hashed_password": hashed_new_password})
        self.save_database()
        return True

    def check_credentials(self, data):
        username = self.ensure_username(data)
        password = data.get('password')
        if username not in self.database:
            return False

        stored_hashed_password = self.database[username]["hashed_password"]
        entered_hashed_password = self.hash_password(password)
        print(stored_hashed_password == entered_hashed_password)
        return stored_hashed_password == entered_hashed_password

    def change_settings(self, data):
        username = self.ensure_username(data)
        if not self.check_credentials(data):
            return False

        self.database[username]['settings'] = data.get('data')
        self.save_database()
        return True

    def get_settings(self, data):
        username = self.ensure_username(data)
        if not self.check_credentials(data):
            return None

        send_back = self.database[username].get('settings')
        return send_back

    def change_history(self, data):
        username = self.ensure_username(data)
        if not self.check_credentials(data):
            return False

        self.database[username]['history'] = data.get('data')
        self.save_database()
        return True

    def get_history(self, data):
        username = self.ensure_username(data)
        if not self.check_credentials(data):
            return None

        send_back = self.database[username].get('history')
        return send_back

    def get_email(self, data):
        username = self.ensure_username(data)
        if not self.check_credentials(data):
            return None

        send_back = self.database[username].get('email')
        return send_back

    def get_name(self, data):
        username = self.ensure_username(data)
        if not self.check_credentials(data):
            return None

        send_back = self.ensure_username(data)
        return send_back

    def save_database(self):
        if os.environ.get('PRODUCTION') == "YES":
            server = pycouchdb.Server("http://admin:admin@localhost:5984/")
            db = server.database("interstellar_ai")
            db.save(self.database)

        else:
            with open("database.json", 'w') as file:
                print("saving")
                json.dump(self.database, file)

    def load_database(self):
        if os.environ.get('PRODUCTION') == "YES":
            server = pycouchdb.Server("http://admin:admin@localhost:5984/")
            db = server.database("interstellar_ai")
            if db:
                self.database = db
            else:
                server.create("interstellar_ai")
                db = server.database("interstellar_ai")
                db.save(self.database)
        else:
            try:
                with open("database.json", 'r') as file:
                    self.database = json.load(file)
            except FileNotFoundError:
                pass