praetorian-cli 2.2.4__py3-none-any.whl → 2.2.6__py3-none-any.whl

praetorian_cli/handlers/add.py

@@ -5,7 +5,7 @@ import click
 
 from praetorian_cli.handlers.chariot import chariot
 from praetorian_cli.handlers.cli_decorators import cli_handler, praetorian_only
-from praetorian_cli.handlers.utils import error
+from praetorian_cli.handlers.utils import error, parse_configuration_value
 from praetorian_cli.sdk.model.globals import AddRisk, Asset, Seed, Kind
 
 
@@ -274,40 +274,33 @@ def setting(sdk, name, value):
 @add.command()
 @cli_handler
 @click.option('-n', '--name', required=True, help='Name of the configuration')
-@click.option('-e', '--entry', required=True, multiple=True, help='Key-value pair in format key=value. Can be specified multiple times to set multiple values.')
+@click.option('-e', '--entry', required=False, multiple=True,
+              help='Key-value pair in format key=value. Can be specified multiple times to set multiple values.')
+@click.option('--string', 'string_value', required=False,
+              help='Set the configuration value to a string')
+@click.option('--integer', 'integer_value', required=False,
+              help='Set the configuration value to an integer')
+@click.option('--float', 'float_value', required=False,
+              help='Set the configuration value to a floating point number')
 @praetorian_only
-def configuration(sdk, name, entry):
+def configuration(sdk, name, entry, string_value, integer_value, float_value):
     """ Add a configuration
 
-    This command adds, or overwrites if exists, a name-value configuration.
+    This command adds, or overwrites if exists, a configuration value.
+
+    Configuration values can be provided as a mapping of key-value pairs using
+    ``--entry`` (the previous behavior), or as primitive values using
+    ``--string``, ``--integer``, or ``--float``.
 
     \b
     Example usages:
         - praetorian chariot add configuration --name "nuclei" --entry extra-tags=http,sql --entry something=else
+        - praetorian chariot add configuration --name "billing-status" --string PAID_MS
+        - praetorian chariot add configuration --name "request-timeout" --integer 60
+        - praetorian chariot add configuration --name "scoring-threshold" --float 0.85
     """
-    config_dict = {}
-    for item in entry:
-        if '=' not in item:
-            click.echo(f"Error: Entry '{item}' is not in the format key=value")
-            return
-
-        if item.count('=') > 1:
-            click.echo(f"Error: Entry '{item}' contains multiple '=' characters. Format should be key=value")
-            return
-
-        key, value = item.split('=', 1)
-
-        if not key:
-            click.echo("Error: Key cannot be empty")
-            return
-
-        if not value:
-            click.echo("Error: Value cannot be empty")
-            return
-
-        config_dict[key] = value
-
-    sdk.configurations.add(name, config_dict)
+    config_value = parse_configuration_value(entry, string_value, integer_value, float_value)
+    sdk.configurations.add(name, config_value)
 
 
 @add.command()

praetorian_cli/handlers/agent.py

@@ -69,3 +69,32 @@ def tools(sdk, allowed):
     """
     for  tool in dict.keys(sdk.agents.list_mcp_tools(allowed)):
         click.echo(tool)
+
+@agent.command()
+@cli_handler
+def conversation(sdk):
+    """ Interactive conversation with Chariot AI assistant
+    
+    Start an interactive chat session with the Chariot AI assistant.
+    The AI can help you query security data, understand findings,
+    and provide insights about your attack surface.
+    
+    \b
+    Commands within conversation:
+        - help    Show available commands and query examples
+        - clear   Clear the screen  
+        - new     Start a new conversation
+        - quit    Exit the conversation
+    
+    \b
+    Example queries:
+        - "Find all active assets"
+        - "Show me critical risks"
+        - "What assets do we have for example.com?"
+        
+    \b
+    Usage:
+        praetorian chariot agent conversation
+    """
+    from praetorian_cli.ui.conversation import run_textual_conversation
+    run_textual_conversation(sdk)

praetorian_cli/handlers/get.py

@@ -147,7 +147,8 @@ def file(chariot, name, path):
 @cli_handler
 @click.argument('name')
 @click.option('-path', '--path', default=os.getcwd(), help='Download path. Default: save to current directory')
-def definition(chariot, name, path):
+@click.option('--global', 'global_', is_flag=True, help='Fetch from global definitions instead of user-specific')
+def definition(chariot, name, path, global_):
     """ Download a definition using the risk name
 
     \b
