praetorian-cli 2.2.1__py3-none-any.whl → 2.2.3__py3-none-any.whl

praetorian_cli/sdk/entities/seeds.py

@@ -1,5 +1,6 @@
 from praetorian_cli.handlers.utils import error
 from praetorian_cli.sdk.model.globals import Seed, Kind
+from praetorian_cli.sdk.model.query import Query, Node, Filter, KIND_TO_LABEL
 
 
 class Seeds:
@@ -9,97 +10,148 @@ class Seeds:
     def __init__(self, api):
         self.api = api
 
-    def add(self, dns, status=Seed.PENDING.value):
+    def add(self, status=Seed.PENDING.value, seed_type=Kind.ASSET.value, **kwargs):
         """
-        Add a seed to the account.
-
-        :param dns: The DNS name, IP address, or CIDR range to add as a seed. Accepts domain names (e.g., 'example.com'), IP addresses (e.g., '192.168.1.1'), or CIDR ranges (e.g., '192.168.1.0/24')
-        :type dns: str
-        :param status: Seed status from Seed enum ('A' for Active, 'F' for Frozen, 'D' for Deleted, 'P' for Pending, 'FR' for Frozen Rejected)
-        :type status: str
+        Add a seed of specified type with dynamic fields.
+        
+        :param status: Status for backward compatibility  
+        :type status: str or None
+        :param type: Asset type (e.g., 'asset', 'addomain', etc.)
+        :type type: str
+        :param kwargs: Dynamic fields for the asset type
         :return: The seed that was added
         :rtype: dict
         """
-        return self.api.upsert('seed', dict(dns=dns, status=status))
+        # Handle status if provided
+        kwargs['status'] = status
+            
+        # Build payload with type wrapper
+        payload = {
+            'type': seed_type,
+            'model': kwargs
+        }
+
+        return self.api.upsert('seed', payload)
 
     def get(self, key):
         """
         Get details of a seed by key.
 
-        :param key: Entity key in format #seed#{type}#{dns} where type is 'domain', 'ip', or 'cidr' and dns is the seed value
+        :param key: Entity key (e.g., '#asset#example.com#example.com')
         :type key: str
         :return: The seed matching the specified key, or None if not found
         :rtype: dict or None
         """
-        return self.api.search.by_exact_key(key, False)
-
-    def update(self, key, status):
+        
+        # Create a Filter for the key field
+        key_filter = Filter(
+            field=Filter.Field.KEY,
+            operator=Filter.Operator.EQUAL,
+            value=key
+        )
+        
+        # Create a Node with Seed label and key filter
+        node = Node(
+            labels=[Node.Label.SEED],
+            filters=[key_filter]
+        )
+        
+        # Create the Query object
+        query = Query(node=node)
+        
+        # Call by_query with the constructed Query object
+        results_tuple = self.api.search.by_query(query)
+        if not results_tuple:
+            return None
+        
+        results, _ = results_tuple
+        if len(results) == 0:
+            return None
+        return results[0]
+
+    def update(self, key, status=None):
         """
-        Update a seed's status.
-
-        Note: The seed PUT endpoint is different from other PUT endpoints. This method
-        internally uses the DNS of the original seed rather than the key for the update operation.
-
-        :param key: Entity key in format #seed#{type}#{dns} where type is 'domain', 'ip', or 'cidr' and dns is the seed value
+        Update seed fields dynamically.
+        
+        :param key: Seed/Asset key (e.g., '#seed#domain#example.com' or '#asset#domain#example.com')
         :type key: str
-        :param status: Seed status from Seed enum ('A' for Active, 'F' for Frozen, 'D' for Deleted, 'P' for Pending, 'FR' for Frozen Rejected)
-        :type status: str
+        :param status: Status for backward compatibility (can be positional)
+        :type status: str or None
+        :param kwargs: Fields to update
         :return: The updated seed, or None if the seed was not found
         :rtype: dict or None
         """
-        seed = self.api.search.by_exact_key(key)
+            
+        seed = self.get(key)  # This already handles old key format conversion
         if seed:
