dsProject/dsAiTeachingModel/utils/CookieUtil.py

from fastapi.responses import Response
from fastapi.requests import Request
import logging


class CookieUtil:
    @staticmethod
    def set_cookie(res: Response, key: str, value: str, max_age: int = 3600, path: str = "/", secure: bool = False,
                   httponly: bool = True):
        """
        设置cookie

        :param res: FastAPI的Response对象
        :param key: cookie的键
        :param value: cookie的值
        :param max_age: cookie的有效时间（秒），默认1小时（3600秒）
        :param path: cookie的路径，默认为根路径("/")
        :param secure: 是否仅通过HTTPS传输，默认False
        :param httponly: 是否仅通过HTTP请求访问，默认True
        """
        res.set_cookie(
            key=key,
            value=value,
            httponly=httponly,
            secure=secure,
            max_age=max_age,
            path=path
        )
        logging.info(f"设置 cookie: {key}={value}")

    @staticmethod
    def get_cookie(req: Request, key: str) -> str:
        """
        从请求中获取cookie的值

        :param req: FastAPI的Request对象
        :param key: cookie的键
        :return: cookie的值，如果未找到则返回None
        """
        token_value = req.cookies.get(key)
        logging.info(f"从 cookie 中获取的 {key}: {token_value}")
        return token_value

    @staticmethod
    def remove_cookie(res: Response, key: str, path: str = "/"):
        """
        移除cookie

        :param res: FastAPI的Response对象
        :param key: cookie的键
        :param path: cookie的路径，默认为根路径("/")
        """
        res.delete_cookie(key, path=path)
        logging.info(f"移除 cookie: {key}")


# 示例使用
if __name__ == "__main__":
    # 创建一个示例响应对象
    response = Response(content="示例响应")

    # 设置一个cookie
    CookieUtil.set_cookie(response, key="auth_token", value="your_token_value")

    # 创建一个示例请求对象（通常从FastAPI请求中获取）
    request = Request(scope={
        "type": "http",
        "headers": [
            (b"cookie", b"auth_token=your_token_value")
        ]
    })

    # 获取cookie
    token = CookieUtil.get_cookie(request, key="auth_token")
    print(token)

    # 移除cookie
    CookieUtil.remove_cookie(response, key="auth_token")

    # 打印示例响应头
    print(response.headers)