@@ -158,8 +159,9 @@ def definition(chariot, name, path):
     Example usage:
         - praetorian chariot get definition jira-unauthenticated-user-picker
         - praetorian chariot get definition CVE-2024-23049
+        - praetorian chariot get definition CVE-2024-23049 --global
      """
-    downloaded_path = chariot.definitions.get(name, path)
+    downloaded_path = chariot.definitions.get(name, path, global_=global_)
     click.echo(f'Saved definition at {downloaded_path}')
 
 

praetorian_cli/handlers/utils.py

@@ -3,6 +3,65 @@ import json
 import click
 
 
+def parse_configuration_value(
+    entries = [], s_val=None, i_val=None, f_val=None):
+    """Return a configuration value derived from CLI inputs."""
+    has_entries = len(entries) > 0
+    typed_values = {
+        'string': s_val,
+        'integer': i_val,
+        'float': f_val,
+    }
+
+    if has_entries and any(value is not None for value in typed_values.values()):
+        error('--entry cannot be combined with --string, --integer, or --float')
+
+    if has_entries:
+        return _parse_entry_dict(entries)
+
+    provided = [(name, value) for name, value in typed_values.items() if value is not None]
+
+    if not provided:
+        error('Provide configuration data via --entry, --string, --integer, or --float')
+    if len(provided) > 1:
+        error('Specify only one of --string, --integer, or --float')
+
+    value_type, raw_value = provided[0]
+    return _cast_typed_value(value_type, raw_value)
+
+
+def _parse_entry_dict(entries):
+    parsed = {}
+
+    for item in entries:
+        key, value = _split_entry(item)
+        if not key:
+            error(f'Key cannot be empty: {item}')
+        if not value:
+            error(f'Value cannot be empty: {item}')
+        parsed[key] = value
+    return parsed
+
+
+def _split_entry(item):
+    if '=' not in item:
+        error(f"Entry '{item}' is not in the format key=value")
+    if item.count('=') > 1:
+        error(f"Entry '{item}' contains multiple '=' characters. Format should be key=value")
+    return item.split('=', 1)
+
+
+def _cast_typed_value(value_type, raw_value):
+    try:
+        if value_type == 'integer':
+            return int(raw_value)
+        if value_type == 'float':
+            return float(raw_value)
+    except ValueError:
+        error(f'{value_type} must be a valid {value_type}')
+    return raw_value
+
+
 def render_list_results(list_results, details):
     list_data, offset = list_results
     if details:

praetorian_cli/sdk/entities/configurations.py

@@ -27,7 +27,7 @@ class Configurations:
                 "Please contact your Praetorian representative for assistance."
             )
 
-    def add(self, name, value: dict):
+    def add(self, name, value):
         """
         Add or update a configuration.
 

praetorian_cli/sdk/entities/definitions.py

@@ -29,7 +29,7 @@ class Definitions:
             definition_name = os.path.basename(local_filepath)
         return self.api.files.add(local_filepath, f'definitions/{definition_name}')
 
-    def get(self, definition_name, download_directory=os.getcwd()):
+    def get(self, definition_name, download_directory=os.getcwd(), global_=False):
         """
         Download a risk definition file from the definitions folder.
 
@@ -37,15 +37,24 @@ class Definitions:
         :type definition_name: str
         :param download_directory: The directory to save the downloaded file (defaults to current working directory)
         :type download_directory: str
+        :param global_: If True, fetch from global definitions instead of user-specific
+        :type global_: bool
         :return: The local file path where the definition was saved
         :rtype: str
         """