-            # the seed PUT endpoint is different from other PUT endpoints. This one has to
-            # take the DNS of the original seed, instead of the key of the seed record.
-            # TODO, 2024-12-23, peter: check with Noah as to why. Ideally, we should
-            # standardize to how other endpoints do it
-            return self.api.upsert('seed', dict(dns=seed['dns'], status=status))
+            update_payload = {
+                'key': key,
+                'status': status
+            }
+            
+            return self.api.upsert('seed', update_payload)
         else:
-            error(f'Seed {key} is not found.')
+            error(f'Seed {key} not found.')
 
     def delete(self, key):
         """
-        Delete a seed by setting its status to DELETED.
-
-        Note: This method does not actually delete the seed from the database. Instead,
-        it sets the seed's status to DELETED ('D'), which marks it as deleted while
-        preserving the record for audit purposes.
-
-        :param key: Entity key in format #seed#{type}#{dns} where type is 'domain', 'ip', or 'cidr' and dns is the seed value
+        Delete seed (supports both old and new key formats).
+        
+        :param key: Seed/Asset key (e.g., '#asset#domain#example.com')
         :type key: str
         :return: The seed that was marked as deleted, or None if the seed was not found
         :rtype: dict or None
         """
-        seed = self.api.search.by_exact_key(key)
+        seed = self.get(key)  # This already handles old key format conversion
+        
         if seed:
-            # TODO, 2024-12-23, peter: check with Noah why this is different from
-            # deleting assets and risks
-            return self.api.upsert('seed', dict(dns=seed['dns'], status=Seed.DELETED.value))
+            delete_payload = {
+                'key': key,
+                'status': Seed.DELETED.value
+            }
+            
+            return self.api.upsert('seed', delete_payload)
         else:
-            error(f'Seed {key} is not found.')
+            error(f'Seed {key} not found.')
 
-    def list(self, type='', prefix_filter='', offset=None, pages=100000) -> tuple:
+    def list(self, seed_type=Kind.SEED.value, key_prefix='', pages=100000) -> tuple:
         """
-        List seeds with optional filtering.
-
-        :param type: The type of seed to filter by ('domain', 'ip', 'cidr'). If empty, returns all seed types
-        :type type: str
-        :param prefix_filter: Supply this to perform prefix-filtering of the seed DNS/IP values after the seed type portion of the key
-        :type prefix_filter: str
-        :param offset: The offset of the page you want to retrieve results. If this is not supplied, this function retrieves from the first page
-        :type offset: str or None
+        List seeds by querying assets with 'Seed' label.
+        
+        :param seed_type: Optional asset seed_type filter (e.g., 'asset', 'addomain')
+        :seed_type seed_type: str or None
+        :param key_prefix: Filter by key prefix
+        :seed_type key_prefix: str
         :param pages: The number of pages of results to retrieve. <mcp>Start with one page of results unless specifically requested.</mcp>
-        :type pages: int
+        :seed_type pages: int
         :return: A tuple containing (list of seeds, next page offset)
-        :rtype: tuple
+        :rseed_type: tuple
         """
-        prefix_term = '#seed#'
-        if type:
-            prefix_term = f'{prefix_term}{type}#'
-        if prefix_filter:
-            prefix_term = f'{prefix_term}{prefix_filter}'
 
