cycode 3.8.10.dev1__py3-none-any.whl → 3.9.0__py3-none-any.whl

cycode/cli/apps/ai_guardrails/scan/scan_command.py

@@ -0,0 +1,134 @@
+"""
+Scan command for AI guardrails.
+
+This command handles AI IDE hooks by reading JSON from stdin and outputting
+a JSON response to stdout. It scans prompts, file reads, and MCP tool calls
+for secrets before they are sent to AI models.
+
+Supports multiple IDEs with different hook event types. The specific hook events
+supported depend on the IDE being used (e.g., Cursor supports beforeSubmitPrompt,
+beforeReadFile, beforeMCPExecution).
+"""
+
+import sys
+from typing import Annotated
+
+import click
+import typer
+
+from cycode.cli.apps.ai_guardrails.scan.handlers import get_handler_for_event
+from cycode.cli.apps.ai_guardrails.scan.payload import AIHookPayload
+from cycode.cli.apps.ai_guardrails.scan.policy import load_policy
+from cycode.cli.apps.ai_guardrails.scan.response_builders import get_response_builder
+from cycode.cli.apps.ai_guardrails.scan.types import AiHookEventType
+from cycode.cli.apps.ai_guardrails.scan.utils import output_json, safe_json_parse
+from cycode.cli.exceptions.custom_exceptions import HttpUnauthorizedError
+from cycode.cli.utils.get_api_client import get_ai_security_manager_client, get_scan_cycode_client
+from cycode.cli.utils.sentry import add_breadcrumb
+from cycode.logger import get_logger
+
+logger = get_logger('AI Guardrails')
+
+
+def _get_auth_error_message(error: Exception) -> str:
+    """Get user-friendly message for authentication errors."""
+    if isinstance(error, click.ClickException):
+        # Missing credentials
+        return f'{error.message} Please run `cycode configure` to set up your credentials.'
+
+    if isinstance(error, HttpUnauthorizedError):
+        # Invalid/expired credentials
+        return (
+            'Unable to authenticate to Cycode. Your credentials are invalid or have expired. '
+            'Please run `cycode configure` to update your credentials.'
+        )
+
+    # Fallback
+    return 'Authentication failed. Please run `cycode configure` to set up your credentials.'
+
+
+def _initialize_clients(ctx: typer.Context) -> None:
+    """Initialize API clients.
+
+    May raise click.ClickException if credentials are missing,
+    or HttpUnauthorizedError if credentials are invalid.
+    """
+    scan_client = get_scan_cycode_client(ctx)
+    ctx.obj['client'] = scan_client
+
+    ai_security_client = get_ai_security_manager_client(ctx)
+    ctx.obj['ai_security_client'] = ai_security_client
+
+
+def scan_command(
+    ctx: typer.Context,
+    ide: Annotated[
+        str,
+        typer.Option(
+            '--ide',
+            help='IDE that sent the payload (e.g., "cursor"). Defaults to cursor.',
+            hidden=True,
+        ),
+    ] = 'cursor',
+) -> None:
+    """Scan content from AI IDE hooks for secrets.
+
+    This command reads a JSON payload from stdin containing hook event data
+    and outputs a JSON response to stdout indicating whether to allow or block the action.
+
+    The hook event type is determined from the event field in the payload (field name
+    varies by IDE). Each IDE may support different hook events for scanning prompts,
+    file access, and tool executions.
+
+    Example usage (from IDE hooks configuration):
+        { "command": "cycode ai-guardrails scan" }
+    """
+    add_breadcrumb('ai-guardrails-scan')
+
+    stdin_data = sys.stdin.read().strip()
+    payload = safe_json_parse(stdin_data)
+
+    tool = ide.lower()
+    response_builder = get_response_builder(tool)
+
+    if not payload:
+        logger.debug('Empty or invalid JSON payload received')
+        output_json(response_builder.allow_prompt())
+        return
+
+    unified_payload = AIHookPayload.from_payload(payload, tool=tool)
+    event_name = unified_payload.event_name
+    logger.debug('Processing AI guardrails hook', extra={'event_name': event_name, 'tool': tool})
+
+    workspace_roots = payload.get('workspace_roots', ['.'])
+    policy = load_policy(workspace_roots[0])
+
+    try:
+        _initialize_clients(ctx)
+
+        handler = get_handler_for_event(event_name)
+        if handler is None:
+            logger.debug('Unknown hook event, allowing by default', extra={'event_name': event_name})
+            output_json(response_builder.allow_prompt())
+            return
+
+        response = handler(ctx, unified_payload, policy)
+        logger.debug('Hook handler completed', extra={'event_name': event_name, 'response': response})
+        output_json(response)
+
+    except (click.ClickException, HttpUnauthorizedError) as e:
+        error_message = _get_auth_error_message(e)
+        if event_name == AiHookEventType.PROMPT:
+            output_json(response_builder.deny_prompt(error_message))
+            return
+        output_json(response_builder.deny_permission(error_message, 'Authentication required'))
+
+    except Exception as e:
+        logger.error('Hook handler failed', exc_info=e)
+        if policy.get('fail_open', True):
+            output_json(response_builder.allow_prompt())
+            return
+        if event_name == AiHookEventType.PROMPT:
+            output_json(response_builder.deny_prompt('Cycode guardrails error - blocking due to fail-closed policy'))
+            return
+        output_json(response_builder.deny_permission('Cycode guardrails error', 'Blocking due to fail-closed policy'))