-        content = self.api.files.get_utf8(f'definitions/{definition_name}')
+        try:
+            content = self.api.files.get_utf8(f'definitions/{definition_name}', _global=global_)
+        except Exception as e:
+            if global_:
+                raise Exception(f'Global definition {definition_name} not found or inaccessible.')
+            else:
+                raise
         download_path = os.path.join(download_directory, definition_name)
         with open(download_path, 'w') as file:
             file.write(content)
         return download_path
 
+
     def list(self, name_filter='', offset=None, pages=100000) -> tuple:
         """
         List the definition names, optionally prefix-filtered by a definition name.

praetorian_cli/sdk/entities/files.py

@@ -46,30 +46,35 @@ class Files:
 
         return download_path
 
-    def get(self, chariot_filepath) -> bytes:
+    def get(self, chariot_filepath, _global=False) -> bytes:
         """
         Download a file from Chariot storage into memory as bytes.
 
         :param chariot_filepath: Path of the file in Chariot storage to download
         :type chariot_filepath: str
+        :param _global: If True, fetch from global storage instead of user-specific
+        :type _global: bool
         :return: The file content as bytes
         :rtype: bytes
         :raises Exception: If the file does not exist in Chariot storage
         """
-        self.raise_if_missing(chariot_filepath)
-        return self.api.download(chariot_filepath)
+        if not _global:
+            self.raise_if_missing(chariot_filepath)
+        return self.api.download(chariot_filepath, global_=_global)
 
-    def get_utf8(self, chariot_filepath) -> str:
+    def get_utf8(self, chariot_filepath, _global=False) -> str:
         """
         Download a file from Chariot storage into memory as a UTF-8 string.
 
         :param chariot_filepath: Path of the file in Chariot storage to download
         :type chariot_filepath: str
+        :param _global: If True, fetch from global storage instead of user-specific
+        :type _global: bool
         :return: The file content as a UTF-8 decoded string
         :rtype: str
         :raises Exception: If the file does not exist in Chariot storage
         """
-        return self.get(chariot_filepath).decode('utf-8')
+        return self.get(chariot_filepath, _global=_global).decode('utf-8')
 
     def list(self, prefix_filter='', offset=None, pages=100000) -> tuple:
         """

praetorian_cli/sdk/entities/risks.py

@@ -1,5 +1,5 @@
 from praetorian_cli.sdk.model.globals import Kind
-from praetorian_cli.sdk.model.query import Relationship, Node, Query, risk_of_key, ASSET_NODE, ATTRIBUTE_NODE
+from praetorian_cli.sdk.model.query import Relationship, Node, Query, risk_of_key, ASSET_NODE, ATTRIBUTE_NODE, Filter
 
 
 class Risks:
@@ -85,12 +85,12 @@ class Risks:
 
         return self.api.delete_by_key('risk', key, body)
 
-    def list(self, prefix_filter='', offset=None, pages=100000) -> tuple:
+    def list(self, contains_filter='', offset=None, pages=100000) -> tuple:
         """
         List risks with optional filtering and pagination.
 
-        :param prefix_filter: Prefix filter to apply after the "#risk#" portion of the key. Risk keys follow format '#risk#{asset_dns}#{risk_name}'
-        :type prefix_filter: str
+        :param contains_filter: Filter to apply to the risk key. Ensure the risk's key contains the filter.
+        :type contains_filter: str
         :param offset: The offset of the page you want to retrieve results. If not supplied, retrieves from the first page
         :type offset: str or None
         :param pages: The number of pages of results to retrieve. <mcp>Start with one page of results unless specifically requested.</mcp>
@@ -98,7 +98,12 @@ class Risks:
         :return: A tuple containing (list of matching risks, next page offset)
         :rtype: tuple
         """
