From 2284f437b95f129ab1459236d2de34b69e24102d Mon Sep 17 00:00:00 2001
From: ngfchl <ngfchl@126.com>
Date: Tue, 3 Jan 2023 18:37:23 +0800
Subject: [PATCH] =?UTF-8?q?=E5=8D=87=E7=BA=A7=20=E4=BF=AE=E6=94=B9ptools.y?=
 =?UTF-8?q?aml=E4=B8=BAptools.toml=EF=BC=8C=E5=B9=B6=E6=B7=BB=E5=8A=A0?=
 =?UTF-8?q?=E8=87=AA=E5=8A=A8=E8=BD=AC=E6=8D=A2=E5=8A=9F=E8=83=BD?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 auto_pt/views.py       | 17 +++++------------
 pt_site/UtilityTool.py | 29 +++++++++++++++++++++++++++++
 pt_site/views.py       | 14 ++++++++++++--
 3 files changed, 46 insertions(+), 14 deletions(-)

diff --git a/auto_pt/views.py b/auto_pt/views.py
index eac794c..9d24dfe 100644
--- a/auto_pt/views.py
+++ b/auto_pt/views.py
@@ -1099,17 +1099,11 @@ def get_site_torrents(request):
 
 def get_config_api(request):
     file_path = os.path.join(BASE_DIR, 'db/ptools.toml')
-    yaml_file_path = os.path.join(BASE_DIR, 'db/ptools.yaml')
-    if not os.path.exists(file_path):
-        data = ''
-        if os.path.exists(yaml_file_path):
-            with open('db/ptools.yaml', 'r') as yaml_file:
-                data = yaml.load(yaml_file, Loader=yaml.FullLoader)
-                logger.info(f'原始文档{data}')
-        with open(file_path, 'w') as toml_f:
-            toml_f.write('')
-            toml.dump(data, toml_f)
-            logger.info(f'配置文件生成成功！')
+    res = pt_spider.generate_config_file()
+    if res.code != 0:
+        return JsonResponse(data=CommonResponse.error(
+            msg=res.msg
+        ).to_dict(), safe=False)
     try:
         with open(file_path, 'rb') as f:
             response = HttpResponse(f)
@@ -1118,7 +1112,6 @@ def get_config_api(request):
                 data=response.content.decode('utf8')
             ).to_dict(), safe=False)
     except Exception as e:
-        raise
         return JsonResponse(data=CommonResponse.error(
             msg='获取配置文件信息失败！'
         ).to_dict(), safe=False)
diff --git a/pt_site/UtilityTool.py b/pt_site/UtilityTool.py
index eaa88a5..3dc8553 100644
--- a/pt_site/UtilityTool.py
+++ b/pt_site/UtilityTool.py
@@ -1,6 +1,7 @@
 import contextlib
 import json
 import logging
+import os
 import random
 import re
 import ssl
@@ -16,6 +17,7 @@ import dateutil.parser
 import opencc
 import qbittorrentapi
 import requests
+import toml
 import transmission_rpc
 import urllib3.util.ssl_
 import yaml
@@ -30,6 +32,7 @@ from wxpusher import WxPusher
 from auto_pt.models import Notify, OCR
 from pt_site.models import MySite, SignIn, TorrentInfo, SiteStatus, Site
 from ptools.base import TorrentBaseInfo, PushConfig, CommonResponse, StatusCodeEnum, DownloaderCategory
+from ptools.settings import BASE_DIR
 
 urllib3.util.ssl_.DEFAULT_CIPHERS = 'ALL'
 
@@ -2282,3 +2285,29 @@ class PtSpider:
                 msg=message,
                 data=0
             )
+
+    @staticmethod
+    def generate_config_file():
+        file_path = os.path.join(BASE_DIR, 'db/ptools.toml')
+        yaml_file_path = os.path.join(BASE_DIR, 'db/ptools.yaml')
+        try:
+            if not os.path.exists(file_path):
+                data = ''
+                if os.path.exists(yaml_file_path):
+                    with open('db/ptools.yaml', 'r') as yaml_file:
+                        data = yaml.load(yaml_file, Loader=yaml.FullLoader)
+                        logger.info(f'原始文档{data}')
+                with open(file_path, 'w') as toml_f:
+                    toml_f.write('')
+                    toml.dump(data, toml_f)
+                    logger.info(f'配置文件生成成功！')
+                    return CommonResponse.success(
+                        msg='配置文件生成成功！',
+                    )
+            return CommonResponse.success(
+                msg='配置文件文件已存在！',
+            )
+        except Exception as e:
+            return CommonResponse.success(
+                msg=f'初始化失败！{e}',
+            )
diff --git a/pt_site/views.py b/pt_site/views.py
index a3c6408..acebf52 100644
--- a/pt_site/views.py
+++ b/pt_site/views.py
@@ -9,6 +9,7 @@ import time
 from concurrent.futures.thread import ThreadPoolExecutor
 
 import requests
+import toml
 import yaml
 from apscheduler.schedulers.background import BackgroundScheduler
 from django_apscheduler.jobstores import DjangoJobStore
@@ -301,10 +302,19 @@ def auto_upgrade():
 
 def auto_update_license():
     """auto_update_license"""
-    with open('db/ptools.yaml', 'r') as f:
-        data = yaml.load(f, Loader=yaml.FullLoader)
+    res = pt_spider.generate_config_file()
+    if res.code != 0:
+        return CommonResponse.error(
+            msg=res.msg
+        )
+    with open('db/ptools.toml', 'r') as f:
+        data = toml.load(f)
         print(data)
     pt_helper = data.get('pt_helper')
+    if len(pt_helper) <= 0:
+        return CommonResponse.error(
+            msg='请先配置小助手相关信息再进行操作！'
+        )
     host = pt_helper.get('host')
     username = pt_helper.get('username')
     password = pt_helper.get('password')