Lib.common

  1"""
  2工具
  3"""
  4import os.path
  5import shutil
  6import sys
  7import threading
  8import time
  9import traceback
 10import uuid
 11from collections import OrderedDict
 12from io import BytesIO
 13
 14import requests
 15
 16from .constants import *
 17from .utils import Logger
 18
 19logger = Logger.get_logger()
 20
 21
 22class LimitedSizeDict(OrderedDict):
 23    """
 24    带有限制大小的字典
 25    """
 26
 27    def __init__(self, max_size):
 28        self._max_size = max_size
 29        super().__init__()
 30
 31    def __setitem__(self, key, value):
 32        if key in self:
 33            del self[key]
 34        elif len(self) >= self._max_size:
 35            oldest_key = next(iter(self))
 36            del self[oldest_key]
 37        super().__setitem__(key, value)
 38
 39
 40def restart() -> None:
 41    """
 42    MRB2重启
 43    Returns:
 44        None
 45    """
 46    # 获取当前解释器路径
 47    p = sys.executable
 48    try:
 49        # 启动新程序(解释器路径, 当前程序)
 50        os.execl(p, p, *sys.argv)
 51    except OSError:
 52        # 关闭当前程序
 53        sys.exit()
 54
 55
 56def download_file_to_cache(url: str, headers=None, file_name: str = "",
 57                           download_path: str = None, stream=False, fake_headers: bool = True) -> str | None:
 58    """
 59    下载文件到缓存
 60    Args:
 61        url: 下载的url
 62        headers: 下载请求的请求头
 63        file_name: 文件名
 64        download_path: 下载路径
 65        stream: 是否使用流式传输
 66        fake_headers: 是否使用自动生成的假请求头
 67    Returns:
 68        文件路径
 69    """
 70    if headers is None:
 71        headers = {}
 72
 73    if fake_headers:
 74        headers["User-Agent"] = ("Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) "
 75                                 "Chrome/113.0.0.0 Safari/537.36 Edg/113.0.1774.42")
 76        headers["Accept-Language"] = "zh-CN,zh;q=0.9,en;q=0.8,da;q=0.7,ko;q=0.6"
 77        headers["Accept-Encoding"] = "gzip, deflate, br"
 78        headers["Accept"] = ("text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,image/apng,*/*;q=0.8,"
 79                             "application/signed-exchange;v=b3;q=0.7")
 80        headers["Connection"] = "keep-alive"
 81        headers["Upgrade-Insecure-Requests"] = "1"
 82        headers["Cache-Control"] = "max-age=0"
 83        headers["Sec-Fetch-Dest"] = "document"
 84        headers["Sec-Fetch-Mode"] = "navigate"
 85        headers["Sec-Fetch-Site"] = "none"
 86        headers["Sec-Fetch-User"] = "?1"
 87        headers["Sec-Ch-Ua"] = "\"Chromium\";v=\"113\", \"Not-A.Brand\";v=\"24\", \"Microsoft Edge\";v=\"113\""
 88        headers["Sec-Ch-Ua-Mobile"] = "?0"
 89        headers["Sec-Ch-Ua-Platform"] = "\"Windows\""
 90        headers["Host"] = url.split("/")[2]
 91
 92    # 路径拼接
 93    if file_name == "":
 94        file_name = uuid.uuid4().hex + ".cache"
 95
 96    if download_path is None:
 97        file_path = os.path.join(CACHE_PATH, file_name)
 98    else:
 99        file_path = os.path.join(download_path, file_name)