-        return self.api.search.by_key_prefix(f'#risk#{prefix_filter}', offset, pages)
+        filters = []
+        if contains_filter:
+            contains_filter_filter = Filter(field=Filter.Field.KEY, operator=Filter.Operator.CONTAINS, value=contains_filter)
+            filters.append(contains_filter_filter)
+        query = Query(Node([Node.Label.RISK], filters=filters))
+        return self.api.search.by_query(query, pages)
 
     def attributes(self, key):
         """

praetorian_cli/sdk/entities/search.py

@@ -1,6 +1,6 @@
 import json
 from praetorian_cli.sdk.model.query import Query
-from praetorian_cli.sdk.model.globals import EXACT_FLAG, DESCENDING_FLAG, GLOBAL_FLAG, Kind
+from praetorian_cli.sdk.model.globals import EXACT_FLAG, DESCENDING_FLAG, GLOBAL_FLAG, USER_FLAG, Kind
 class Search:
 
     def __init__(self, api):
@@ -17,7 +17,7 @@ class Search:
         """
         return self.api.count(dict(key=search_term))
 
-    def by_key_prefix(self, key_prefix, offset=None, pages=100000) -> tuple:
+    def by_key_prefix(self, key_prefix, offset=None, pages=100000, user=False) -> tuple:
         """
         Search for entities by key prefix. <mcp>If the response is too large, make your query more specific.<mcp>
 
@@ -30,7 +30,7 @@ class Search:
         :return: A tuple containing (list of matching entities, next page offset)
         :rtype: tuple
         """
-        return self.by_term(key_prefix, None, offset, pages)
+        return self.by_term(key_prefix, None, offset, pages, user=user)
 
     def by_exact_key(self, key, get_attributes=False) -> {}:
         """
@@ -119,7 +119,7 @@ class Search:
         return self.by_term(f'dns:{dns_prefix}', kind, offset, pages)
 
     def by_term(self, search_term, kind=None, offset=None, pages=100000, exact=False, descending=False,
-                global_=False) -> tuple:
+                global_=False, user=False) -> tuple:
         """
         Search for a given kind by term.
 
@@ -151,6 +151,8 @@ class Search:
             params |= DESCENDING_FLAG
         if global_:
             params |= GLOBAL_FLAG
+        if user:
+            params |= USER_FLAG
 
         results = self.api.my(params, pages)
 

praetorian_cli/sdk/model/globals.py

@@ -114,3 +114,4 @@ class Kind(Enum):
 EXACT_FLAG = {'exact': 'true'}
 DESCENDING_FLAG = {'desc': 'true'}
 GLOBAL_FLAG = {'global': 'true'}
+USER_FLAG = {'user': 'true'}

praetorian_cli/sdk/test/test_conversation.py

