From 3f29464dbdfa6563ad8e2dae950abc9961738bf6 Mon Sep 17 00:00:00 2001 From: RockChinQ <1010553892@qq.com> Date: Wed, 13 Dec 2023 14:53:39 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E6=A0=87=E8=AF=86=E7=AC=A6=E7=94=9F?= =?UTF-8?q?=E6=88=90=E5=99=A8=E6=A8=A1=E5=9D=97?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .gitignore | 3 +- main.py | 5 ++ pkg/audit/identifier.py | 83 ++++++++++++++++++++++++ pkg/database/manager.py | 2 +- tests/identifier_test/host_identifier.py | 43 ++++++++++++ 5 files changed, 134 insertions(+), 2 deletions(-) create mode 100644 pkg/audit/identifier.py create mode 100644 tests/identifier_test/host_identifier.py diff --git a/.gitignore b/.gitignore index 9d4595de..d671df3a 100644 --- a/.gitignore +++ b/.gitignore @@ -30,4 +30,5 @@ qcapi claude.json bard.json /*yaml -!/docker-compose.yaml \ No newline at end of file +!/docker-compose.yaml +res/instance_id.json \ No newline at end of file diff --git a/main.py b/main.py index 043500ef..155620a5 100644 --- a/main.py +++ b/main.py @@ -145,6 +145,11 @@ async def start_process(first_time_init=False): global known_exception_caught import pkg.utils.context + # 计算host和instance标识符 + import pkg.audit.identifier + pkg.audit.identifier.init() + pkg.audit.identifier.print_out() + # 加载配置 cfg_inst: pymodule_cfg.PythonModuleConfigFile = pymodule_cfg.PythonModuleConfigFile( 'config.py', diff --git a/pkg/audit/identifier.py b/pkg/audit/identifier.py new file mode 100644 index 00000000..334b7e3b --- /dev/null +++ b/pkg/audit/identifier.py @@ -0,0 +1,83 @@ +import os +import uuid +import json +import time + + +identifier = { + 'host_id': '', + 'instance_id': '', + 'host_create_ts': 0, + 'instance_create_ts': 0, +} + +HOST_ID_FILE = os.path.expanduser('~/.qchatgpt/host_id.json') +INSTANCE_ID_FILE = 'res/instance_id.json' + +def init(): + global identifier + + if not os.path.exists(os.path.expanduser('~/.qchatgpt')): + os.mkdir(os.path.expanduser('~/.qchatgpt')) + + if not os.path.exists(HOST_ID_FILE): + new_host_id = 'host_'+str(uuid.uuid4()) + new_host_create_ts = int(time.time()) + + with open(HOST_ID_FILE, 'w') as f: + json.dump({ + 'host_id': new_host_id, + 'host_create_ts': new_host_create_ts + }, f) + + identifier['host_id'] = new_host_id + identifier['host_create_ts'] = new_host_create_ts + else: + loaded_host_id = '' + loaded_host_create_ts = 0 + + with open(HOST_ID_FILE, 'r') as f: + file_content = json.load(f) + loaded_host_id = file_content['host_id'] + loaded_host_create_ts = file_content['host_create_ts'] + + identifier['host_id'] = loaded_host_id + identifier['host_create_ts'] = loaded_host_create_ts + + # 检查实例 id + if os.path.exists(INSTANCE_ID_FILE): + instance_id = {} + with open(INSTANCE_ID_FILE, 'r') as f: + instance_id = json.load(f) + + if instance_id['host_id'] != identifier['host_id']: # 如果实例 id 不是当前主机的，删除 + os.remove(INSTANCE_ID_FILE) + + if not os.path.exists(INSTANCE_ID_FILE): + new_instance_id = 'instance_'+str(uuid.uuid4()) + new_instance_create_ts = int(time.time()) + + with open(INSTANCE_ID_FILE, 'w') as f: + json.dump({ + 'host_id': identifier['host_id'], + 'instance_id': new_instance_id, + 'instance_create_ts': new_instance_create_ts + }, f) + + identifier['instance_id'] = new_instance_id + identifier['instance_create_ts'] = new_instance_create_ts + else: + loaded_instance_id = '' + loaded_instance_create_ts = 0 + + with open(INSTANCE_ID_FILE, 'r') as f: + file_content = json.load(f) + loaded_instance_id = file_content['instance_id'] + loaded_instance_create_ts = file_content['instance_create_ts'] + + identifier['instance_id'] = loaded_instance_id + identifier['instance_create_ts'] = loaded_instance_create_ts + +def print_out(): + global identifier + print(identifier) diff --git a/pkg/database/manager.py b/pkg/database/manager.py index f410b418..ad44d512 100644 --- a/pkg/database/manager.py +++ b/pkg/database/manager.py @@ -91,7 +91,7 @@ class DatabaseManager: `json` text not null ) """) - print('Database initialized.') + # print('Database initialized.') # session持久化 def persistence_session(self, subject_type: str, subject_number: int, create_timestamp: int, diff --git a/tests/identifier_test/host_identifier.py b/tests/identifier_test/host_identifier.py new file mode 100644 index 00000000..64834931 --- /dev/null +++ b/tests/identifier_test/host_identifier.py @@ -0,0 +1,43 @@ +import os +import uuid +import json + +# 向 ~/.qchatgpt 写入一个 标识符 + +if not os.path.exists(os.path.expanduser('~/.qchatgpt')): + os.mkdir(os.path.expanduser('~/.qchatgpt')) + +identifier = { + "host_id": "host_"+str(uuid.uuid4()), +} + +if not os.path.exists(os.path.expanduser('~/.qchatgpt/host.json')): + print('create ~/.qchatgpt/host.json') + with open(os.path.expanduser('~/.qchatgpt/host.json'), 'w') as f: + json.dump(identifier, f) +else: + print('load ~/.qchatgpt/host.json') + with open(os.path.expanduser('~/.qchatgpt/host.json'), 'r') as f: + identifier = json.load(f) + +print(identifier) + +instance_id = { + "host_id": identifier['host_id'], + "instance_id": "instance_"+str(uuid.uuid4()), +} + +# 实例 id +if os.path.exists("res/instance_id.json"): + with open("res/instance_id.json", 'r') as f: + instance_id = json.load(f) + + if instance_id['host_id'] != identifier['host_id']: + os.remove("res/instance_id.json") + +if not os.path.exists("res/instance_id.json"): + print('create res/instance_id.json') + with open("res/instance_id.json", 'w') as f: + json.dump(instance_id, f) + +print(instance_id) \ No newline at end of file