Merge pull request 'Backend merge' (#6) from React-Group/interstellar_ai:main into main

Reviewed-on: https://interstellardevelopment.org/code/code/sageTheDm/interstellar_ai/pulls/6

.gitignore

@@ -38,6 +38,7 @@ next-env.d.ts
 
 .idea/
 venv/
+__pycache__/
 
 key.pem
 cert.pem

py/ai.py

@@ -12,13 +12,9 @@ class AI:
             options={"temperature": 0.5},
         )
 
-        for i in messages:
-            print(i)
-
         return_class.ai_response[access_token] = ""
 
         for chunk in stream:
-            print(chunk['message']['content'])
             return_class.ai_response[access_token] += chunk['message']['content']
 
     @staticmethod

py/api.py

@@ -3,20 +3,23 @@ from flask_cors import CORS
 import secrets
 from ai import AI
 from db import DB
+from OpenSSL import crypto
 
 
 class API:
     def __init__(self):
+        self.crypt_size = 4096
         self.app = Flask(__name__)
         self.ai_response = {}
         self.ai = AI()
         self.db = DB()
+        self.db.load_database()
         CORS(self.app)
 
     def run(self):
         @self.app.route('/interstellar/api/ai_create', methods=['GET'])
         def create_ai():
-            access_token = secrets.token_urlsafe(4096)
+            access_token = secrets.token_urlsafe(self.crypt_size)
             self.ai_response[access_token] = ""
             return jsonify({'status': 200, 'access_token': access_token})
 
@@ -24,11 +27,16 @@ class API:
         def send_ai():
             data = request.get_json()
             messages = data.get('messages')
+            model_type = data.get('model_type')
             ai_model = data.get('ai_model')
             access_token = data.get('access_token')
             if access_token not in self.ai_response:
                 return jsonify({'status': 401, 'error': 'Invalid access token'})
-            self.ai.process_local(ai_model, messages, self, access_token)
+
+            if model_type == "local":
+                self.ai.process_local(ai_model, messages, self, access_token)
+            if model_type == "mistral":
+                self.ai.process_mistralai(ai_model, messages, self, access_token)
             return jsonify({'status': 200})
 
         @self.app.route('/interstellar/api/ai_get', methods=['GET'])
@@ -44,14 +52,49 @@ class API:
             data = request.args.get('data')
             if action == "create_account":
                 self.db.add_user(data)
-            if action == "change_password":
+            elif action == "change_password":
                 self.db.update_password(data)
-            if action == "get_data":
+            elif action == "get_data":
-                self.db.get_additional_info(data)
+                self.db.get_data(data)
-            if action == "check_credentials":
+            elif action == "change_data":
+                self.db.change_data(data)
+            elif action == "check_credentials":
                 self.db.check_credentials(data)
 
-        ssl_context = ('cert.pem', 'key.pem')
+        email_address = "emailAddress"
+        common_name = "commonName"
+        country_name = "NT"
+        locality_name = "localityName"
+        state_or_province_name = "stateOrProvinceName"
+        organization_name = "organizationName"
+        organization_unit_name = "organizationUnitName"
+        serial_number = 0
+        validity_start_in_seconds = 0
+        validity_end_in_seconds = 10 * 365 * 24 * 60 * 60
+        k = crypto.PKey()
+        k.generate_key(crypto.TYPE_RSA, 4096)
+        cert = crypto.X509()
+        cert.get_subject().C = country_name
+        cert.get_subject().ST = state_or_province_name
+        cert.get_subject().L = locality_name
+        cert.get_subject().O = organization_name
+        cert.get_subject().OU = organization_unit_name
+        cert.get_subject().CN = common_name
+        cert.get_subject().emailAddress = email_address
+        cert.set_serial_number(serial_number)
+        cert.gmtime_adj_notBefore(validity_start_in_seconds)
+        cert.gmtime_adj_notAfter(validity_end_in_seconds)
+        cert.set_issuer(cert.get_subject())
+        cert.set_pubkey(k)
+        cert.sign(k, 'sha512')
+
+        with open("cert.pem", "wt") as f:
+            f.write(crypto.dump_certificate(crypto.FILETYPE_PEM, cert).decode("utf-8"))
+
+        with open("key.pem", "wt") as f:
+            f.write(crypto.dump_privatekey(crypto.FILETYPE_PEM, k).decode("utf-8"))
+
+        ssl_context = ("cert.pem", "key.pem")
         self.app.run(debug=True, host='0.0.0.0', port=5000, ssl_context=ssl_context)
 
 

py/db.py

@@ -1,5 +1,5 @@
-import json
 import hashlib
+import json
 
 
 class DB:
@@ -19,15 +19,25 @@ class DB:
         user_data = {"hashed_password": hashed_password}
         self.database[username] = user_data
 
+    def change_data(self, data):
+        username = data.get['username']
+        data = data.get['data']
+        if not self.check_credentials(data):
+            return False
+
+        self.database[username]['data'] = data
+        self.save_database()
+        return True
+
     def update_password(self, data):
         username = data.get['username']
-        old_password = data.get['old_password']
         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):
@@ -40,12 +50,21 @@ class DB:
         entered_hashed_password = self.hash_password(password)
         return stored_hashed_password == entered_hashed_password
 
-    def get_additional_info(self, data):
+    def get_data(self, data):
         username = data.get['username']
-        password = data.get['password']
         if not self.check_credentials(data):
             return None
 
-        send_back = self.database[username]
+        send_back = self.database[username].get['data']
-        del send_back['hashed_password']
         return send_back
+
+    def save_database(self):
+        with open("database.json", 'w') as file:
+            json.dump(self.database, file)
+
+    def load_database(self):
+        try:
+            with open("database.json", 'r') as file:
+                self.database = json.load(file)
+        except FileNotFoundError:
+            pass

py/install.sh

@@ -1,4 +1,3 @@
-openssl req -x509 -newkey rsa:4096 -nodes -keyout key.pem -out cert.pem -days 365
 python -m venv venv
 source venv/bin/activate
 pip install -r requirements.txt

py/requirements.txt

@@ -2,3 +2,4 @@ flask
 flask-cors
 ollama
 mistralai
+pyOpenSSL