cycode/cli/apps/ai_guardrails/scan/types.py

@@ -0,0 +1,54 @@
+"""Type definitions for AI guardrails."""
+
+import sys
+
+if sys.version_info >= (3, 11):
+    from enum import StrEnum
+else:
+    from enum import Enum
+
+    class StrEnum(str, Enum):
+        def __str__(self) -> str:
+            return self.value
+
+
+class AiHookEventType(StrEnum):
+    """Canonical event types for AI guardrails.
+
+    These are IDE-agnostic event types. Each IDE's specific event names
+    are mapped to these canonical types using the mapping dictionaries below.
+    """
+
+    PROMPT = 'Prompt'
+    FILE_READ = 'FileRead'
+    MCP_EXECUTION = 'McpExecution'
+
+
+# IDE-specific event name mappings to canonical types
+CURSOR_EVENT_MAPPING = {
+    'beforeSubmitPrompt': AiHookEventType.PROMPT,
+    'beforeReadFile': AiHookEventType.FILE_READ,
+    'beforeMCPExecution': AiHookEventType.MCP_EXECUTION,
+}
+
+
+class AIHookOutcome(StrEnum):
+    """Outcome of an AI hook event evaluation."""
+
+    ALLOWED = 'allowed'
+    BLOCKED = 'blocked'
+    WARNED = 'warned'
+
+
+class BlockReason(StrEnum):
+    """Reason why an AI hook event was blocked.
+
+    These are categorical reasons sent to the backend for tracking/analytics,
+    separate from the detailed user-facing messages.
+    """
+
+    SECRETS_IN_PROMPT = 'secrets_in_prompt'
+    SECRETS_IN_FILE = 'secrets_in_file'
+    SECRETS_IN_MCP_ARGS = 'secrets_in_mcp_args'
+    SENSITIVE_PATH = 'sensitive_path'
+    SCAN_FAILURE = 'scan_failure'

cycode/cli/apps/ai_guardrails/scan/utils.py

@@ -0,0 +1,72 @@
+"""
+Utility functions for AI guardrails.
+
+Includes JSON parsing, path matching, and text handling utilities.
+"""
+
+import json
+import os
+from pathlib import Path
+
+from cycode.cli.apps.ai_guardrails.scan.policy import get_policy_value
+
+
+def safe_json_parse(s: str) -> dict:
+    """Parse JSON string, returning empty dict on failure."""
+    try:
+        return json.loads(s) if s else {}
+    except (json.JSONDecodeError, TypeError):
+        return {}
+
+
+def truncate_utf8(text: str, max_bytes: int) -> str:
+    """Truncate text to max bytes while preserving valid UTF-8."""
+    if not text:
+        return ''
+    encoded = text.encode('utf-8')
+    if len(encoded) <= max_bytes:
+        return text
+    return encoded[:max_bytes].decode('utf-8', errors='ignore')
+
+
+def normalize_path(file_path: str) -> str:
+    """Normalize path to prevent traversal attacks."""
+    if not file_path:
+        return ''
+    normalized = os.path.normpath(file_path)
+    # Reject paths that attempt to escape outside bounds
+    if normalized.startswith('..'):
+        return ''
+    return normalized
+
+
+def matches_glob(file_path: str, pattern: str) -> bool:
+    """Check if file path matches a glob pattern.
+
+    Case-insensitive matching for cross-platform compatibility.
+    """
+    normalized = normalize_path(file_path)
+    if not normalized or not pattern:
+        return False
+
+    path = Path(normalized)
+    # Try case-sensitive first
+    if path.match(pattern):
+        return True
+
+    # Then try case-insensitive by lowercasing both path and pattern
+    path_lower = Path(normalized.lower())
+    return path_lower.match(pattern.lower())
+
+
+def is_denied_path(file_path: str, policy: dict) -> bool:
+    """Check if file path is in the denylist."""
+    if not file_path:
+        return False
+    globs = get_policy_value(policy, 'file_read', 'deny_globs', default=[])
+    return any(matches_glob(file_path, g) for g in globs)
+
+
+def output_json(obj: dict) -> None:
+    """Write JSON response to stdout (for IDE to read)."""
+    print(json.dumps(obj), end='')  # noqa: T201

