added authorization for executing main pyscript via callback

kairon/async_callback/router/pyscript_callback.py

@@ -96,7 +96,11 @@ async def execute_async_action_standalone(request: Request, token: str) -> BSRes
     return await process_router_message(token, None, request.method, request)
 
 @router.post("/main_pyscript/execute-python")
-async def trigger_restricted_python(payload: PyscriptPayload):
+async def trigger_restricted_python(
+        request: Request,
+        payload: PyscriptPayload
+):
+    await CallbackAuthenticator.verify(request)
     try:
         result = CallbackUtility.main_pyscript_handler({
             "source_code": payload.source_code,

kairon/shared/actions/utils.py

@@ -25,7 +25,7 @@
 from ..admin.processor import Sysadmin
 from ..cloud.utils import CloudUtility
 from ..constants import KAIRON_USER_MSG_ENTITY, PluginTypes, EventClass
-from ..data.constant import REQUEST_TIMESTAMP_HEADER, TASK_TYPE
+from ..data.constant import REQUEST_TIMESTAMP_HEADER, TASK_TYPE, TOKEN_TYPE
 from ..data.data_objects import Slots, KeyVault
 from ..plugins.factory import PluginFactory
 from ..rest_client import AioRestClient
@@ -691,14 +691,28 @@ def run_pyscript(source_code: Text, context: dict):
             result = lambda_response["Payload"].get('body')
         else:
             callback_url=Utility.environment['async_callback_action']['pyscript']['url']
+            lifespan = Utility.environment['events']['executor']['dynamic_token_lifespan']
+            from ..auth import Authentication
+
+            claims = {"sub": "action-server", "callback": True}
+
+            token = Authentication.create_access_token(
+                data=claims,
+                token_type=TOKEN_TYPE.DYNAMIC.value,
+                token_expire=lifespan
+            )
+            headers = {
+                       "Content-Type": "application/json",
+                       "Authorization": f"Bearer {token}"
+                       }
             resp = Utility.execute_http_request(
                 "POST",
                 http_url=callback_url,
                 request_body={
                     "source_code": source_code,
                     "predefined_objects": context
                 },
-                headers={"Content-Type": "application/json"}
+                headers=headers
             )
             if resp.get('statusCode') != 200:
                 raise ActionFailure(f'Pyscript evaluation failed: {resp}')

tests/integration_test/callback_service_test.py

@@ -661,22 +661,65 @@ async def test_execute_python_success(mock_handler):
     await app.start()
     client = TestClient(app)
 
+    # Create a mock token with three segments
+    mock_token = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJnQUFBQUFCb0xjUi00QnRVM1FFOUdBRFdONVg0N29Yay03VDdxN2dOQks1dlFqM3VnQkoybnZlLVBrVGVTUzZYRm55dkRjM0RmMmNob2JUSFAyM0dDTnI3STZ2eWwyZ0tUYkF3eXFFNERKVDU1TWNsWmV5LWdSRGR4dG55TS01M1R0cjVQLU9iQ3lKTmpDcTFaa1NkRFFQNWk0eDZsN3R1c0E9PSIsInZlcnNpb24iOiIyLjAifQ.HEL3BSt1uO2uRfwwPTVV5UgYvwXsSoyERqZ9Gev2X-U"
+
+    mock_handler.return_value = {"output": "Execution successful", "success": True}
+
     payload = {
         "source_code": "bot_response=100",
         "predefined_objects": {"x": 1}
     }
 
-    # Simulate response from handler
-    mock_handler.return_value = {"output": "Execution successful", "success": True}
+    headers = {
+        "Authorization": f"Bearer {mock_token}"
+    }
 
-    response = await client.post("/main_pyscript/execute-python", content=JSONContent(payload))
+    response = await client.post("/main_pyscript/execute-python", headers=headers, content=JSONContent(payload))
     json_response = await response.json()
     print(json_response)
 
+    # Assertions
     assert response.status == 200
     assert json_response["output"] == "Execution successful"
     assert json_response["success"] is True
 
+@pytest.mark.asyncio
+@patch("kairon.async_callback.utils.CallbackUtility.main_pyscript_handler")
+async def test_execute_python_failure(mock_handler):
+    await app.start()
+    client = TestClient(app)
+
+    # Generate a real token using the create_access_token method
+    data = {"sub": "action-server", "callback": True}
+    token = Authentication.create_access_token(
+        data=data,
+        token_type="dynamic",  # Use the dynamic token type if needed
+        token_expire=30  # Set expiration time
+    )
+
+    # Simulate an exception being raised in the handler
+    mock_handler.side_effect = Exception("Simulated exception")
+
+    payload = {
+        "source_code": "bot_response=100",
+        "predefined_objects": {"x": 1}
+    }
+
+    # Send the request with the real token in the header
+    headers = {
+        "Authorization": f"Bearer {token}"  # Use the real token created above
+    }
+
+    response = await client.post("/main_pyscript/execute-python", headers=headers, content=JSONContent(payload))
+    json_response = await response.json()
+    print(json_response)
+
+    # Assertions to check if the exception handling block is reached
+    assert response.status == 422  # This should match the status code in your error response
+    assert json_response["success"] is False
+    assert "error" in json_response  # Ensure the error message is included in the response
+    assert json_response["error"] == "Simulated exception"  # Check if the simulated exception message is returned
 
 @pytest.mark.asyncio
 @patch("kairon.async_callback.utils.CallbackUtility.main_pyscript_handler")