Add Infinity Portal Application mismatch detection and fix request trace headers #10

Author: chkp-yuvalpo

.flake8

@@ -0,0 +1,7 @@
+[flake8]
+max-line-length = 150
+exclude = .git,__pycache__,venv
+
+per-file-ignores =
+    # imported but unused
+    __init__.py: F401

chkp_harmony_endpoint_management_sdk/core/session_manager.py

@@ -1,7 +1,7 @@
 import json
 import threading
 import time
-from typing import Dict, Optional, Any, Callable
+from typing import Any
 from unitsnet_py import Duration
 from chkp_harmony_endpoint_management_sdk.classes import harmony_endpoint_saas_options
 from chkp_harmony_endpoint_management_sdk.classes.harmony_endpoint_saas_options import HarmonyEndpointSaaSOptions
@@ -18,6 +18,7 @@
 import uuid
 from enum import Enum
 from urllib.parse import urlparse
+import jwt
 import requests
 
 class WorkMode(Enum):
@@ -32,6 +33,8 @@ class WorkMode(Enum):
 
 SOURCE_HEADER = 'harmony-endpoint-py-sdk'
 
+VERIFY_CONTENT = False
+
 class SessionManager:
     def __init__(self):
         self.__session_operations: SessionOperations = None
@@ -47,7 +50,7 @@ def __init__(self):
         The CI token expiration
         """
 
-        self.__on_premise_portal_auth = OnPremisePortalAuth = None
+        self.__on_premise_portal_auth: OnPremisePortalAuth = None
         """
         The CI token expiration
         """
@@ -189,9 +192,10 @@ def __perform_ci_login(self):
             if not response_json['success']:
                 error_logger(f'Failed to login to CI GW for session "{self.__session_id}" url "{auth_url}", error payload: {response_json}')
                 raise response_json
-            logger(f'Preforming CI login to session id "{self.__session_id}" succeeded');
+            logger(f'Preforming CI login to session id "{self.__session_id}" succeeded')
 
             self.__infinity_portal_token = response_json['data']['token']
+            self.__assert_token_is_for_correct_application(self.__infinity_portal_token)
             self.__next_ci_expiration = Duration.from_seconds(time.time()) + Duration.from_seconds(response_json['data']['expiresIn'])
         except Exception as e:
             error_logger(f'Failed to login to CI GW for session "{self.__session_id}" url "{auth_url}", error: {e}')
@@ -367,9 +371,42 @@ def __validate_premise_params(self, on_premise_portal_auth: OnPremisePortalAuth)
                 message=message,
                 error_scope=HarmonyErrorScope.INVALID_PARAMS,
             )
-    
+
+    def __assert_token_is_for_correct_application(self, bearer_token: str) -> None:
+        if not bearer_token:
+            error_logger('No bearer token was given. Ignoring. Requests may fail')
+            return
+
+        try:
+            # Token verification is NOT required here - this is just to determine whether
+            # the Infinity Portal bearer is for the 'Endpoint' application
+            decoded_token = jwt.decode(bearer_token, verify=VERIFY_CONTENT)
+            if not decoded_token:
+                error_logger('Bearer decoding yielded nothing. Ignoring. Requests may fail')
+                return
+
+            app_id = decoded_token['appId']
+            if not app_id:
+                error_logger('An Application ID claim was not present in the bearer token. Ignoring. Requests may fail')
+                return
+
+            endpoint_app_id = '12345678-8888-1234-1234-123456789123'
+            if app_id != endpoint_app_id:
+                error_logger(f"Target application is incorrect - expected '{endpoint_app_id}' but got '{app_id}'. Raising an error")
+                raise HarmonyApiException(
+                    message=(
+                            "The provided API key must be for the 'Endpoint' service. Please refer to the documentation at "
+                            'https://app.swaggerhub.com/apis/Check-Point/web-mgmt-external-api-production for more details'
+                    ),
+                    error_scope=HarmonyErrorScope.INVALID_PARAMS,
+                )
+
+        except jwt.InvalidTokenError:
+            error_logger('The given token could not be decoded. Ignoring. Requests may fail')
+            return
+
     def connect_cloud(self, infinity_portal_auth: InfinityPortalAuth, session_operations: SessionOperations):
-        self.__work_mode  = WorkMode.CLOUD
+        self.__work_mode = WorkMode.CLOUD
         self.__sdk_connection_state = SDKConnectionState.CONNECTING
         self.__validate_cloud_params(infinity_portal_auth)
         self.__session_operations = session_operations
@@ -383,9 +420,8 @@ def connect_cloud(self, infinity_portal_auth: InfinityPortalAuth, session_operat
 
         self.__activate_keep_alive()
 
-
     def connect_saas(self, infinity_portal_auth: InfinityPortalAuth, harmony_endpoint_saas_options: HarmonyEndpointSaaSOptions, session_operations: SessionOperations):
-        self.__work_mode  = WorkMode.SAAS
+        self.__work_mode = WorkMode.SAAS
         self.__sdk_connection_state = SDKConnectionState.CONNECTING
         self.__validate_cloud_params(infinity_portal_auth)
         self.__harmony_endpoint_saas_options = harmony_endpoint_saas_options

requirements.txt

@@ -8,5 +8,6 @@ python-dateutil==2.8.2
 python-dotenv==1.0.0
 requests==2.31.0
 typing-extensions==4.8.0
+pyjwt==2.8.0
 unitsnet-py>=0.1.82
 urllib3==2.0.7

sdk_generator/templates/endpoint.handlebars

@@ -349,7 +349,7 @@ class BaseApi(api_client.Api):
         if prefix_separator_iterator is None:
                 prefix_separator_iterator = parameter.get_prefix_separator_iterator()
 
-        serialized_data = parameter.serialize('{{operationId}}', prefix_separator_iterator)
+        serialized_data = parameter.serialize('{{operationIdOriginal}}', prefix_separator_iterator)
         for serialized_value in serialized_data.values():
             used_path += serialized_value
 
@@ -407,8 +407,8 @@ class BaseApi(api_client.Api):
         host = self._get_host_oapg('{{operationId}}', _servers, host_index)
     {{/if}}
 
-        _headers.add("x-mgmt-session-id", session_id)
-        _headers.add("x-mgmt-request-id", request_id)
+        _headers.add("x-mgmt-data-session-id", session_id)
+        _headers.add("x-mgmt-data-request-id", request_id)
 
         source_header = self.api_client.configuration.api_key['session']['source_header']
         if source_header:

setup.py

@@ -19,13 +19,14 @@
     'python-dotenv==1.0.0',
     'requests==2.31.0',
     'typing-extensions==4.8.0',
+    'pyjwt==2.8.0',
     'unitsnet-py>=0.1.82',
     'urllib3==2.0.7',
 ]
 
 setup_kwargs = {
     'name': "chkp-harmony-endpoint-management-sdk",
-    'version': '1.1.35',
+    'version': '1.1.37',
     'keywords': 'python, harmony, endpoint, sdk, checkpoint',
     'license': 'MIT',
     'description': 'Harmony Endpoint Official Python SDK',