From b86bdb74210fde3fc5905defc7b81427400e562f Mon Sep 17 00:00:00 2001
From: Patrick_Pluto <patrick_pluto@noreply.codeberg.org>
Date: Thu, 19 Sep 2024 12:46:52 +0200
Subject: [PATCH] Security Update.

---
 .gitignore          |  3 +++
 py/api.py           | 20 +++++++++++++++-----
 py/install.sh       |  1 +
 py/requirements.txt |  1 +
 4 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/.gitignore b/.gitignore
index e28c376..bb3b79a 100644
--- a/.gitignore
+++ b/.gitignore
@@ -38,3 +38,6 @@ next-env.d.ts
 
 .idea/
 venv/
+
+key.pem
+cert.pem
diff --git a/py/api.py b/py/api.py
index 6eb7cc0..d372b22 100644
--- a/py/api.py
+++ b/py/api.py
@@ -1,5 +1,7 @@
 from flask import Flask, request, jsonify
+from flask_cors import CORS
 import ollama
+import secrets
 
 
 class AI:
@@ -19,14 +21,16 @@ class AI:
 class API:
     def __init__(self):
         self.app = Flask(__name__)
-        self.ai_response = []
+        self.ai_response = {}
         self.ai = AI()
+        CORS(self.app)
 
     def run(self):
         @self.app.route('/interstellar/api/ai_create', methods=['GET'])
         def create_ai():
-            self.ai_response.append("")
-            return jsonify({'status': 200, 'access_token': len(self.ai_response) - 1})
+            access_token = secrets.token_urlsafe(4096)
+            self.ai_response[access_token] = ""
+            return jsonify({'status': 200, 'access_token': access_token})
 
         @self.app.route('/interstellar/api/ai_send', methods=['POST'])
         def send_ai():
@@ -35,17 +39,23 @@ class API:
             ai_model = data.get('ai_model')
             system_prompt = data.get('system_prompt')
             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, message, system_prompt, self, access_token)
             return jsonify({'status': 200})
 
         @self.app.route('/interstellar/api/ai_get', methods=['GET'])
         def get_ai():
             data = request.args.get('access_token')
-            return jsonify({'status': 200, 'response': self.ai_response[int(data)]})
+            if data not in self.ai_response:
+                return jsonify({'status': 401, 'error': 'Invalid access token'})
+            return jsonify({'status': 200, 'response': self.ai_response[data]})
 
-        self.app.run(debug=True)
+        ssl_context = ('cert.pem', 'key.pem')
+        self.app.run(debug=True, host='0.0.0.0', port=5000, ssl_context=ssl_context)
 
 
 if __name__ == '__main__':
     api = API()
     api.run()
+
diff --git a/py/install.sh b/py/install.sh
index 1fbdcba..31bb21e 100644
--- a/py/install.sh
+++ b/py/install.sh
@@ -1,3 +1,4 @@
+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
diff --git a/py/requirements.txt b/py/requirements.txt
index 731af42..1e95d6f 100644
--- a/py/requirements.txt
+++ b/py/requirements.txt
@@ -1,2 +1,3 @@
 flask
+flask-cors
 ollama
\ No newline at end of file