100
101    # 路径不存在特判
102    if not os.path.exists(CACHE_PATH):
103        os.makedirs(CACHE_PATH)
104
105    try:
106        # 下载
107        if stream:
108            with open(file_path, "wb") as f, requests.get(url, stream=True, headers=headers) as res:
109                for chunk in res.iter_content(chunk_size=64 * 1024):
110                    if not chunk:
111                        break
112                    f.write(chunk)
113        else:
114            # 不使用流式传输
115            res = requests.get(url, headers=headers)
116
117            with open(file_path, "wb") as f:
118                f.write(res.content)
119    except requests.exceptions.RequestException as e:
120        logger.warning(f"下载文件失败: {e}")
121        if os.path.exists(file_path):
122            os.remove(file_path)
123        return None
124
125    return file_path
126
127
128# 删除缓存文件
129def clean_cache() -> None:
130    """
131    清理缓存
132    Returns:
133        None
134    """
135    if os.path.exists(CACHE_PATH):
136        try:
137            shutil.rmtree(CACHE_PATH, ignore_errors=True)
138        except Exception as e:
139            logger.warning("删除缓存时报错，报错信息: %s" % repr(e))
140
141
142# 函数缓存
143def function_cache(max_size: int, expiration_time: int = -1):
144    """
145    函数缓存
146    Args:
147        max_size: 最大大小
148        expiration_time: 过期时间
149    Returns:
150        None
151    """
152    cache = LimitedSizeDict(max_size)
153
154    def cache_decorator(func):
155        """
156        缓存装饰器
157        Args:
158            @param func:
159        Returns:
160            None
161        """
162
163        def wrapper(*args, **kwargs):
164            key = str(func.__name__) + str(args) + str(kwargs)
165            if key in cache and (expiration_time == -1 or time.time() - cache[key][1] < expiration_time):
166                return cache[key][0]
167            result = func(*args, **kwargs)
168            cache[key] = (result, time.time())
169            return result
170
171        def clear_cache():
172            """清理缓存"""
173            cache.clear()
174
175        def get_cache():
176            """获取缓存"""
177            return dict(cache)
178
179        def original_func(*args, **kwargs):
180            """调用原函数"""
181            return func(*args, **kwargs)
182
183        wrapper.clear_cache = clear_cache
184        wrapper.get_cache = get_cache
185        wrapper.original_func = original_func
186        return wrapper
187
188    return cache_decorator
189
190
191def thread_lock(func):
192    """
193    线程锁装饰器
194    """
195    thread_lock = threading.Lock()
196
197    def wrapper(*args, **kwargs):
198        with thread_lock:
199            return func(*args, **kwargs)
200
201    return wrapper
202
203
204def finalize_and_cleanup():
205    """
206    结束运行
207    @return:
208    """
209    logger.info("MuRainBot即将关闭，正在删除缓存")
210
211    clean_cache()
212
213    logger.warning("MuRainBot结束运行！")
214    logger.info("再见！\n")
215
216
217@thread_lock
218def save_exc_dump(description: str = None, path: str = None):
219    """
220    保存异常堆栈
221    Args:
222        description: 保存的dump描述，为空则默认
223        path: 保存的路径，为空则自动根据错误生成
224    """
225    try:
226        import coredumpy
227    except ImportError:
228        logger.warning("coredumpy未安装，无法保存异常堆栈")
229        return
230
231    try:
232        exc_type, exc_value, exc_traceback = sys.exc_info()
233        if not exc_traceback:
234            raise Exception("No traceback found")
235
236        # 遍历 traceback 链表，找到最后一个 frame (异常最初发生的位置)
237        current_tb = exc_traceback
238        frame = current_tb.tb_frame
239        while current_tb:
240            frame = current_tb.tb_frame
241            current_tb = current_tb.tb_next
242
243        i = 0
244        while True:
245            if i > 0:
246                path_ = os.path.join(DUMPS_PATH,
247                                     f"coredumpy_"
248                                     f"{time.strftime('%Y%m%d%H%M%S')}_"
249                                     f"{frame.f_code.co_name}_{i}.dump")
250            else:
251                path_ = os.path.join(DUMPS_PATH,
252                                     f"coredumpy_"
253                                     f"{time.strftime('%Y%m%d%H%M%S')}_"
254                                     f"{frame.f_code.co_name}.dump")
255            if not os.path.exists(path_):
256                break
257            i += 1
258
259        for _ in ['?', '*', '"', '<', '>']:
260            path_ = path_.replace(_, "")
261
262        kwargs = {
263            "frame": frame,
264            "path": os.path.normpath(path_)
265        }
266        if description:
267            kwargs["description"] = description
268        if path:
269            kwargs["path"] = path
270
271        coredumpy.dump(**kwargs)
272    except Exception as e:
273        logger.error(f"保存异常堆栈时发生错误: {repr(e)}\n"
274                     f"{traceback.format_exc()}")
275        return None
276
277    return kwargs["path"]
278
279
280def bytes_io_to_file(
281        io_bytes: BytesIO,
282        file_name: str | None = None,
283        file_type: str | None = None,
284        save_dir: str = CACHE_PATH
285):
286    """
287    将BytesIO对象保存成文件，并返回路径
288    Args:
289        io_bytes: BytesIO对象
290        file_name: 要保存的文件名，与file_type选一个填即可
291        file_type: 文件类型(扩展名)，与file_name选一个填即可
292        save_dir: 保存的文件夹
293
294    Returns:
295        保存的文件路径
296    """
297    if not isinstance(io_bytes, BytesIO):
298        raise TypeError("bytes_io_to_file: 输入类型错误")
299    if file_name is None:
300        if file_type is None:
301            file_type = "cache"
302        file_name = uuid.uuid4().hex + "." + file_type
303    if not os.path.exists(save_dir):
304        os.makedirs(save_dir)
305
306    with open(os.path.join(save_dir, file_name), "wb") as f:
307        f.write(io_bytes.getvalue())
308    return os.path.join(save_dir, file_name)