#!/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}")