diff --git a/source/app/iris_engine/module_handler/module_handler.py b/source/app/iris_engine/module_handler/module_handler.py index 2e66840c4..b55604e90 100644 --- a/source/app/iris_engine/module_handler/module_handler.py +++ b/source/app/iris_engine/module_handler/module_handler.py @@ -18,6 +18,7 @@ # along with this program; if not, write to the Free Software Foundation, # Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. import traceback +from datetime import datetime, date import base64 import importlib @@ -28,6 +29,27 @@ from app import app from app.blueprints.iris_user import iris_current_user + + +def _serialize_value(obj): + if obj is None: + return None + if isinstance(obj, (datetime, date)): + return obj.isoformat() + if isinstance(obj, dict): + return {k: _serialize_value(v) for k, v in obj.items()} + if isinstance(obj, (list, tuple)): + return [_serialize_value(item) for item in obj] + if hasattr(obj, '__dict__'): + result = {} + for key, value in obj.__dict__.items(): + if key.startswith('_sa_'): + continue + result[key] = _serialize_value(value) + return result + if hasattr(obj, '__iter__'): + return str(obj) + return obj from app.logger import logger from app import celery from app import db @@ -427,13 +449,20 @@ def task_hook_wrapper(self, module_name, hook_name, hook_ui_name, data, init_use """ try: # Data is serialized, so deserialized - signature, pdata = data.encode("utf-8").split(b" ") - is_verified = hmac_verify(signature, pdata) - if is_verified is False: - logger.warning("data argument has not been correctly serialised") - raise Exception('Unable to instantiate target module. Data has not been correctly serialised') - - deser_data = loads(base64.b64decode(pdata)) + # Support both formats: + # 1. Dict (when using auth serializer with security enabled) + # 2. String with HMAC signature (legacy format) + if isinstance(data, dict): + # Auth serializer - data is already a dict + deser_data = data + else: + # Legacy format with HMAC signature + signature, pdata = data.encode("utf-8").split(b" ") + is_verified = hmac_verify(signature, pdata) + if is_verified is False: + logger.warning("data argument has not been correctly serialised") + raise Exception('Unable to instantiate target module. Data has not been correctly serialised') + deser_data = loads(base64.b64decode(pdata)) except Exception as e: logger.exception(e) @@ -487,6 +516,13 @@ def task_hook_wrapper(self, module_name, hook_name, hook_ui_name, data, init_use logger.exception(e) task_status = IStatus.I2Error(message=msg, logs=[traceback.format_exc()], user=init_user, caseid=caseid) + if hasattr(task_status, '__dict__'): + return { + 'code': task_status.code, + 'message': task_status.message, + 'data': _serialize_value(task_status.data), + 'logs': task_status.logs + } return task_status diff --git a/source/app/iris_engine/tasker/celery.py b/source/app/iris_engine/tasker/celery.py index bc353f9d8..d05ff3da7 100644 --- a/source/app/iris_engine/tasker/celery.py +++ b/source/app/iris_engine/tasker/celery.py @@ -16,11 +16,14 @@ # along with this program; if not, write to the Free Software Foundation, # Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. + +import os from celery import Celery from celery.security import setup_security from kombu.serialization import register from app.configuration import CeleryConfig + def _patch_celery_cert_datetime(): import datetime from celery.security.certificate import Certificate @@ -39,25 +42,70 @@ def _patched_has_expired(self): def _register_auth_serializer(): import json + from datetime import datetime, date + + def _serialize_value(obj): + if obj is None: + return None + if isinstance(obj, (datetime, date)): + return obj.isoformat() + if isinstance(obj, dict): + return {k: _serialize_value(v) for k, v in obj.items()} + if isinstance(obj, (list, tuple)): + return [_serialize_value(item) for item in obj] + if hasattr(obj, '__dict__'): + result = {} + for key, value in obj.__dict__.items(): + if key.startswith('_sa_'): + continue + result[key] = _serialize_value(value) + return result + if hasattr(obj, '__iter__'): + return str(obj) + return obj + + class _CeleryJsonEncoder(json.JSONEncoder): + def default(self, obj): + return _serialize_value(obj) def _encode_auth(data): - return json.dumps(data), 'application/auth' + return json.dumps(data, cls=_CeleryJsonEncoder).encode('utf-8'), 'application/auth' def _decode_auth(data): + if isinstance(data, bytes): + data = data.decode('utf-8') return json.loads(data) register('auth', _encode_auth, _decode_auth, content_type='application/auth') -def make_celery(name): - _register_auth_serializer() +def _check_certificate_files(): + key_path = CeleryConfig.security_key + cert_path = CeleryConfig.security_certificate + store_path = CeleryConfig.security_cert_store + + if not key_path or not cert_path or not store_path: + return False + + if not os.path.exists(key_path): + return False + if not os.path.exists(cert_path): + return False + if not os.path.exists(store_path): + return False + + return True + +def make_celery(name): celery_app = Celery( name, config_source=CeleryConfig ) - if CeleryConfig.security_key and CeleryConfig.security_certificate: + _register_auth_serializer() + + if _check_certificate_files(): _patch_celery_cert_datetime() setup_security( allowed_serializers=['auth'], diff --git a/source/app/iris_engine/tasker/tasks.py b/source/app/iris_engine/tasker/tasks.py index b437c54c2..06c00c5ef 100644 --- a/source/app/iris_engine/tasker/tasks.py +++ b/source/app/iris_engine/tasker/tasks.py @@ -32,7 +32,10 @@ @task_prerun.connect def on_task_init(*args, **kwargs): - db.engine.dispose() + try: + db.engine.dispose() + except Exception: + pass def task_case_update(module, pipeline, pipeline_args, caseid): diff --git a/source/app/util.py b/source/app/util.py index 111db8247..029b285f5 100644 --- a/source/app/util.py +++ b/source/app/util.py @@ -79,8 +79,18 @@ def add_obj_history_entry(obj, action, commit=False): return obj + def hmac_sign(data): - key = bytes(current_app.config.get("SECRET_KEY"), "utf-8") + import os + key = os.environ.get("SECRET_KEY") + if not key: + try: + key = current_app.config.get("SECRET_KEY") + except: + pass + if not key: + raise ValueError("SECRET_KEY not available") + key = bytes(key, "utf-8") h = hmac.HMAC(key, hashes.SHA256()) h.update(data) signature = base64.b64encode(h.finalize()) @@ -89,8 +99,17 @@ def hmac_sign(data): def hmac_verify(signature_enc, data): + import os signature = base64.b64decode(signature_enc) - key = bytes(current_app.config.get("SECRET_KEY"), "utf-8") + key = os.environ.get("SECRET_KEY") + if not key: + try: + key = current_app.config.get("SECRET_KEY") + except: + pass + if not key: + return False + key = bytes(key, "utf-8") h = hmac.HMAC(key, hashes.SHA256()) h.update(data)