From 9888dcefc95cba9571a9a0170716f2efe87160e5 Mon Sep 17 00:00:00 2001
From: pengzhile <pengzhile@gmail.com>
Date: Thu, 13 Apr 2023 19:15:37 +0800
Subject: [PATCH] support access token switching in server mode

Signed-off-by: pengzhile <pengzhile@gmail.com>
---
 README.md                  |  1 +
 src/pandora/__init__.py    |  2 +-
 src/pandora/bots/server.py | 24 ++++++++++++++++++------
 3 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/README.md b/README.md
index 4957f75e..b2bcb0e1 100644
--- a/README.md
+++ b/README.md
@@ -137,6 +137,7 @@
 ## HTTP服务文档
 
 * 如果你以`http`服务方式启动，现在你可以打开一个极简版的`ChatGPT`了。通过你指定的`http://ip:port`来访问。
+* 通过`http://ip:port/?token=xxx`，传递一个Token的名字，可以切换到对应的`Access Token`。
 * API文档见：[doc/HTTP-API.md](https://github.com/pengzhile/pandora/blob/master/doc/HTTP-API.md)
 
 ## 操作命令
diff --git a/src/pandora/__init__.py b/src/pandora/__init__.py
index d1a82ef6..e5f75a4d 100644
--- a/src/pandora/__init__.py
+++ b/src/pandora/__init__.py
@@ -1,3 +1,3 @@
 # -*- coding: utf-8 -*-
 
-__version__ = '1.0.0'
+__version__ = '1.0.1'
diff --git a/src/pandora/bots/server.py b/src/pandora/bots/server.py
index a789dc72..a462b048 100644
--- a/src/pandora/bots/server.py
+++ b/src/pandora/bots/server.py
@@ -1,6 +1,7 @@
 # -*- coding: utf-8 -*-
 
 import logging
+from datetime import timedelta
 from os.path import join, abspath, dirname
 
 from flask import Flask, jsonify, make_response, request, Response, render_template
@@ -102,18 +103,29 @@ def __handle_error(self, e):
             'message': str(e.original_exception if self.debug and hasattr(e, 'original_exception') else e.name)
         }), 500)
 
+    @staticmethod
+    def __set_cookie(resp, token_key, max_age):
+        resp.set_cookie('token-key', token_key, max_age=max_age, path='/', domain=None, httponly=True, samesite='Lax')
+
     @staticmethod
     def __get_token_key():
-        return request.headers.get('X-Use-Token', None)
+        return request.headers.get('X-Use-Token', request.cookies.get('token-key'))
 
     def chat(self, conversation_id=None):
         query = {'chatId': [conversation_id]} if conversation_id else {}
 
-        return render_template('chat.html',
-                               pandora_base=request.url_root.strip('/'),
-                               pandora_sentry=self.sentry,
-                               query=query
-                               )
+        token_key = request.args.get('token')
+        rendered = render_template('chat.html',
+                                   pandora_base=request.url_root.strip('/'),
+                                   pandora_sentry=self.sentry,
+                                   query=query
+                                   )
+        resp = make_response(rendered)
+
+        if token_key:
+            self.__set_cookie(resp, token_key, timedelta(days=30))
+
+        return resp
 
     @staticmethod
     def session():