cycode/cli/apps/ai_guardrails/status_command.py

@@ -0,0 +1,92 @@
+"""Status command for AI guardrails hooks."""
+
+import os
+from pathlib import Path
+from typing import Annotated, Optional
+
+import typer
+from rich.table import Table
+
+from cycode.cli.apps.ai_guardrails.command_utils import console, validate_and_parse_ide, validate_scope
+from cycode.cli.apps.ai_guardrails.hooks_manager import get_hooks_status
+from cycode.cli.utils.sentry import add_breadcrumb
+
+
+def status_command(
+    ctx: typer.Context,
+    scope: Annotated[
+        str,
+        typer.Option(
+            '--scope',
+            '-s',
+            help='Check scope: "user", "repo", or "all" for both.',
+        ),
+    ] = 'all',
+    ide: Annotated[
+        str,
+        typer.Option(
+            '--ide',
+            help='IDE to check status for (e.g., "cursor"). Defaults to cursor.',
+        ),
+    ] = 'cursor',
+    repo_path: Annotated[
+        Optional[Path],
+        typer.Option(
+            '--repo-path',
+            help='Repository path for repo-scoped status (defaults to current directory).',
+            exists=True,
+            file_okay=False,
+            dir_okay=True,
+            resolve_path=True,
+        ),
+    ] = None,
+) -> None:
+    """Show AI guardrails hook installation status.
+
+    Displays the current status of Cycode AI guardrails hooks for the specified IDE.
+
+    Examples:
+        cycode ai-guardrails status                # Show both user and repo status
+        cycode ai-guardrails status --scope user   # Show only user-level status
+        cycode ai-guardrails status --scope repo   # Show only repo-level status
+        cycode ai-guardrails status --ide cursor   # Check status for Cursor IDE
+    """
+    add_breadcrumb('ai-guardrails-status')
+
+    # Validate inputs (status allows 'all' scope)
+    validate_scope(scope, allowed_scopes=('user', 'repo', 'all'))
+    if repo_path is None:
+        repo_path = Path(os.getcwd())
+    ide_type = validate_and_parse_ide(ide)
+
+    scopes_to_check = ['user', 'repo'] if scope == 'all' else [scope]
+
+    for check_scope in scopes_to_check:
+        status = get_hooks_status(check_scope, repo_path if check_scope == 'repo' else None, ide=ide_type)
+
+        console.print()
+        console.print(f'[bold]{check_scope.upper()} SCOPE[/]')
+        console.print(f'Path: {status["hooks_path"]}')
+
+        if not status['file_exists']:
+            console.print('[dim]No hooks.json file found[/]')
+            continue
+
+        if status['cycode_installed']:
+            console.print('[green]✓ Cycode AI guardrails: INSTALLED[/]')
+        else:
+            console.print('[yellow]○ Cycode AI guardrails: NOT INSTALLED[/]')
+
+        # Show hook details
+        table = Table(show_header=True, header_style='bold')
+        table.add_column('Hook Event')
+        table.add_column('Cycode Enabled')
+        table.add_column('Total Hooks')
+
+        for event, info in status['hooks'].items():
+            enabled = '[green]Yes[/]' if info['enabled'] else '[dim]No[/]'
+            table.add_row(event, enabled, str(info['total_entries']))
+
+        console.print(table)
+
+    console.print()

cycode/cli/apps/ai_guardrails/uninstall_command.py

