#!/usr/bin/env python

"""
Copyright (c) 2006-2024 sqlmap developers (https://sqlmap.org/)
See the file 'LICENSE' for copying permission
"""#!/usr/bin/env python

"""
Copyright (c) 2006-2024 sqlmap developers (https://sqlmap.org/)
See the file 'LICENSE' for copying permission
"""

import threading

from lib.core.data import logger
from lib.core.enums import CUSTOM_LOGGING 
from lib.core.enums import TIMEOUT_STATE

def timeout(func, args=None, kwargs=None, duration=1, default=None):
    """
    带超时控制的函数执行装饰器
    参数说明:
    func: 要执行的目标函数
    args: 传递给目标函数的位置参数,默认为None
    kwargs: 传递给目标函数的关键字参数,默认为None  
    duration: 超时时间,单位为秒,默认为1秒
    default: 超时后的默认返回值,默认为None
    """
    
    class InterruptableThread(threading.Thread):
        """
        可中断的线程类,继承自threading.Thread
        用于在独立线程中执行目标函数
        """
        def __init__(self):
            threading.Thread.__init__(self)
            self.result = None  # 存储函数执行结果
            self.timeout_state = None  # 存储执行状态
            
        def run(self):
            """
            线程执行的主要逻辑:
            1. 尝试执行目标函数并保存结果
            2. 如果执行成功,设置状态为NORMAL
            3. 如果发生异常,记录日志并设置状态为EXCEPTION
            """
            try:
                # 执行目标函数,处理参数
                self.result = func(*(args or ()), **(kwargs or {}))
                self.timeout_state = TIMEOUT_STATE.NORMAL
            except Exception as ex:
                # 异常处理:记录日志,返回默认值
                logger.log(CUSTOM_LOGGING.TRAFFIC_IN, ex)
                self.result = default
                self.timeout_state = TIMEOUT_STATE.EXCEPTION

    # 创建并启动线程
    thread = InterruptableThread()
    thread.start()
    # 等待线程执行,最多等待duration秒
    thread.join(duration)

    # 判断线程是否还在运行
    if thread.is_alive():
        # 如果超时,返回默认值和超时状态
        return default, TIMEOUT_STATE.TIMEOUT
    else:
        # 如果执行完成,返回执行结果和执行状态
        return thread.result, thread.timeout_state


import threading

from lib.core.data import logger
from lib.core.enums import CUSTOM_LOGGING
from lib.core.enums import TIMEOUT_STATE

def timeout(func, args=None, kwargs=None, duration=1, default=None):
    class InterruptableThread(threading.Thread):
        def __init__(self):
            threading.Thread.__init__(self)
            self.result = None
            self.timeout_state = None

        def run(self):
            try:
                self.result = func(*(args or ()), **(kwargs or {}))
                self.timeout_state = TIMEOUT_STATE.NORMAL
            except Exception as ex:
                logger.log(CUSTOM_LOGGING.TRAFFIC_IN, ex)
                self.result = default
                self.timeout_state = TIMEOUT_STATE.EXCEPTION

    thread = InterruptableThread()
    thread.start()
    thread.join(duration)

    if thread.is_alive():
        return default, TIMEOUT_STATE.TIMEOUT
    else:
        return thread.result, thread.timeout_state