-        return self.api.search.by_key_prefix(prefix_term, offset, pages)
+        if seed_type in KIND_TO_LABEL:
+            seed_type = KIND_TO_LABEL[seed_type]
+        elif not seed_type:
+            seed_type = Node.Label.SEED
+        else:
+            raise ValueError(f'Invalid seed type: {seed_type}')
+
+        node = Node(
+            labels=[seed_type],
+            filters=[]
+        )
+
+        key_filter = Filter(
+                field=Filter.Field.KEY,
+                operator=Filter.Operator.STARTS_WITH,
+                value=key_prefix
+            )
+
+        if key_prefix:
+            node.filters.append(key_filter)
+        
+        query = Query(node=node)
+        
+        return self.api.search.by_query(query, pages)
 
     def attributes(self, key):
         """

praetorian_cli/sdk/mcp_server.py

@@ -8,7 +8,6 @@ from mcp.server.lowlevel import Server
 from mcp.server.stdio import stdio_server
 from mcp.types import Tool, TextContent
 
-
 class MCPServer:
     def __init__(self, chariot_instance, allowable_tools: Optional[List[str]] = None):
         self.chariot = chariot_instance
@@ -157,7 +156,7 @@ class MCPServer:
             tools = []
             for tool_name, tool_info in self.discovered_tools.items():
                 parameters = self._extract_parameters_from_doc(tool_info['doc'], tool_info['signature'])
-                
+
                 properties = {}
                 required = []
                 
@@ -170,7 +169,7 @@ class MCPServer:
                         "description": param_info["description"]
                     }
                     
-                    if param_info["required"]:
+                    if param_info.get("required", False):
                         required.append(param_name)
                 
                 tool_schema = {

praetorian_cli/sdk/model/aegis.py

@@ -0,0 +1,156 @@
+"""
+Aegis agent data models and structures.
+
+This module contains dataclass definitions for Aegis agent entities,
+including network interfaces, tunnel status, health checks, and agent metadata.
+"""
+from dataclasses import dataclass
+from typing import Optional, List, Dict, Any
+from datetime import datetime
+
+
+@dataclass
+class NetworkInterface:
+    """Represents a network interface on an agent"""
+    name: str
+    ip_addresses: List[str]
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> 'NetworkInterface':
+        return cls(
+            name=data.get('name', ''),
+            ip_addresses=data.get('ip_addresses', [])
+        )
+
+
+@dataclass 
+class CloudflaredStatus:
+    """Represents Cloudflared tunnel status"""
+    hostname: Optional[str] = None
+    tunnel_name: Optional[str] = None
+    authorized_users: Optional[str] = None
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> 'CloudflaredStatus':
+        return cls(
+            hostname=data.get('hostname'),
+            tunnel_name=data.get('tunnel_name'),
+            authorized_users=data.get('authorized_users')
+        )
+
+
+@dataclass
+class HealthCheck:
+    """Represents agent health check data"""
+    cloudflared_status: Optional[CloudflaredStatus] = None
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> 'HealthCheck':
+        cf_data = data.get('cloudflared_status')
+        return cls(
+            cloudflared_status=CloudflaredStatus.from_dict(cf_data) if cf_data else None
+        )
+
+
+@dataclass
+class Agent:
+    """Represents an Aegis agent"""
+    client_id: str = 'N/A'
+    hostname: str = 'Unknown'
+    fqdn: str = 'N/A'
+    os: str = 'unknown'
+    os_version: str = ''
+    architecture: str = 'Unknown'
+    last_seen_at: Optional[int] = None
+    network_interfaces: List[NetworkInterface] = None
+    health_check: Optional[HealthCheck] = None
+    key: Optional[str] = None
+    
+    def __post_init__(self):
+        if self.network_interfaces is None:
+            self.network_interfaces = []
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> 'Agent':
+        """Create an Agent from dictionary data"""
+        network_interfaces = []
+        for iface_data in data.get('network_interfaces', []):
+            if isinstance(iface_data, dict):
+                network_interfaces.append(NetworkInterface.from_dict(iface_data))
+        
+        health_data = data.get('health_check')
+        health_check = HealthCheck.from_dict(health_data) if health_data else None
+        
+        return cls(
+            client_id=data.get('client_id', 'N/A'),
+            hostname=data.get('hostname', 'Unknown'),
+            fqdn=data.get('fqdn', 'N/A'),
+            os=data.get('os', 'unknown'),
+            os_version=data.get('os_version', ''),
+            architecture=data.get('architecture', 'Unknown'),
+            last_seen_at=data.get('last_seen_at'),
+            network_interfaces=network_interfaces,
+            health_check=health_check,
+            key=data.get('key')
+        )
+    
+    
+    @property
+    def has_tunnel(self) -> bool:
+        """Check if agent has an active Cloudflare tunnel"""
+        return (self.health_check is not None and 
+                self.health_check.cloudflared_status is not None and
+                self.health_check.cloudflared_status.hostname is not None)
+    
+    @property
+    def is_online(self) -> bool:
+        """Check if agent is currently online (last seen within 60 seconds)"""
+        if not self.last_seen_at:
+            return False
+        
+        current_time = datetime.now().timestamp()
+        last_seen_seconds = (self.last_seen_at / 1000000 
+                           if self.last_seen_at > 1000000000000 
+                           else self.last_seen_at)
+        
+        return (current_time - last_seen_seconds) < 60
+    
+    @property 
+    def ip_addresses(self) -> List[str]:
+        """Get all non-loopback IP addresses"""
+        ips = []
+        for iface in self.network_interfaces:
+            if iface.name != 'lo':  # Skip loopback
+                ips.extend(iface.ip_addresses)
+        return [ip for ip in ips if ip]  # Filter empty strings
+    
+    def __str__(self) -> str:
+        """Return a simple string representation of the agent"""
+        status = "🔗" if self.has_tunnel else "○"
+        return f"{status} {self.hostname} ({self.client_id})"
+    
+    def to_detailed_string(self) -> str:
+        """Return a detailed string representation of the agent"""
+        os_info = f"{self.os} {self.os_version}".strip()
+        
+        lines = [
+            f"\n{self.hostname} ({self.client_id})",
+            f"  OS: {os_info}",
+            f"  Architecture: {self.architecture}",
+            f"  FQDN: {self.fqdn}"
+        ]
+        
+        if self.has_tunnel:
+            cf_status = self.health_check.cloudflared_status
+            lines.append(f"  Tunnel: {cf_status.tunnel_name}")
+            lines.append(f"  Public hostname: {cf_status.hostname}")
+            if cf_status.authorized_users:
+                lines.append(f"  Authorized users: {cf_status.authorized_users}")
+        else:
+            lines.append("  Tunnel: Not configured")
+        
+        ips = self.ip_addresses
+        if ips:
+            lines.append(f"  IP addresses: {', '.join(ips)}")
+        
+        return '\n'.join(lines)

praetorian_cli/sdk/model/query.py

@@ -91,7 +91,7 @@ class Node:
         ASSET = 'Asset'
         REPOSITORY = 'Repository'
         INTEGRATION = 'Integration'
-        ADDOMAIN = 'Addomain'
+        ADDOMAIN = 'ADDomain'
         ATTRIBUTE = 'Attribute'
         RISK = 'Risk'
         PRESEED = 'Preseed'

praetorian_cli/sdk/model/utils.py

@@ -15,14 +15,11 @@ def integration_key(dns, name):
 def risk_key(dns, name):
     return f'#risk#{dns}#{name}'
 
-
 def attribute_key(name, value, source_key):
     return f'#attribute#{name}#{value}{source_key}'
 
-
-def seed_key(type, dns):
-    return f'#seed#{type}#{dns}'
-
+def seed_asset_key(dns):
+    return f'#asset#{dns}#{dns}'
 
 def preseed_key(type, title, value):
     return f'#preseed#{type}#{title}#{value}'
@@ -32,6 +29,3 @@ def setting_key(name):
 
 def configuration_key(name):
     return f'#configuration#{name}'
-
-def seed_status(type, status_code):
-    return f'{type}#{status_code}'

praetorian_cli/sdk/test/pytest.ini

@@ -3,3 +3,4 @@ addopts = -p no:cacheprovider
 markers =
     coherence: these are tests on product workflows via the backend API
     cli: these are tests on the CLI interface
+    tui: these are tests on the TUI interface

praetorian_cli/sdk/test/test_asset.py

@@ -42,7 +42,7 @@ class TestAsset:
         assert any([a['group'] == self.asset_dns for a in deleted_assets])
     
     def test_add_ad_domain(self):
-        asset = self.sdk.assets.add(self.ad_domain_name, self.ad_domain_name, status=Asset.ACTIVE.value, surface='test-surface', type=Kind.ADDOMAIN.value)
+        asset = self.sdk.assets.add(self.ad_domain_name, self.ad_object_id, status=Asset.ACTIVE.value, surface='test-surface', type=Kind.ADDOMAIN.value)
         assert asset['key'] == self.ad_domain_key
         assert len(asset['attackSurface']) == 1
         assert 'test-surface' in asset['attackSurface']
@@ -51,7 +51,7 @@ class TestAsset:
     def test_get_ad_domain(self):
         asset = self.sdk.assets.get(self.ad_domain_key)
         assert asset['key'] == self.ad_domain_key
-        assert asset['name'] == self.ad_domain_name
+        assert asset['domain'] == self.ad_domain_name
         assert asset['status'] == Asset.ACTIVE.value
     
     def test_list_ad_domain(self):

praetorian_cli/sdk/test/test_seed.py

@@ -1,7 +1,6 @@
 import pytest
 
-from praetorian_cli.sdk.model.globals import Seed
-from praetorian_cli.sdk.model.utils import seed_status
+from praetorian_cli.sdk.model.globals import Seed, Kind
 from praetorian_cli.sdk.test.utils import make_test_values, clean_test_entities, setup_chariot
 
 
@@ -12,30 +11,30 @@ class TestSeed:
         self.sdk = setup_chariot()
         make_test_values(self)
 
-    def test_add_seed(self):
-        seed = self.sdk.seeds.add(self.seed_dns)
-        assert seed['key'] == self.seed_key
+    def test_add_asset_seed(self):
+        seed = self.sdk.seeds.add(dns=self.seed_asset_dns)
+        assert seed['key'] == self.seed_asset_key
 
     def test_get_seed(self):
         a = self.get_seed()
-        assert a['dns'] == self.seed_dns
-        assert a['status'] == seed_status('domain', Seed.PENDING.value)
+        assert a['dns'] == self.seed_asset_dns
+        assert a['status'] == Seed.PENDING.value
 
     def test_list_seed(self):
-        results, _ = self.sdk.seeds.list('domain', self.seed_dns)
+        results, _ = self.sdk.seeds.list(Kind.ASSET.value, f"#asset#{self.seed_asset_dns}")
         assert len(results) == 1
-        assert results[0]['dns'] == self.seed_dns
+        assert results[0]['dns'] == self.seed_asset_dns
 
     def test_update_seed(self):
-        self.sdk.seeds.update(self.seed_key, Seed.ACTIVE.value)
-        assert self.get_seed()['status'] == seed_status('domain', Seed.ACTIVE.value)
+        self.sdk.seeds.update(self.seed_asset_key, Seed.ACTIVE.value)
+        assert self.get_seed()['status'] == Seed.ACTIVE.value
 
     def test_delete_seed(self):
-        self.sdk.seeds.delete(self.seed_key)
-        assert self.sdk.seeds.get(self.seed_key)['status'] == seed_status('domain', Seed.DELETED.value)
+        self.sdk.seeds.delete(self.seed_asset_key)
+        assert self.sdk.seeds.get(self.seed_asset_key)['status'] == Seed.DELETED.value
 
     def get_seed(self):
-        return self.sdk.seeds.get(self.seed_key)
+        return self.sdk.seeds.get(self.seed_asset_key)
 
     def teardown_class(self):
         clean_test_entities(self.sdk, self)

praetorian_cli/sdk/test/test_z_cli.py

@@ -5,7 +5,6 @@ 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 seed_status
 from praetorian_cli.sdk.test.utils import epoch_micro, random_ip, make_test_values, clean_test_entities, setup_chariot
 
 
@@ -21,11 +20,12 @@ class TestZCli:
         self.verify(f'add asset -n {o.asset_name} -d {o.asset_dns}')
 
         self.verify('list assets -p all', [o.asset_key])
-        self.verify(f'list assets -f "{o.asset_dns}"', [o.asset_key])
-        self.verify(f'list assets -f "{o.asset_dns}" -p first', [o.asset_key])
-        self.verify(f'list assets -f "{o.asset_dns}" -p all', [o.asset_key])
-        self.verify(f'list assets -f "{o.asset_dns}" -d', [o.asset_key, '"key"', '"data"'])
+        self.verify(f'list assets -f "#asset#{o.asset_dns}"', [o.asset_key])
+        self.verify(f'list assets -f "#asset#{o.asset_dns}" -p first', [o.asset_key])
+        self.verify(f'list assets -f "#asset#{o.asset_dns}" -p all', [o.asset_key])
+        self.verify(f'list assets -f "#asset#{o.asset_dns}" -d', [o.asset_key, '"key"', '"data"'])
 
+        self.verify(f'list assets -f "#asset#{epoch_micro()}"')
         self.verify(f'list assets -f {epoch_micro()}')
 
         self.verify(f'get asset "{o.asset_key}"', [o.asset_key, f'"status": "{Asset.ACTIVE.value}"'])
@@ -43,30 +43,28 @@ class TestZCli:
     def test_seed_cli(self):
         o = make_test_values(lambda: None)
 
-        self.verify(f'add seed -d {o.seed_dns}')
+        self.verify(f'add seed --field dns:{o.seed_asset_dns}')
 
-        self.verify('list seeds -p all', [o.seed_key])
-        self.verify('list seeds -t domain -p all', [o.seed_key])
-        self.verify(f'list seeds -t domain -f "{o.seed_dns}"', [o.seed_key])
-        self.verify(f'list seeds -t domain -f "{o.seed_dns}" -p first', [o.seed_key])
-        self.verify(f'list seeds -t domain -f "{o.seed_dns}" -p all', [o.seed_key])
-        self.verify(f'list seeds -t domain -f "{o.seed_dns}" -p first', [o.seed_key])
-        self.verify(f'list seeds -t domain -f "{o.seed_dns}" -d', [o.seed_dns, '"key"', '"data"'])
-        self.verify(f'list seeds -t ip -f "{o.seed_dns}"')
-        self.verify(f'list seeds -f "{o.seed_dns}"', [],
-                    ["When the DNS filter is specified, you also need to specify the type of the filter"])
+        self.verify('list seeds -p all', [o.seed_asset_key])
+        self.verify('list seeds -t asset -p all', [o.seed_asset_key])
+        self.verify(f'list seeds -t asset -f "#asset#{o.seed_asset_dns}"', [o.seed_asset_key])
+        self.verify(f'list seeds -t asset -f "#asset#{o.seed_asset_dns}" -p first', [o.seed_asset_key])
+        self.verify(f'list seeds -t asset -f "#asset#{o.seed_asset_dns}" -p all', [o.seed_asset_key])
+        self.verify(f'list seeds -t asset -f "#asset#{o.seed_asset_dns}" -d', [o.seed_asset_dns, '"key"', '"data"'])
+        self.verify(f'list seeds -t notatype -f "#asset#{o.seed_asset_dns}"', [], ['Invalid seed type: notatype'])
+        self.verify(f'list seeds -f "#asset#{o.seed_asset_dns}"', [o.seed_asset_key])
 
-        self.verify(f'list seeds -t domain -f {epoch_micro()}')
+        self.verify(f'list seeds -t asset -f {epoch_micro()}')
 
-        self.verify(f'get seed "{o.seed_key}"',
-                    [o.seed_key, f'"status": "{seed_status("domain", Seed.PENDING.value)}"'])
+        self.verify(f'get seed "{o.seed_asset_key}"',
+                    [o.seed_asset_key, f'"status": "{Seed.PENDING.value}"'])
 
-        self.verify(f'update seed -s {Seed.ACTIVE.value} "{o.seed_key}"')
-        self.verify(f'get seed "{o.seed_key}"',
-                    [o.seed_key, f'"status": "{seed_status("domain", Seed.ACTIVE.value)}"'])
+        self.verify(f'update seed -s {Seed.ACTIVE.value} "{o.seed_asset_key}"')
+        self.verify(f'get seed "{o.seed_asset_key}"',
+                    [o.seed_asset_key, f'"status": "{Seed.ACTIVE.value}"'])
 
-        self.verify(f'delete seed "{o.seed_key}"')
-        self.verify(f'get seed "{o.seed_key}"', [f'"status": "{seed_status("domain", Seed.DELETED.value)}"'])
+        self.verify(f'delete seed "{o.seed_asset_key}"')
+        self.verify(f'get seed "{o.seed_asset_key}"', [f'"status": "{Seed.DELETED.value}"'])
 
         clean_test_entities(self.sdk, o)