@@ -0,0 +1,195 @@
+import pytest
+import json
+import time
+
+from praetorian_cli.sdk.test.utils import make_test_values, clean_test_entities, setup_chariot
+
+
+@pytest.mark.coherence
+class TestConversation:
+    """Test conversation functionality with the Chariot agent"""
+
+    def setup_class(self):
+        self.sdk = setup_chariot()
+        make_test_values(self)
+        self.conversation_id = None
+        self.message_keys = []
+
+    def test_start_conversation(self):
+        """Test starting a new conversation with the agent"""
+        url = self.sdk.url("/planner")
+        payload = {
+            "message": f"Test conversation started at {int(time.time())}",
+            "mode": "query"
+        }
+        
+        response = self.sdk._make_request("POST", url, json=payload)
+        
+        # Should get successful response
+        assert response.status_code == 200
+        
+        result = response.json()
+        assert "conversation" in result
+        assert "uuid" in result["conversation"]
+        
+        self.conversation_id = result["conversation"]["uuid"]
+        assert self.conversation_id is not None
+        assert len(self.conversation_id) > 0
+
+    def test_send_message_with_existing_conversation(self):
+        """Test sending a message to an existing conversation"""
+        # First, start a new conversation
+        url = self.sdk.url("/planner")
+        initial_payload = {
+            "message": f"Initial message for follow-up test at {int(time.time())}",
+            "mode": "query"
+        }
+        
+        response = self.sdk._make_request("POST", url, json=initial_payload)
+        assert response.status_code == 200
+        
+        result = response.json()
+        conversation_id = result["conversation"]["uuid"]
+        
+        # Now send a follow-up message to the existing conversation
+        follow_up_payload = {
+            "message": f"Follow-up message in existing conversation at {int(time.time())}",
+            "mode": "query",
+            "conversationId": conversation_id
+        }
+        
+        response = self.sdk._make_request("POST", url, json=follow_up_payload)
+        
+        # Should get successful response
+        assert response.status_code == 200
+
+    def test_read_conversation_messages(self):
+        """Test reading messages from a conversation"""
+        # First, create a conversation and send a message
+        url = self.sdk.url("/planner")
+        payload = {
+            "message": f"Test message for reading at {int(time.time())}",
+            "mode": "query"
+        }
+        
+        response = self.sdk._make_request("POST", url, json=payload)
+        assert response.status_code == 200
+        
+        result = response.json()
+        conversation_id = result["conversation"]["uuid"]
+        
+        # Wait a moment for message to be stored
+        time.sleep(2)
+        
+        # Now try to read messages from this conversation
+        messages, offset = self.sdk.search.by_key_prefix(
+            f"#message#{conversation_id}#", user=True
+        )
+        
+        # Should have at least the user message
+        assert isinstance(messages, list)
+        # Note: messages might be empty immediately after creation due to async processing
+        # This is still a valid test as it verifies the search functionality works
+
+    def test_read_conversations_list(self):
+        """Test reading list of conversations"""
+        # First create a test conversation
+        url = self.sdk.url("/planner")
+        payload = {
+            "message": f"Test conversation for listing at {int(time.time())}",
+            "mode": "query"
+        }
+        
+        response = self.sdk._make_request("POST", url, json=payload)
+        assert response.status_code == 200
+        
+        # Wait a moment for conversation to be stored
+        time.sleep(1)
+        
+        # Now search for conversations
+        conversations, offset = self.sdk.search.by_key_prefix(
+            "#conversation#", user=True
+        )
+        
+        # Should be able to search for conversations (list might be empty or contain conversations)
+        assert isinstance(conversations, list)
+
+    def test_conversation_api_error_handling(self):
+        """Test handling of malformed requests"""
+        url = self.sdk.url("/planner")
+        # Send malformed payload (missing required fields)
+        malformed_payload = {
+            "invalid_field": "should cause error"
+            # Missing required "message" field
+        }
+        
+        response = self.sdk._make_request("POST", url, json=malformed_payload)
+        
+        # Should get a 4xx error for bad request
+        assert response.status_code >= 400
+
+    def test_conversation_modes(self):
+        """Test different conversation modes (query vs agent)"""
+        url = self.sdk.url("/planner")
+        
+        # Test query mode
+        query_payload = {
+            "message": f"Test query mode at {int(time.time())}",
+            "mode": "query"
+        }
+        
+        response = self.sdk._make_request("POST", url, json=query_payload)
+        assert response.status_code == 200
+        
+        result = response.json()
+        assert "conversation" in result
+        
+        # Test agent mode
+        agent_payload = {
+            "message": f"Test agent mode at {int(time.time())}",
+            "mode": "agent"
+        }
+        
+        response = self.sdk._make_request("POST", url, json=agent_payload)
+        assert response.status_code == 200
+        
+        result = response.json()
+        assert "conversation" in result
+
+    def test_message_polling_integration(self):
+        """Test integration of message polling to check for new messages"""
+        # Create a conversation first
+        url = self.sdk.url("/planner")
+        payload = {
+            "message": f"Test message for polling at {int(time.time())}",
+            "mode": "query"
+        }
+        
+        response = self.sdk._make_request("POST", url, json=payload)
+        assert response.status_code == 200
+        
+        result = response.json()
+        conversation_id = result["conversation"]["uuid"]
+        
+        # Test polling for messages (first call)
+        messages, offset = self.sdk.search.by_key_prefix(
+            f"#message#{conversation_id}#", user=True
+        )
+        
+        initial_count = len(messages) if messages else 0
+        
+        # Wait a moment and check again (simulating polling)
+        time.sleep(1)
+        
+        messages, offset = self.sdk.search.by_key_prefix(
+            f"#message#{conversation_id}#", user=True
+        )
+        
+        # Should still return valid result (messages may or may not have changed)
+        assert isinstance(messages, list)
+
+    def teardown_class(self):
+        """Clean up test data"""
+        # Note: Conversations and messages are typically read-only in tests
+        # Real cleanup would depend on having delete methods for conversations
+        pass