@@ -0,0 +1,73 @@
+"""Uninstall command for AI guardrails hooks."""
+
+from pathlib import Path
+from typing import Annotated, Optional
+
+import typer
+
+from cycode.cli.apps.ai_guardrails.command_utils import (
+    console,
+    resolve_repo_path,
+    validate_and_parse_ide,
+    validate_scope,
+)
+from cycode.cli.apps.ai_guardrails.consts import IDE_CONFIGS
+from cycode.cli.apps.ai_guardrails.hooks_manager import uninstall_hooks
+from cycode.cli.utils.sentry import add_breadcrumb
+
+
+def uninstall_command(
+    ctx: typer.Context,
+    scope: Annotated[
+        str,
+        typer.Option(
+            '--scope',
+            '-s',
+            help='Uninstall scope: "user" for user-level hooks, "repo" for repository-level hooks.',
+        ),
+    ] = 'user',
+    ide: Annotated[
+        str,
+        typer.Option(
+            '--ide',
+            help='IDE to uninstall hooks from (e.g., "cursor"). Defaults to cursor.',
+        ),
+    ] = 'cursor',
+    repo_path: Annotated[
+        Optional[Path],
+        typer.Option(
+            '--repo-path',
+            help='Repository path for repo-scoped uninstallation (defaults to current directory).',
+            exists=True,
+            file_okay=False,
+            dir_okay=True,
+            resolve_path=True,
+        ),
+    ] = None,
+) -> None:
+    """Remove AI guardrails hooks from supported IDEs.
+
+    This command removes Cycode hooks from the IDE's hooks configuration.
+    Other hooks (if any) will be preserved.
+
+    Examples:
+        cycode ai-guardrails uninstall                    # Remove user-level hooks
+        cycode ai-guardrails uninstall --scope repo       # Remove repo-level hooks
+        cycode ai-guardrails uninstall --ide cursor       # Uninstall from Cursor IDE
+    """
+    add_breadcrumb('ai-guardrails-uninstall')
+
+    # Validate inputs
+    validate_scope(scope)
+    repo_path = resolve_repo_path(scope, repo_path)
+    ide_type = validate_and_parse_ide(ide)
+    ide_name = IDE_CONFIGS[ide_type].name
+    success, message = uninstall_hooks(scope, repo_path, ide=ide_type)
+
+    if success:
+        console.print(f'[green]✓[/] {message}')
+        console.print()
+        console.print(f'[dim]Restart {ide_name} for changes to take effect.[/]')
+    else:
+        console.print(f'[red]✗[/] {message}', style='bold red')
+        raise typer.Exit(1)

cycode/cli/apps/scan/code_scanner.py

@@ -91,7 +91,7 @@ def _should_use_sync_flow(command_scan_type: str, scan_type: str, sync_option: b
     if not sync_option and scan_type != consts.IAC_SCAN_TYPE:
         return False
 
-    if command_scan_type not in {'path', 'repository'}:
+    if command_scan_type not in {'path', 'repository', 'ai_guardrails'}:
         return False
 
     if scan_type == consts.IAC_SCAN_TYPE:

cycode/cli/cli_types.py

@@ -86,6 +86,10 @@ class SeverityOption(StrEnum):
     def get_member_emoji(name: str) -> str:
         return _SEVERITY_EMOJIS.get(name.lower(), _SEVERITY_DEFAULT_EMOJI)
 
+    @staticmethod
+    def get_member_unicode_emoji(name: str) -> str:
+        return _SEVERITY_UNICODE_EMOJIS.get(name.lower(), _SEVERITY_DEFAULT_UNICODE_EMOJI)
+
     def __rich__(self) -> str:
         color = self.get_member_color(self.value)
         return f'[{color}]{self.value.upper()}[/]'
@@ -117,3 +121,12 @@ _SEVERITY_EMOJIS = {
     SeverityOption.HIGH.value: ':red_circle:',
     SeverityOption.CRITICAL.value: ':exclamation_mark:',  # double_exclamation_mark is not red
 }
+
+_SEVERITY_DEFAULT_UNICODE_EMOJI = '⚪'
+_SEVERITY_UNICODE_EMOJIS = {
+    SeverityOption.INFO.value: '🔵',
+    SeverityOption.LOW.value: '🟡',
+    SeverityOption.MEDIUM.value: '🟠',
+    SeverityOption.HIGH.value: '🔴',
+    SeverityOption.CRITICAL.value: '❗',
+}

cycode/cli/utils/get_api_client.py

@@ -3,11 +3,17 @@ from typing import TYPE_CHECKING, Optional, Union
 import click
 
 from cycode.cli.user_settings.credentials_manager import CredentialsManager
-from cycode.cyclient.client_creator import create_import_sbom_client, create_report_client, create_scan_client
+from cycode.cyclient.client_creator import (
+    create_ai_security_manager_client,
+    create_import_sbom_client,
+    create_report_client,
+    create_scan_client,
+)
 
 if TYPE_CHECKING:
     import typer
 
+    from cycode.cyclient.ai_security_manager_client import AISecurityManagerClient
     from cycode.cyclient.import_sbom_client import ImportSbomClient
     from cycode.cyclient.report_client import ReportClient
     from cycode.cyclient.scan_client import ScanClient
@@ -19,7 +25,7 @@ def _get_cycode_client(
     client_secret: Optional[str],
     hide_response_log: bool,
     id_token: Optional[str] = None,
-) -> Union['ScanClient', 'ReportClient']:
+) -> Union['ScanClient', 'ReportClient', 'ImportSbomClient', 'AISecurityManagerClient']:
     if client_id and id_token:
         return create_client_func(client_id, None, hide_response_log, id_token)
 
