Merge pull request #27 from nianhua99/feat-support_custom_add_user

Feat support custom add user

Dockerfile

@@ -4,6 +4,7 @@ COPY . /app
 ENV TZ Asia/Shanghai
 RUN pip install -r requirements.txt --no-cache-dir -i https://mirrors.aliyun.com/pypi/simple/ --trusted-host mirrors.aliyun.com
 EXPOSE 8182
+# 数据库初始化/升级
+RUN flask db upgrade
 # 启动waitress
-#ENTRYPOINT ["python","app.py"]
 ENTRYPOINT ["waitress-serve","--port", "8182", "--call","app:create_app"]

README.md

@@ -34,6 +34,8 @@ $ pip3 install -r requirements.txt
 $ export PANDORA_NEXT_DOMAIN=https://www.baidu.com
 # 修改以下路径为你本机PandoraNext的路径，确保路径中包含config.json
 $ export PANDORA_NEXT_PATH=/path/to/pandora
+# 数据库初始化
+$ flask db upgrade
 # 启动
 $ python3 waitress_run.py
 # 或者在后台启动
@@ -67,13 +69,14 @@ $ nohup python3 waitress_run.py &
 >	  }```
 ## Todo
 - [x] 展示Pandora额度信息
-- [ ] 生成指定账号下各Share Token的用量情况柱状图
+- [x] 生成指定账号下各Share Token的用量情况柱状图
+- [x] 支持预置Token、Refresh Token
 - [ ] Русская адаптация
 - [ ] 支持管理Pool Token
 - [ ] 支持编辑
 - [ ] 支持更多PandoraNext配置
 - [ ] 支持更多验证码
-- [ ] ~~代码优化~~
+- [x] ~~代码优化~~
 ## Star History
 
 [![Star History Chart](https://api.star-history.com/svg?repos=nianhua99/PandoraNext-Helper&type=Date)](https://star-history.com/#nianhua99/PandoraNext-Helper&Date)

app.py

@@ -2,12 +2,12 @@
 import os
 import re
 import secrets
-import sqlite3
 
 from apscheduler.jobstores.sqlalchemy import SQLAlchemyJobStore
-from flask import Flask, g, redirect, url_for
+from flask import Flask, redirect, url_for
 from flask_bootstrap import Bootstrap5
 from flask_login import LoginManager
+from flask_migrate import Migrate
 from flask_moment import Moment
 from flask_apscheduler import APScheduler
 from loguru import logger
@@ -41,57 +41,6 @@ def context_api_prefix():
     return dict(api_prefix=app.config['proxy_api_prefix'])
 
 
-def init_db():
-    with app.app_context():
-        db = connect_db()
-        db.execute('''
-            create table if not exists users
-            (
-                id            INTEGER
-                    primary key autoincrement,
-                email         TEXT not null,
-                password      TEXT not null,
-                session_token TEXT,
-                access_token  TEXT,
-                share_list    TEXT,
-                create_time   datetime,
-                update_time   datetime,
-                shared        INT default 0
-            )
-        ''')
-        db.commit()
-
-
-def connect_db():
-    db = getattr(g, '_database', None)
-    if db is None:
-        db = g._database = sqlite3.connect(os.path.join(app.config['pandora_path'], DATABASE))
-        db.row_factory = sqlite3.Row
-    return db
-
-
-def query_db(query, args=(), one=False):
-    # 查看g对象是否存在db属性，如果不存在则创建
-    if not hasattr(g, 'db'):
-        g.db = connect_db()
-    cur = g.db.execute(query, args)
-    rv = [dict((cur.description[idx][0], value)
-               for idx, value in enumerate(row)) for row in cur.fetchall()]
-    return (rv[0] if rv else None) if one else rv
-
-
-@app.before_request
-def before_request():
-    g.db = connect_db()
-
-
-@app.after_request
-def after_request(result):
-    if hasattr(g, 'db'):
-        g.db.close()
-    return result
-
-
 def check_require_config():
     PANDORA_NEXT_PATH = os.getenv('PANDORA_NEXT_PATH')
     # 如果PANDORA_NEXT_PATH 为空则检查/data下是否存在config.json
@@ -148,26 +97,47 @@ def check_require_config():
 
 from auth import auth
 from main import main
+from model import db
 
 check_require_config()
-init_db()
 
-#scheduler jobstore
+# scheduler jobstore
 app.config['SCHEDULER_JOBSTORES'] = {
     'default': SQLAlchemyJobStore(url='sqlite:///' + os.path.join(app.config['pandora_path'], DATABASE))
 }
 scheduler = APScheduler()
 scheduler.init_app(app)
 scheduler.start()
 
+app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///' + os.path.join(app.config['pandora_path'], DATABASE)
+app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = True
+db.init_app(app)
+
+
+def include_object(object, name, type_, reflected, compare_to):
+    if (
+            type_ == "table" and name == "apscheduler_jobs"
+    ):
+        return False
+    else:
+        return True
+
+
+migrate = Migrate(include_object=include_object)
+migrate.init_app(app, db)
+
+
+def format_datetime(value):
+    """Format a datetime to a string."""
+    if value is None:
+        return ""
+    return value.strftime('%Y-%m-%d %H:%M:%S')
+
 
 def create_app():
     app.register_blueprint(auth.auth_bp, url_prefix='/' + app.config['proxy_api_prefix'])
     app.register_blueprint(main.main_bp, url_prefix='/' + app.config['proxy_api_prefix'])
-    # 设置日志等级
-    import logging
-    logging.basicConfig()
-    logging.getLogger('apscheduler').setLevel(logging.DEBUG)
+    app.jinja_env.filters['datetime'] = format_datetime
     return app
 
 

login_tools.py

@@ -19,9 +19,11 @@ def login(username, password):
         'Content-Type': 'application/x-www-form-urlencoded'
     }
     response = requests.request("POST", host + "/api/auth/login", headers=headers, data=payload)
+    logger.info("登录结果：{},{}", response.text, response.status_code)
     if response.status_code != 200:
+        if response.status_code == 404:
+            raise Exception("接口不存在,请检查Pandora是否配置正确")
         raise Exception(response.text)
-    logger.info("登录结果：{}", response.json())
     return response.json()
 
 
@@ -33,7 +35,42 @@ def get_access_token(session_token):
         'Content-Type': 'application/x-www-form-urlencoded'
     }
     response = requests.request("POST", host + "/api/auth/session", headers=headers, data=payload)
-    logger.info("获取access_token结果：{},{}", response.json(), response.status_code)
+    logger.info("获取access_token结果：{},{}", response.text, response.status_code)
+    if response.status_code != 200:
+        raise Exception("获取access_token失败")
+    return response.json()
+
+
+# 获取Refresh Token
+def get_refresh_token(username, password):
+    host = get_host()
+    headers = {
+        'Content-Type': 'application/x-www-form-urlencoded'
+    }
+    payload = {
+        'username': username,
+        'password': password
+    }
+    response = requests.request("POST", host + "/api/auth/login2", headers=headers, data=payload)
+    logger.info("获取refresh_token结果：{},{}", response.text, response.status_code)
+
+    if response.status_code != 200:
+        raise Exception("获取refresh_token失败")
+
+    return response.json()
+
+
+# 使用refresh_token获取access_token /api/auth/refresh
+def get_access_token_by_refresh_token(refresh_token):
+    host = get_host()
+    headers = {
+        'Content-Type': 'application/x-www-form-urlencoded'
+    }
+    payload = {
+        'refresh_token': refresh_token
+    }
+    response = requests.request("POST", host + "/api/auth/refresh", headers=headers, data=payload)
+    logger.info("获取access_token结果：{},{}", response.text, response.status_code)
     if response.status_code != 200:
         raise Exception("获取access_token失败")
     return response.json()