praetorian_cli/sdk/test/test_z_cli.py

@@ -5,6 +5,7 @@ from subprocess import run
 import pytest
 
 from praetorian_cli.sdk.model.globals import AddRisk, Asset, Risk, Seed, Preseed
+from praetorian_cli.sdk.model.utils import configuration_key
 from praetorian_cli.sdk.test.utils import epoch_micro, random_ip, make_test_values, clean_test_entities, setup_chariot
 
 
@@ -273,6 +274,24 @@ class TestZCli:
 
         self.verify(f'delete configuration "{o.configuration_key}"', ignore_stdout=True)
 
+        string_name = f'{o.configuration_name}-string'
+        string_key = configuration_key(string_name)
+        self.verify(f'add configuration --name "{string_name}" --string PAID_MS')
+        self.verify(f'get configuration "{string_key}"', expected_stdout=[string_key, string_name, '"value": "PAID_MS"'])
+        self.verify(f'delete configuration "{string_key}"', ignore_stdout=True)
+
+        integer_name = f'{o.configuration_name}-integer'
+        integer_key = configuration_key(integer_name)
+        self.verify(f'add configuration --name "{integer_name}" --integer 123')
+        self.verify(f'get configuration "{integer_key}"', expected_stdout=[integer_key, integer_name, '"value": 123'])
+        self.verify(f'delete configuration "{integer_key}"', ignore_stdout=True)
+
+        float_name = f'{o.configuration_name}-float'
+        float_key = configuration_key(float_name)
+        self.verify(f'add configuration --name "{float_name}" --float 1.5')
+        self.verify(f'get configuration "{float_key}"', expected_stdout=[float_key, float_name, '"value": 1.5'])
+        self.verify(f'delete configuration "{float_key}"', ignore_stdout=True)
+
     def test_webapplication_cli(self):
         o = make_test_values(lambda: None)
         self.verify(f'add asset --dns "{o.webapp_name}" --name "{o.webapp_url}" --type webapplication')
@@ -339,6 +358,11 @@ class TestZCli:
 
         self.verify('link --help', ignore_stdout=True)
         self.verify('link account --help', ignore_stdout=True)
+        
+        # Agent commands
+        self.verify('agent --help', ignore_stdout=True)
+        self.verify('agent conversation --help', ignore_stdout=True)
+        self.verify('agent mcp --help', ignore_stdout=True)
         self.verify('link webpage-source --help', ignore_stdout=True)
 
         self.verify('unlink --help', ignore_stdout=True)

praetorian_cli/ui/conversation/__init__.py

@@ -0,0 +1,3 @@
+from .textual_chat import run_textual_conversation
+
+__all__ = ['run_textual_conversation']