From 20530fb0d893703119adea0b21291d82d76484ad Mon Sep 17 00:00:00 2001 From: Ilya Siamionau Date: Fri, 12 Apr 2024 00:42:37 +0200 Subject: [PATCH 1/2] CM-34777 - Add correlation ID --- cycode/cli/printers/printer_base.py | 4 ++ cycode/cyclient/cycode_client_base.py | 26 ++--------- cycode/cyclient/headers.py | 46 +++++++++++++++++++ .../cli/exceptions/test_handle_scan_errors.py | 2 +- tests/cyclient/test_client_base.py | 4 +- 5 files changed, 59 insertions(+), 23 deletions(-) create mode 100644 cycode/cyclient/headers.py diff --git a/cycode/cli/printers/printer_base.py b/cycode/cli/printers/printer_base.py index 97a0aff5..cc354082 100644 --- a/cycode/cli/printers/printer_base.py +++ b/cycode/cli/printers/printer_base.py @@ -5,6 +5,7 @@ import click from cycode.cli.models import CliError, CliResult +from cycode.cyclient.headers import get_correlation_id if TYPE_CHECKING: from cycode.cli.models import LocalScanResult @@ -46,3 +47,6 @@ def print_exception(self, e: Optional[BaseException] = None) -> None: message = f'Error: {traceback_message}' click.secho(message, err=True, fg=self.RED_COLOR_NAME) + + correlation_message = f'Correlation ID: {get_correlation_id()}' + click.secho(correlation_message, err=True, fg=self.RED_COLOR_NAME) diff --git a/cycode/cyclient/cycode_client_base.py b/cycode/cyclient/cycode_client_base.py index 6cfcd8c3..a4c5ab63 100644 --- a/cycode/cyclient/cycode_client_base.py +++ b/cycode/cyclient/cycode_client_base.py @@ -1,33 +1,17 @@ -import platform from typing import ClassVar, Dict, Optional from requests import Response, exceptions, request -from cycode import __version__ from cycode.cli.exceptions.custom_exceptions import HttpUnauthorizedError, NetworkError -from cycode.cli.user_settings.configuration_manager import ConfigurationManager from cycode.cyclient import config, logger - - -def get_cli_user_agent() -> str: - """Return base User-Agent of CLI. - - Example: CycodeCLI/0.2.3 (OS: Darwin; Arch: arm64; Python: 3.8.16; InstallID: *uuid4*) - """ - app_name = 'CycodeCLI' - version = __version__ - - os = platform.system() - arch = platform.machine() - python_version = platform.python_version() - - install_id = ConfigurationManager().get_or_create_installation_id() - - return f'{app_name}/{version} (OS: {os}; Arch: {arch}; Python: {python_version}; InstallID: {install_id})' +from cycode.cyclient.headers import get_cli_user_agent, get_correlation_id class CycodeClientBase: - MANDATORY_HEADERS: ClassVar[Dict[str, str]] = {'User-Agent': get_cli_user_agent()} + MANDATORY_HEADERS: ClassVar[Dict[str, str]] = { + 'User-Agent': get_cli_user_agent(), + 'X-Correlation-Id': get_correlation_id(), + } def __init__(self, api_url: str) -> None: self.timeout = config.timeout diff --git a/cycode/cyclient/headers.py b/cycode/cyclient/headers.py new file mode 100644 index 00000000..53eba880 --- /dev/null +++ b/cycode/cyclient/headers.py @@ -0,0 +1,46 @@ +import platform +from typing import Optional +from uuid import uuid4 + +from cycode import __version__ +from cycode.cli.user_settings.configuration_manager import ConfigurationManager +from cycode.cyclient import logger + + +def get_cli_user_agent() -> str: + """Return base User-Agent of CLI. + + Example: CycodeCLI/0.2.3 (OS: Darwin; Arch: arm64; Python: 3.8.16; InstallID: *uuid4*) + """ + app_name = 'CycodeCLI' + version = __version__ + + os = platform.system() + arch = platform.machine() + python_version = platform.python_version() + + install_id = ConfigurationManager().get_or_create_installation_id() + + return f'{app_name}/{version} (OS: {os}; Arch: {arch}; Python: {python_version}; InstallID: {install_id})' + + +class _CorrelationId: + _id: Optional[str] = None + + def get_correlation_id(self) -> str: + """Get correlation ID. + + Notes: + Used across all requests to correlate logs and metrics. + It doesn't depend on client instances. + Lifetime is the same as the process. + """ + if self._id is None: + # example: 16fd2706-8baf-433b-82eb-8c7fada847da + self._id = str(uuid4()) + logger.debug(f'Correlation ID: {self._id}') + + return self._id + + +get_correlation_id = _CorrelationId().get_correlation_id diff --git a/tests/cli/exceptions/test_handle_scan_errors.py b/tests/cli/exceptions/test_handle_scan_errors.py index 7d63802b..d473801f 100644 --- a/tests/cli/exceptions/test_handle_scan_errors.py +++ b/tests/cli/exceptions/test_handle_scan_errors.py @@ -59,7 +59,7 @@ def test_handle_exception_verbose(monkeypatch: 'MonkeyPatch') -> None: ctx = click.Context(click.Command('path'), obj={'verbose': True, 'output': 'text'}) def mock_secho(msg: str, *_, **__) -> None: - assert 'Error:' in msg + assert 'Error:' in msg or 'Correlation ID:' in msg monkeypatch.setattr(click, 'secho', mock_secho) diff --git a/tests/cyclient/test_client_base.py b/tests/cyclient/test_client_base.py index d0b00563..d9e871d1 100644 --- a/tests/cyclient/test_client_base.py +++ b/tests/cyclient/test_client_base.py @@ -1,10 +1,12 @@ from cycode.cyclient import config -from cycode.cyclient.cycode_client_base import CycodeClientBase, get_cli_user_agent +from cycode.cyclient.cycode_client_base import CycodeClientBase +from cycode.cyclient.headers import get_cli_user_agent, get_correlation_id def test_mandatory_headers() -> None: expected_headers = { 'User-Agent': get_cli_user_agent(), + 'X-Correlation-Id': get_correlation_id(), } client = CycodeClientBase(config.cycode_api_url) From e5f66a488a46d7423d21c2dd76b9fc1e740f8da2 Mon Sep 17 00:00:00 2001 From: Ilya Siamionau Date: Fri, 12 Apr 2024 12:33:13 +0200 Subject: [PATCH 2/2] expose some methods for agent --- cycode/cyclient/config.py | 36 +++++++++++++++++++----------------- 1 file changed, 19 insertions(+), 17 deletions(-) diff --git a/cycode/cyclient/config.py b/cycode/cyclient/config.py index da3c6a18..ade271d1 100644 --- a/cycode/cyclient/config.py +++ b/cycode/cyclient/config.py @@ -1,12 +1,12 @@ import logging import os import sys -from typing import Optional +from typing import Optional, Union from urllib.parse import urlparse from cycode.cli import consts from cycode.cli.user_settings.configuration_manager import ConfigurationManager -from cycode.cyclient.config_dev import DEV_MODE_ENV_VAR_NAME, DEV_TENANT_ID_ENV_VAR_NAME +from cycode.cyclient import config_dev def _set_io_encodings() -> None: @@ -37,7 +37,7 @@ def _set_io_encodings() -> None: DEFAULT_CONFIGURATION = { consts.TIMEOUT_ENV_VAR_NAME: 300, consts.LOGGING_LEVEL_ENV_VAR_NAME: logging.INFO, - DEV_MODE_ENV_VAR_NAME: 'False', + config_dev.DEV_MODE_ENV_VAR_NAME: 'false', } configuration = dict(DEFAULT_CONFIGURATION, **os.environ) @@ -45,12 +45,14 @@ def _set_io_encodings() -> None: _CREATED_LOGGERS = set() -def get_logger(logger_name: Optional[str] = None) -> logging.Logger: - config_level = _get_val_as_string(consts.LOGGING_LEVEL_ENV_VAR_NAME) - level = logging.getLevelName(config_level) +def get_logger_level() -> Optional[Union[int, str]]: + config_level = get_val_as_string(consts.LOGGING_LEVEL_ENV_VAR_NAME) + return logging.getLevelName(config_level) + +def get_logger(logger_name: Optional[str] = None) -> logging.Logger: new_logger = logging.getLogger(logger_name) - new_logger.setLevel(level) + new_logger.setLevel(get_logger_level()) _CREATED_LOGGERS.add(new_logger) @@ -62,16 +64,16 @@ def set_logging_level(level: int) -> None: created_logger.setLevel(level) -def _get_val_as_string(key: str) -> str: +def get_val_as_string(key: str) -> str: return configuration.get(key) -def _get_val_as_bool(key: str, default: str = '') -> bool: +def get_val_as_bool(key: str, default: str = '') -> bool: val = configuration.get(key, default) - return val.lower() in ('true', '1') + return val.lower() in {'true', '1'} -def _get_val_as_int(key: str) -> Optional[int]: +def get_val_as_int(key: str) -> Optional[int]: val = configuration.get(key) if val: return int(val) @@ -79,7 +81,7 @@ def _get_val_as_int(key: str) -> Optional[int]: return None -def _is_valid_url(url: str) -> bool: +def is_valid_url(url: str) -> bool: try: urlparse(url) return True @@ -92,12 +94,12 @@ def _is_valid_url(url: str) -> bool: configuration_manager = ConfigurationManager() cycode_api_url = configuration_manager.get_cycode_api_url() -if not _is_valid_url(cycode_api_url): +if not is_valid_url(cycode_api_url): cycode_api_url = consts.DEFAULT_CYCODE_API_URL -timeout = _get_val_as_int(consts.CYCODE_CLI_REQUEST_TIMEOUT_ENV_VAR_NAME) +timeout = get_val_as_int(consts.CYCODE_CLI_REQUEST_TIMEOUT_ENV_VAR_NAME) if not timeout: - timeout = _get_val_as_int(consts.TIMEOUT_ENV_VAR_NAME) + timeout = get_val_as_int(consts.TIMEOUT_ENV_VAR_NAME) -dev_mode = _get_val_as_bool(DEV_MODE_ENV_VAR_NAME) -dev_tenant_id = _get_val_as_string(DEV_TENANT_ID_ENV_VAR_NAME) +dev_mode = get_val_as_bool(config_dev.DEV_MODE_ENV_VAR_NAME) +dev_tenant_id = get_val_as_string(config_dev.DEV_TENANT_ID_ENV_VAR_NAME)