refactor: add hmac_signed_serializer for celery.

Author: liqiang-fit2cloud

apps/maxkb/settings/lib.py

@@ -71,9 +71,9 @@
 CELERY_result_backend = CELERY_BROKER_URL
 CELERY_timezone = CONFIG.get_time_zone()
 CELERY_ENABLE_UTC = False
-CELERY_task_serializer = 'pickle'
-CELERY_result_serializer = 'pickle'
-CELERY_accept_content = ['json', 'pickle']
+CELERY_task_serializer = 'hmac_signed_serializer'
+CELERY_result_serializer = 'hmac_signed_serializer'
+CELERY_accept_content = ['json', 'hmac_signed_serializer']
 CELERY_RESULT_EXPIRES = 600
 CELERY_WORKER_TASK_LOG_FORMAT = '%(asctime).19s %(message)s'
 CELERY_WORKER_LOG_FORMAT = '%(asctime).19s %(message)s'

apps/ops/celery/__init__.py

@@ -3,14 +3,16 @@
 import os
 
 from celery import Celery
-from celery.schedules import crontab
 from kombu import Exchange, Queue
 from maxkb import settings
 from .heartbeat import *
+from .hmac_signed_serializer import register_hmac_signed_serializer
 
 # set the default Django settings module for the 'celery' program.
 os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'maxkb.settings')
 
+register_hmac_signed_serializer()
+
 app = Celery('MaxKB')
 
 configs = {k: v for k, v in settings.__dict__.items() if k.startswith('CELERY')}

apps/ops/celery/hmac_signed_serializer.py

@@ -0,0 +1,38 @@
+import hmac
+import hashlib
+import pickle
+import os
+import socket
+from kombu.serialization import register
+
+_local_secret_key = os.environ.get('MAXKB_HMAC_SIGNED_SERIALIZER_SECRET_KEY', 'default_hmac_signed_serializer_secret_key:' + os.getenv('MAXKB_VERSION', socket.gethostname()))
+try:
+    from xpack import get_md5
+    _local_secret_key = get_md5()
+except ImportError:
+    pass
+
+def secure_dumps(obj):
+    data = pickle.dumps(obj)
+    signature = hmac.new(_local_secret_key.encode(), data, hashlib.sha256).digest()
+    return signature + data
+
+def secure_loads(signed_data):
+    if len(signed_data) < 32:
+        raise ValueError("Invalid signed data packet")
+    signature = signed_data[:32]
+    payload = signed_data[32:]
+    expected_signature = hmac.new(_local_secret_key.encode(), payload, hashlib.sha256).digest()
+    if hmac.compare_digest(signature, expected_signature):
+        return pickle.loads(payload)
+    else:
+        raise ValueError("Security Alert: Task signature mismatch! Potential tampering detected.")
+
+def register_hmac_signed_serializer():
+    register(
+        'hmac_signed_serializer',
+        secure_dumps,
+        secure_loads,
+        content_type='application/x-spickle',
+        content_encoding='binary'
+    )