feat: remove the using of pickle(bkpaas-auth) (#266)

Author: piglei

sdks/bkpaas-auth/AGENTS.md

@@ -0,0 +1,22 @@
+## Context
+
+You are in the bkpaas-auth repo, helping implement features, fix bugs, and refactor existing code.
+
+## Source code
+
+* bkpaas-auth is a Django app that helps implement user authentication, used by BlueKing systems.
+* The main project in `bkpaas_auth/`.
+* Unit tests are placed in 'tests/' directory, following pytest conventions.
+
+## Coding style
+
+* For Python files, follow PEP-8.
+* For Python files, run `ruff format` to format after edits.
+
+## Common workflows
+
+### Running tests
+
+* Run all tests: `poetry run pytest -s`
+* Run some tests: `poetry run pytest -s tests/filename.py`
+* ALWAYS prefer specifying test files for efficiency.

sdks/bkpaas-auth/bkpaas_auth/backends.py

@@ -1,15 +1,13 @@
 # -*- coding: utf-8 -*-
 import inspect
 import logging
-import pickle
 from typing import Dict, Optional, Union
 
 from django.conf import settings
 from django.contrib.auth import get_user_model
 from django.contrib.auth.models import AnonymousUser
 from django.core.exceptions import ImproperlyConfigured
 from django.http import HttpRequest
-from django.utils.encoding import force_bytes
 
 from bkpaas_auth.conf import bkauth_settings
 from bkpaas_auth.core.constants import ProviderType
@@ -108,11 +106,15 @@ def get_token_from_session(self, request: HttpRequest) -> Optional[LoginToken]:
         if "user_token" not in request.session:
             return None
 
+        raw_user_token = request.session["user_token"]
+        if not isinstance(raw_user_token, str) or not raw_user_token.startswith("{"):
+            logger.warning("ignore legacy or invalid session user_token payload")
+            return None
+
         try:
-            user_token_pickled = force_bytes(request.session["user_token"], "latin1")
-            user_token: LoginToken = pickle.loads(user_token_pickled)
+            user_token: LoginToken = LoginToken.parse_json(raw_user_token)
         except Exception:
-            logger.exception("pickle loads user_token failed")
+            logger.exception("deserialize user_token failed")
             return None
 
         # token 已经过期则不返回，否则会出现 403

sdks/bkpaas-auth/bkpaas_auth/core/token.py

@@ -5,7 +5,7 @@
 import json
 import logging
 from abc import abstractmethod
-from typing import NamedTuple, Optional
+from typing import Any, ClassVar, NamedTuple, Optional, Tuple
 
 from django.utils.timezone import now
 from django.utils.translation import get_language
@@ -23,10 +23,21 @@
 from bkpaas_auth.core.services import get_app_credentials
 from bkpaas_auth.core.user_info import BkUserInfo, RtxUserInfo, UserInfo
 from bkpaas_auth.models import User
-from bkpaas_auth.utils import scrub_data
+from bkpaas_auth.utils import deserialize_datetime, scrub_data, serialize_datetime
 
 logger = logging.getLogger(__name__)
 
+# Constants related with serialization of UserInfo in LoginToken.
+#
+# The field name to store the type of user_info instance.
+_USER_INFO_TYPE_FIELD = "_user_info_type"
+# The mapping between user_info type name and the actual class.
+_USER_INFO_TYPES: dict[str, type[UserInfo]] = {
+    UserInfo.__name__: UserInfo,
+    RtxUserInfo.__name__: RtxUserInfo,
+    BkUserInfo.__name__: BkUserInfo,
+}
+
 
 class UserAccount(NamedTuple):
     """
@@ -181,6 +192,12 @@ class LoginToken:
     """Access token object"""
 
     token_timeout_margin = 300
+    _json_fields: ClassVar[Tuple[str, ...]] = (
+        "login_token",
+        "expires_at",
+        "issued_at",
+        "user_info",
+    )
 
     def __init__(self, login_token=None, expires_in=None):
         assert login_token, "Must provide token string"
@@ -200,6 +217,49 @@ def make_user(self, provider_type):
         self.user_info.provider_type = provider_type
         return create_user_from_token(self)
 
+    def dump_json(self) -> str:
+        """Serialize the token to JSON string."""
+        user_info_type = type(self.user_info).__name__
+        if user_info_type not in _USER_INFO_TYPES:
+            raise TypeError(f"unsupported user info type: {user_info_type}")
+
+        payload = {}
+        for field in self._json_fields:
+            value = getattr(self, field)
+            match field:
+                case "expires_at" | "issued_at":
+                    value = serialize_datetime(value)
+                case "user_info":
+                    value = json.loads(value.dump_json())
+            payload[field] = value
+
+        payload[_USER_INFO_TYPE_FIELD] = user_info_type
+        return json.dumps(payload)
+
+    @classmethod
+    def parse_json(cls, payload: str | dict[str, Any]) -> "LoginToken":
+        """Parse the token from JSON string or dict."""
+        if isinstance(payload, str):
+            payload = json.loads(payload)
+        if not isinstance(payload, dict):
+            raise TypeError(f"serialized payload must be dict, got: {type(payload)!r}")
+
+        user_info_type = payload.get(_USER_INFO_TYPE_FIELD)
+        if user_info_type not in _USER_INFO_TYPES:
+            raise ValueError(f"unexpected serialized type: {user_info_type!r}")
+
+        # Bypass __init__ so the original issued/expires timestamps survive the round trip.
+        token = cls.__new__(cls)
+        for field in cls._json_fields:
+            value = payload[field]
+            match field:
+                case "expires_at" | "issued_at":
+                    value = deserialize_datetime(value)
+                case "user_info":
+                    value = _USER_INFO_TYPES[user_info_type].parse_json(value)
+            setattr(token, field, value)
+        return token
+
 
 def mocked_create_user_from_token(
     token: LoginToken, provider_type: int = ProviderType.RTX, username: str = bkauth_settings.MOCKED_USER_NAME

sdks/bkpaas-auth/bkpaas_auth/core/user_info.py

@@ -1,5 +1,6 @@
 # -*- coding: utf-8 -*-
-from typing import TYPE_CHECKING
+import json
+from typing import TYPE_CHECKING, Any, ClassVar, Tuple
 
 from bkpaas_auth.core.constants import ProviderType
 from bkpaas_auth.core.encoder import user_id_encoder
@@ -9,44 +10,88 @@
 
 
 class UserInfo:
-    """Base class for Userinfo"""
+    """Base class for UserInfo"""
 
     provider_type: ProviderType
+    _json_fields: ClassVar[Tuple[str, ...]] = (
+        "provider_type",
+        "username",
+        "display_name",
+        "time_zone",
+        "tenant_id",
+    )
 
     def __init__(self, username, **kwargs):
         self.username = username
         self.display_name = kwargs.get("display_name") or username
         self.time_zone = kwargs.get("time_zone")
         self.tenant_id = kwargs.get("tenant_id")
 
-    def provide(self, user: 'User'):
+    def provide(self, user: "User"):
         user.provider_type = self.provider_type
         user.username = self.username
         user.bkpaas_user_id = user_id_encoder.encode(self.provider_type, self.username)
 
         user.update_user_info(self.__dict__)
         return user
 
+    def dump_json(self) -> str:
+        payload = {}
+        for field in self._json_fields:
+            value = getattr(self, field, None)
+            if field == "provider_type" and value is not None:
+                value = int(value)
+            payload[field] = value
+        return json.dumps(payload)
+
+    @classmethod
+    def parse_json(cls, payload: str | dict[str, Any]) -> "UserInfo":
+        data = cls._parse_json_payload(payload)
+        user_info = cls.__new__(cls)
+        for field in cls._json_fields:
+            value = data.get(field)
+            if field == "provider_type" and value is not None:
+                value = ProviderType(value)
+            setattr(user_info, field, value)
+        return user_info
+
     def __eq__(self, other):
         if not isinstance(other, UserInfo):
             return False
         return self.username == other.username
 
+    @staticmethod
+    def _parse_json_payload(payload: str | dict[str, Any]) -> dict[str, Any]:
+        """Parse the JSON payload to dict."""
+        if isinstance(payload, str):
+            payload = json.loads(payload)
+        if not isinstance(payload, dict):
+            raise TypeError(f"serialized payload must be dict, got: {type(payload)!r}")
+        return payload
+
 
 class RtxUserInfo(UserInfo):
     """User info for RTX user"""
 
     provider_type = ProviderType.RTX
     email_suffix = "@tencent.com"
 
+    _json_fields = UserInfo._json_fields + (
+        "nickname",
+        "chinese_name",
+        "email",
+        "phone",
+        "avatar_url",
+    )
+
     def __init__(self, **kwargs):
         super().__init__(kwargs["LoginName"], **kwargs)
-        self.nickname = kwargs['ChineseName']
-        self.chinese_name = kwargs['ChineseName']
-        self.email = f'{self.username}{self.email_suffix}'
+        self.nickname = kwargs["ChineseName"]
+        self.chinese_name = kwargs["ChineseName"]
+        self.email = f"{self.username}{self.email_suffix}"
         # 用户 API 添加了限制，没有申请特殊权限的情况下无法获取手机信息
-        self.phone = kwargs.get('MobilePhoneNumber', '')
-        self.avatar_url = ''
+        self.phone = kwargs.get("MobilePhoneNumber", "")
+        self.avatar_url = ""
 
     def __eq__(self, other):
         if not isinstance(other, RtxUserInfo):
@@ -64,16 +109,23 @@ class BkUserInfo(UserInfo):
     """User info for Bk user"""
 
     provider_type = ProviderType.BK
+    _json_fields = UserInfo._json_fields + (
+        "nickname",
+        "chinese_name",
+        "email",
+        "phone",
+        "avatar_url",
+    )
 
     def __init__(self, **kwargs):
         # bk_username 用户英文ID
         super().__init__(kwargs["bk_username"], **kwargs)
         # chname 用户中文名
-        self.nickname = kwargs['chname']
-        self.chinese_name = kwargs['chname']
-        self.email = kwargs['email']
-        self.phone = kwargs['phone']
-        self.avatar_url = ''
+        self.nickname = kwargs["chname"]
+        self.chinese_name = kwargs["chname"]
+        self.email = kwargs["email"]
+        self.phone = kwargs["phone"]
+        self.avatar_url = ""
 
     def __eq__(self, other):
         if not isinstance(other, BkUserInfo):

sdks/bkpaas-auth/bkpaas_auth/middlewares.py

@@ -1,17 +1,15 @@
 # -*- coding: utf-8 -*-
 import json
 import logging
-import pickle
 import time
 from typing import Dict
-from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
 
 from django.conf import settings
 from django.contrib import auth
 from django.http import HttpRequest, HttpResponse
-from django.utils.deprecation import MiddlewareMixin
-from django.utils.encoding import force_str
 from django.utils import timezone as dj_timezone
+from django.utils.deprecation import MiddlewareMixin
+from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
 
 from bkpaas_auth.backends import UniversalAuthBackend
 from bkpaas_auth.core.constants import ACCESS_PERMISSION_DENIED_CODE
@@ -24,7 +22,7 @@ class CookieLoginMiddleware(MiddlewareMixin):
     """Call auth.login when user credential cookies changes"""
 
     def process_request(self, request):
-        assert hasattr(request, 'session'), (
+        assert hasattr(request, "session"), (
             "The CookieLoginMiddleware requires session middleware "
             "to be installed. Edit your MIDDLEWARE%s setting to insert "
             "'django.contrib.sessions.middleware.SessionMiddleware' before "
@@ -44,7 +42,7 @@ def process_request(self, request):
                 self.authenticate_and_login(request, credentials)
             except AccessPermissionDenied as e:
                 resp = HttpResponse(
-                    json.dumps({'code': ACCESS_PERMISSION_DENIED_CODE, 'detail': str(e)}),
+                    json.dumps({"code": ACCESS_PERMISSION_DENIED_CODE, "detail": str(e)}),
                     content_type="application/json",
                 )
                 resp.status_code = 403
@@ -57,7 +55,7 @@ def should_authenticate(
     ) -> bool:
         """Decide whether to re-authenticate current credentials or not"""
         # Force re-login if credentials is different from last time
-        credentials_been_modified = credentials != request.session.get('auth_credentials', {})
+        credentials_been_modified = credentials != request.session.get("auth_credentials", {})
         if credentials_been_modified:
             return True
 
@@ -71,10 +69,10 @@ def authenticate_and_login(self, request: HttpRequest, credentials: Dict[str, st
         :params request: Current request object
         :params credentials: user credentials, such as uin/skey pair
         """
-        logger.debug('Authenticating credentials...')
+        logger.debug("Authenticating credentials...")
         user = auth.authenticate(request=request, auth_credentials=credentials)
         if user is None or not user.is_authenticated:
-            logger.info('Authentication failed, logout.')
+            logger.info("Authentication failed, logout.")
             auth.logout(request)
             return
 
@@ -83,13 +81,12 @@ def authenticate_and_login(self, request: HttpRequest, credentials: Dict[str, st
             logger.info("User is not validate by UniversalAuthBackend, skip login processes.")
             return
 
-        logger.debug('Authentication finished, username: %s', user.username)
-        request.session['provider_type'] = user.provider_type.value
-        request.session['bkpaas_user_id'] = user.bkpaas_user_id
-        request.session['bkpaas_authenticated_at'] = time.time()
-        request.session['auth_credentials'] = credentials
-        # python3 compatibility
-        request.session['user_token'] = force_str(pickle.dumps(user.token), 'latin1')
+        logger.debug("Authentication finished, username: %s", user.username)
+        request.session["provider_type"] = user.provider_type.value
+        request.session["bkpaas_user_id"] = user.bkpaas_user_id
+        request.session["bkpaas_authenticated_at"] = time.time()
+        request.session["auth_credentials"] = credentials
+        request.session["user_token"] = user.token.dump_json()
 
         # Calling `auth.login` will rotate CSRF token and modify user session, only do this when the authenticated
         # user was different with the user stored in session. Otherwise CSRF token validation may fail due to the

sdks/bkpaas-auth/bkpaas_auth/utils.py

@@ -1,3 +1,4 @@
+import datetime
 from typing import Any, Dict
 
 DEFAULT_SCRUBBED_FIELDS = (
@@ -46,3 +47,15 @@ def _key_is_sensitive(key: str) -> bool:
             else:
                 current_result[key] = value
     return result
+
+
+def serialize_datetime(value: datetime.datetime) -> str:
+    if not isinstance(value, datetime.datetime):
+        raise TypeError(f"datetime value required, got: {type(value)!r}")
+    return value.isoformat()
+
+
+def deserialize_datetime(value: str) -> datetime.datetime:
+    if not isinstance(value, str):
+        raise TypeError(f"datetime payload must be str, got: {type(value)!r}")
+    return datetime.datetime.fromisoformat(value)