@@ -62,6 +68,13 @@ def get_import_sbom_cycode_client(ctx: 'typer.Context', hide_response_log: bool
     return _get_cycode_client(create_import_sbom_client, client_id, client_secret, hide_response_log, id_token)
 
 
+def get_ai_security_manager_client(ctx: 'typer.Context', hide_response_log: bool = True) -> 'AISecurityManagerClient':
+    client_id = ctx.obj.get('client_id')
+    client_secret = ctx.obj.get('client_secret')
+    id_token = ctx.obj.get('id_token')
+    return _get_cycode_client(create_ai_security_manager_client, client_id, client_secret, hide_response_log, id_token)
+
+
 def _get_configured_credentials() -> tuple[str, str]:
     credentials_manager = CredentialsManager()
     return credentials_manager.get_credentials()

cycode/cli/utils/scan_utils.py

@@ -1,9 +1,12 @@
 import os
+from collections import defaultdict
 from typing import TYPE_CHECKING, Optional
 from uuid import UUID, uuid4
 
 import typer
 
+from cycode.cli.cli_types import SeverityOption
+
 if TYPE_CHECKING:
     from cycode.cli.models import LocalScanResult
     from cycode.cyclient.models import ScanConfiguration
@@ -33,3 +36,24 @@ def generate_unique_scan_id() -> UUID:
         return UUID(os.environ['PYTEST_TEST_UNIQUE_ID'])
 
     return uuid4()
+
+
+def build_violation_summary(local_scan_results: list['LocalScanResult']) -> str:
+    """Build violation summary string with severity breakdown and emojis."""
+    detections_count = 0
+    severity_counts = defaultdict(int)
+
+    for local_scan_result in local_scan_results:
+        for document_detections in local_scan_result.document_detections:
+            for detection in document_detections.detections:
+                if detection.severity:
+                    detections_count += 1
+                    severity_counts[SeverityOption(detection.severity)] += 1
+
+    severity_parts = []
+    for severity in reversed(SeverityOption):
+        emoji = SeverityOption.get_member_unicode_emoji(severity)
+        count = severity_counts[severity]
+        severity_parts.append(f'{emoji} {severity.upper()} - {count}')
+
+    return f'Cycode found {detections_count} violations: {" | ".join(severity_parts)}'

cycode/cyclient/ai_security_manager_client.py

@@ -0,0 +1,86 @@
+"""Client for AI Security Manager service."""
+
+from typing import TYPE_CHECKING, Optional
+
+from cycode.cli.exceptions.custom_exceptions import HttpUnauthorizedError
+from cycode.cyclient.cycode_client_base import CycodeClientBase
+from cycode.cyclient.logger import logger
+
+if TYPE_CHECKING:
+    from cycode.cli.apps.ai_guardrails.scan.payload import AIHookPayload
+    from cycode.cli.apps.ai_guardrails.scan.types import AiHookEventType, AIHookOutcome, BlockReason
+    from cycode.cyclient.ai_security_manager_service_config import AISecurityManagerServiceConfigBase
+
+
+class AISecurityManagerClient:
+    """Client for interacting with AI Security Manager service."""
+
+    _CONVERSATIONS_PATH = 'v4/ai-security/interactions/conversations'
+    _EVENTS_PATH = 'v4/ai-security/interactions/events'
+
+    def __init__(self, client: CycodeClientBase, service_config: 'AISecurityManagerServiceConfigBase') -> None:
+        self.client = client
+        self.service_config = service_config
+
+    def _build_endpoint_path(self, path: str) -> str:
+        """Build the full endpoint path including service name/port."""
+        service_name = self.service_config.get_service_name()
+        if service_name:
+            return f'{service_name}/{path}'
+        return path
+
+    def create_conversation(self, payload: 'AIHookPayload') -> Optional[str]:
+        """Creates an AI conversation from hook payload."""
+        conversation_id = payload.conversation_id
+        if not conversation_id:
+            return None
+
+        body = {
+            'id': conversation_id,
+            'ide_user_email': payload.ide_user_email,
+            'model': payload.model,
+            'ide_provider': payload.ide_provider,
+            'ide_version': payload.ide_version,
+        }
+
+        try:
+            self.client.post(self._build_endpoint_path(self._CONVERSATIONS_PATH), body=body)
+        except HttpUnauthorizedError:
+            # Authentication error - re-raise so prompt_command can catch it
+            raise
+        except Exception as e:
+            logger.debug('Failed to create conversation', exc_info=e)
+            # Don't fail the hook if tracking fails (non-auth errors)
+
+        return conversation_id
+
+    def create_event(
+        self,
+        payload: 'AIHookPayload',
+        event_type: 'AiHookEventType',
+        outcome: 'AIHookOutcome',
+        scan_id: Optional[str] = None,
+        block_reason: Optional['BlockReason'] = None,
+    ) -> None:
+        """Create an AI hook event from hook payload."""
+        conversation_id = payload.conversation_id
+        if not conversation_id:
+            logger.debug('No conversation ID available, skipping event creation')
+            return
+
+        body = {
+            'conversation_id': conversation_id,
+            'event_type': event_type,
+            'outcome': outcome,
+            'generation_id': payload.generation_id,
+            'block_reason': block_reason,
+            'cli_scan_id': scan_id,
+            'mcp_server_name': payload.mcp_server_name,
+            'mcp_tool_name': payload.mcp_tool_name,
+        }
+
+        try:
+            self.client.post(self._build_endpoint_path(self._EVENTS_PATH), body=body)
+        except Exception as e:
+            logger.debug('Failed to create AI hook event', exc_info=e)
+            # Don't fail the hook if tracking fails

cycode/cyclient/ai_security_manager_service_config.py

@@ -0,0 +1,27 @@
+"""Service configuration for AI Security Manager."""
+
+
+class AISecurityManagerServiceConfigBase:
+    """Base class for AI Security Manager service configuration."""
+
+    def get_service_name(self) -> str:
+        """Get the service name or port for URL construction.
+
+        In dev mode, returns the port number.
+        In production, returns the service name.
+        """
+        raise NotImplementedError
+
+
+class DevAISecurityManagerServiceConfig(AISecurityManagerServiceConfigBase):
+    """Dev configuration for AI Security Manager."""
+
+    def get_service_name(self) -> str:
+        return '5163/api'
+
+
+class DefaultAISecurityManagerServiceConfig(AISecurityManagerServiceConfigBase):
+    """Production configuration for AI Security Manager."""
+
+    def get_service_name(self) -> str:
+        return ''

cycode/cyclient/client_creator.py

@@ -1,5 +1,10 @@
 from typing import Optional
 
+from cycode.cyclient.ai_security_manager_client import AISecurityManagerClient
+from cycode.cyclient.ai_security_manager_service_config import (
+    DefaultAISecurityManagerServiceConfig,
+    DevAISecurityManagerServiceConfig,
+)
 from cycode.cyclient.config import dev_mode
 from cycode.cyclient.config_dev import DEV_CYCODE_API_URL
 from cycode.cyclient.cycode_dev_based_client import CycodeDevBasedClient
@@ -49,3 +54,18 @@ def create_import_sbom_client(
     else:
         client = CycodeTokenBasedClient(client_id, client_secret)
     return ImportSbomClient(client)
+
+
+def create_ai_security_manager_client(
+    client_id: str, client_secret: Optional[str] = None, _: bool = False, id_token: Optional[str] = None
+) -> AISecurityManagerClient:
+    if dev_mode:
+        client = CycodeDevBasedClient(DEV_CYCODE_API_URL)
+        service_config = DevAISecurityManagerServiceConfig()
+    else:
+        if id_token:
+            client = CycodeOidcBasedClient(client_id, id_token)
+        else:
+            client = CycodeTokenBasedClient(client_id, client_secret)
+        service_config = DefaultAISecurityManagerServiceConfig()
+    return AISecurityManagerClient(client, service_config)