DatangDataPy/cms/rz_token.py

import json
import time
from pathlib import Path

import requests

from log import logger

CONFIG = {
    "base_url": "http://10.172.12.201:31000",
    # Token接口
    "token_path": "/api/identityserverService/connect/token",
    "username": "dskj1",
    "client_id": "044fc37a-859d-4063-b897-2e1c93a52074",
    "client_secret": "56CE732DBE97195E2376F79B1D494008",
    "grant_type": "delegation",
    "token_file": "token_cache.json",
    # 资产台账接口
    "asset_api": "/api/openapiService/asset/v2/assetAccount",
}


# ===================== 【2】Token 管理器（复用之前的，自动缓存+刷新） =====================
class TokenManager:
    def __init__(self, config):
        self.config = config
        self.token_file = Path(config["token_file"])
        self.session = requests.Session()

    def _save_token(self, token_data):
        expire_time = int(time.time()) + token_data["expires_in"] - 60
        token_data["expire_time"] = expire_time
        with open(self.token_file, "w", encoding="utf-8") as f:
            json.dump(token_data, f, ensure_ascii=False, indent=2)

    def _load_token_from_file(self):
        if not self.token_file.exists():
            return None
        try:
            with open(self.token_file, "r", encoding="utf-8") as f:
                return json.load(f)
        except:
            return None

    def _is_token_valid(self, token_data):
        if not token_data:
            return False
        required_fields = ["access_token", "expire_time", "token_type"]
        for field in required_fields:
            if field not in token_data:
                return False
        current_time = int(time.time())
        if token_data["expire_time"] <= current_time:
            return False
        return True

    def _get_new_token(self):
        url = self.config["base_url"] + self.config["token_path"]
        data = {
            "username": self.config["username"],
            "client_id": self.config["client_id"],
            "client_secret": self.config["client_secret"],
            "grant_type": self.config["grant_type"]
        }
        headers = {"Content-Type": "application/x-www-form-urlencoded; charset=utf-8"}
        try:
            response = self.session.post(url, data=data, headers=headers, timeout=15)
            response.raise_for_status()
            result = response.json()
            if "access_token" not in result:
                raise Exception(f"获取Token失败：{result}")
            return result
        except Exception as e:
            raise Exception(f"请求Token异常：{str(e)}")

    def get_auth_header(self):
        token_data = self._load_token_from_file()
        if self._is_token_valid(token_data):
            logger.debug("使用缓存Token")
        else:
            logger.debug("Token无效/已过期，重新获取...")
            token_data = self._get_new_token()
            self._save_token(token_data)
            logger.debug("新Token已保存")
        return f"{token_data['token_type']} {token_data['access_token']}"


# ===================== 使用示例 =====================
if __name__ == "__main__":
    # 初始化管理器
    token_manager = TokenManager(CONFIG)

    # 获取认证头
    auth_header = token_manager.get_auth_header()
    logger.info(f"最终认证头： {auth_header}")

    # 调用其他接口时直接使用：
    # headers = {"Authorization": auth_header}
    # requests.get(url, headers=headers)