78 lines
2.4 KiB
Python
78 lines
2.4 KiB
Python
#!/usr/bin/env python
|
||
# -*- coding:utf-8 -*-
|
||
"""
|
||
token 验证
|
||
Created on 2023/12/06
|
||
@author: Monday
|
||
@group : waditu
|
||
@contact:
|
||
"""
|
||
from typing import Optional
|
||
from tushare.util.upass import get_token, set_token
|
||
import requests
|
||
from requests.adapters import HTTPAdapter, Retry
|
||
from tushare.stock import cons as ct
|
||
|
||
|
||
# 自定义异常类，用于权限不足时抛出异常
|
||
class PermissionError(Exception):
|
||
pass
|
||
|
||
|
||
# 装饰器函数，用于进行权限验证
|
||
def require_permission(event_name, event_detail):
|
||
def decorator(func):
|
||
def wrapper(*args, **kwargs):
|
||
# event_name = kwargs["event_name"]
|
||
# event_detail = kwargs["event_detail"]
|
||
# 检查用户权限
|
||
token = get_token()
|
||
if token:
|
||
try:
|
||
r = verify_token(token, event_name, event_detail).json()
|
||
except Exception as err:
|
||
raise PermissionError(f"验证token出错，{err}")
|
||
if r.get("message") == "success":
|
||
# 如果有足够权限，调用原始函数
|
||
return func(*args, **kwargs)
|
||
else:
|
||
print(f"验证token出错，{r.text}")
|
||
raise PermissionError(f"{r['msg']}")
|
||
else:
|
||
raise PermissionError(ct.TOKEN_ERR_MSG)
|
||
|
||
wrapper.__doc__ = func.__doc__
|
||
return wrapper
|
||
|
||
return decorator
|
||
|
||
|
||
def verify_token(token: Optional[str] = None, event_name: Optional[str] = "string",
|
||
event_detail: Optional[str] = "string"):
|
||
session = requests.Session()
|
||
retries = Retry(total=3, backoff_factor=1)
|
||
session.mount('http://', HTTPAdapter(max_retries=retries))
|
||
session.mount('https://', HTTPAdapter(max_retries=retries))
|
||
headers = {
|
||
"accept": "application/json",
|
||
"Content-Type": "application/json"
|
||
}
|
||
# url = "http://api.tushare.pro/dataapi/sdk-event"
|
||
url = ct.verify_token_url
|
||
data = {"user_token": token,
|
||
"event_name": event_name,
|
||
"event_detail": event_detail}
|
||
response = session.post(url, headers=headers, json=data)
|
||
return response
|
||
|
||
|
||
# 使用装饰器进行权限验证
|
||
@require_permission(event_name="xxx", event_detail="xxx")
|
||
def admin_only_function():
|
||
print("ssss")
|
||
|
||
|
||
if __name__ == '__main__':
|
||
set_token("asdasdasa")
|
||
admin_only_function() # 正常情况，输出执行信息
|