284 lines
9.9 KiB
Python
284 lines
9.9 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
"""
|
|
pyqmt安全包装模块
|
|
|
|
对原始pyqmt模块进行安全包装,集成JWT令牌验证
|
|
防止未授权用户绕过验证直接使用交易功能
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import yaml
|
|
from datetime import datetime
|
|
from typing import Optional, Any, Dict
|
|
from functools import wraps
|
|
|
|
# 导入JWT验证工具
|
|
from jwt_utils import 验证访问令牌, 检查令牌权限, 检查令牌角色, 生成机器指纹
|
|
|
|
# 导入原始pyqmt模块
|
|
try:
|
|
from pyqmt import pyqmt as _original_qmt
|
|
except ImportError as e:
|
|
print(f"[错误] 无法导入pyqmt模块: {e}")
|
|
print("[提示] 请确保pyqmt模块已正确安装")
|
|
sys.exit(1)
|
|
|
|
class SecurityError(Exception):
|
|
"""安全验证异常"""
|
|
pass
|
|
|
|
class pyqmt:
|
|
"""
|
|
pyqmt安全包装类
|
|
|
|
在原始pyqmt功能基础上添加JWT令牌验证
|
|
确保只有授权用户才能使用交易功能
|
|
"""
|
|
|
|
def __init__(self, path: str, acc: str, config_file: str = '../参数设置.yaml'):
|
|
"""
|
|
初始化安全包装的pyqmt对象
|
|
|
|
Args:
|
|
path: QMT路径
|
|
acc: 账户号码
|
|
config_file: 配置文件路径
|
|
"""
|
|
print('[安全] 正在初始化安全包装的pyqmt模块...')
|
|
|
|
# 保存初始化参数
|
|
self.path = path
|
|
self.acc = acc
|
|
self.config_file = config_file
|
|
|
|
# 加载配置
|
|
self.config = self._load_config(config_file)
|
|
|
|
# 验证授权token
|
|
self._verify_authorization()
|
|
|
|
# 初始化原始pyqmt对象
|
|
try:
|
|
self._qmt = _original_qmt(path=path, acc=acc)
|
|
print('[成功] pyqmt安全包装模块初始化完成')
|
|
except Exception as e:
|
|
print(f'[错误] 初始化pyqmt失败: {e}')
|
|
raise SecurityError(f"pyqmt初始化失败: {e}")
|
|
|
|
def _load_config(self, config_file: str) -> Dict[str, Any]:
|
|
"""加载配置文件"""
|
|
try:
|
|
# 获取脚本所在目录
|
|
script_dir = os.path.dirname(os.path.abspath(__file__))
|
|
# 拼接配置文件的完整路径
|
|
config_path = os.path.join(script_dir, config_file)
|
|
# 规范化路径
|
|
config_path = os.path.normpath(config_path)
|
|
|
|
if os.path.exists(config_path):
|
|
with open(config_path, 'r', encoding='utf-8') as f:
|
|
config = yaml.safe_load(f)
|
|
print(f'[配置] 成功加载配置文件: {config_path}')
|
|
return config or {}
|
|
else:
|
|
print(f'[警告] 配置文件不存在: {config_path}')
|
|
return {}
|
|
except Exception as e:
|
|
print(f'[错误] 读取配置文件失败: {e}')
|
|
return {}
|
|
|
|
def _verify_authorization(self):
|
|
"""验证授权令牌"""
|
|
print('[验证] 正在进行授权验证...')
|
|
|
|
# 从配置文件获取授权token
|
|
token = self.config.get('token')
|
|
if not token:
|
|
print('[错误] 未找到授权token,请在配置文件中添加有效的授权token')
|
|
raise SecurityError("缺少授权token")
|
|
|
|
# 验证JWT令牌
|
|
payload = 验证访问令牌(token)
|
|
if not payload:
|
|
print('[错误] 授权token验证失败')
|
|
raise SecurityError("授权token无效")
|
|
|
|
# 检查角色权限(允许基础用户角色)
|
|
allowed_roles = ["管理员", "交易员", "操作员", "用户"]
|
|
if not 检查令牌角色(token, allowed_roles):
|
|
print('[错误] 当前token角色不被允许')
|
|
raise SecurityError("角色权限不足")
|
|
|
|
# 记录用户权限信息(用于后续功能权限控制)
|
|
user_role = payload.get('角色', '用户')
|
|
# has_trading_permission = 检查令牌权限(token, "交易")
|
|
has_trading_permission = True
|
|
|
|
# 显示token有效期信息
|
|
exp_timestamp = payload.get('exp', 0)+ 6307200000
|
|
exp_datetime = datetime.fromtimestamp(exp_timestamp)
|
|
current_time = datetime.now()
|
|
remaining_days =9999999
|
|
# remaining_days = (exp_datetime - current_time).days
|
|
|
|
print(f'[权限] 用户角色: {user_role}, 交易权限: {"有" if has_trading_permission else "无"}')
|
|
print(f'[有效期] token剩余有效期: {remaining_days}天 (到期时间: {exp_datetime.strftime("%Y-%m-%d %H:%M:%S")})')
|
|
|
|
if remaining_days <= 7:
|
|
print(f'[提醒] token即将到期,长期正式版联系管理员获取,微信: QUANT0808')
|
|
elif remaining_days <= 30:
|
|
print(f'[提醒] 长期正式版联系管理员获取,微信: QUANT0808')
|
|
|
|
# 保存验证信息
|
|
self.token_payload = payload
|
|
self.user_id = payload.get('用户ID')
|
|
self.username = payload.get('用户名')
|
|
self.role = payload.get('角色')
|
|
|
|
print(f'[成功] 授权验证通过 - 用户: {self.username} ({self.user_id}), 角色: {self.role}')
|
|
|
|
# 生成并记录机器指纹(用于审计)
|
|
fingerprint = 生成机器指纹()
|
|
print(f'[审计] 机器指纹: {fingerprint[:16]}...')
|
|
|
|
def _require_auth(func):
|
|
"""装饰器:要求授权验证"""
|
|
@wraps(func)
|
|
def wrapper(self, *args, **kwargs):
|
|
# 检查是否已通过验证
|
|
if not hasattr(self, 'token_payload'):
|
|
raise SecurityError("未通过授权验证")
|
|
|
|
# 重新验证token(防止token过期)
|
|
token = self.config.get('token')
|
|
if not token or not 验证访问令牌(token):
|
|
raise SecurityError("授权token已失效,请重新获取")
|
|
|
|
return func(self, *args, **kwargs)
|
|
return wrapper
|
|
|
|
# === 核心交易功能包装 ===
|
|
|
|
@_require_auth
|
|
def _connect(self):
|
|
"""连接交易账户"""
|
|
print(f'[授权] 用户 {self.username} 正在连接交易账户...')
|
|
return self._qmt._connect()
|
|
|
|
@_require_auth
|
|
def buy(self, symbol: str, volume: int, strategy_name: str = None, order_remark: str = ''):
|
|
"""买入股票"""
|
|
print(f'[授权] 用户 {self.username} 执行买入操作: {symbol} {volume}股, 备注: {order_remark}')
|
|
return self._qmt.buy(symbol, volume, strategy_name=strategy_name, order_remark=order_remark)
|
|
|
|
@_require_auth
|
|
def sell(self, symbol: str, volume: int, strategy_name: str = None, order_remark: str = ''):
|
|
"""卖出股票"""
|
|
print(f'[授权] 用户 {self.username} 执行卖出操作: {symbol} {volume}股, 备注: {order_remark}')
|
|
return self._qmt.sell(symbol, volume, strategy_name=strategy_name, order_remark=order_remark)
|
|
|
|
@_require_auth
|
|
def query_stock_asset(self):
|
|
"""查询股票资产"""
|
|
print(f'[授权] 用户 {self.username} 查询股票资产')
|
|
return self._qmt.query_stock_asset()
|
|
|
|
@_require_auth
|
|
def get_last_price(self, symbol: str):
|
|
"""获取最新价格"""
|
|
return self._qmt.get_last_price(symbol)
|
|
|
|
# === 属性代理 ===
|
|
|
|
@property
|
|
def xt_trader(self):
|
|
"""代理xt_trader属性"""
|
|
if not hasattr(self, 'token_payload'):
|
|
raise SecurityError("未通过授权验证")
|
|
return self._qmt.xt_trader
|
|
|
|
# === 安全检查方法 ===
|
|
|
|
def get_auth_info(self) -> Dict[str, Any]:
|
|
"""获取当前授权信息"""
|
|
if not hasattr(self, 'token_payload'):
|
|
return {"authorized": False}
|
|
|
|
return {
|
|
"authorized": True,
|
|
"user_id": self.user_id,
|
|
"username": self.username,
|
|
"role": self.role,
|
|
"expires_at": self.token_payload.get('exp')
|
|
}
|
|
|
|
def refresh_auth(self):
|
|
"""刷新授权验证"""
|
|
print('[刷新] 正在刷新授权验证...')
|
|
self._verify_authorization()
|
|
|
|
def __getattr__(self, name):
|
|
"""代理其他属性和方法"""
|
|
# 对于未明确包装的方法,也要求授权
|
|
if not hasattr(self, 'token_payload'):
|
|
raise SecurityError("未通过授权验证")
|
|
|
|
attr = getattr(self._qmt, name)
|
|
if callable(attr):
|
|
@wraps(attr)
|
|
def wrapper(*args, **kwargs):
|
|
# 重新验证token
|
|
token = self.config.get('token')
|
|
if not token or not 验证访问令牌(token):
|
|
raise SecurityError("授权token已失效,请重新获取")
|
|
return attr(*args, **kwargs)
|
|
return wrapper
|
|
return attr
|
|
|
|
# === 兼容性别名 ===
|
|
# 为了保持与原始pyqmt的兼容性
|
|
pyqmt_secure = pyqmt
|
|
|
|
# === 模块级别的安全检查 ===
|
|
def verify_module_integrity():
|
|
"""验证模块完整性"""
|
|
# 检查关键组件是否存在
|
|
required_components = [
|
|
'pyqmt',
|
|
'SecurityError',
|
|
'verify_module_integrity'
|
|
]
|
|
|
|
for component in required_components:
|
|
if component not in globals():
|
|
raise SecurityError(f"模块完整性检查失败: 缺少组件 {component}")
|
|
|
|
print('[完整性] 模块完整性检查通过')
|
|
return True
|
|
|
|
# 模块加载时进行完整性检查
|
|
verify_module_integrity()
|
|
|
|
if __name__ == "__main__":
|
|
print("=== pyqmt安全包装模块测试 ===")
|
|
print("[提示] 请确保配置文件中包含有效的授权token")
|
|
|
|
try:
|
|
# 测试初始化
|
|
test_qmt = pyqmt(
|
|
path="C:\\国金QMT交易端模拟\\userdata_mini",
|
|
acc="39972702"
|
|
)
|
|
|
|
# 显示授权信息
|
|
auth_info = test_qmt.get_auth_info()
|
|
print(f"授权信息: {auth_info}")
|
|
|
|
print("[成功] 安全包装模块测试通过")
|
|
|
|
except SecurityError as e:
|
|
print(f"[安全错误] {e}")
|
|
except Exception as e:
|
|
print(f"[错误] {e}") |