binalyze-air-sdk 1.0.2__py3-none-any.whl → 1.0.3__py3-none-any.whl

binalyze_air/models/interact.py

@@ -1,136 +1,267 @@
-"""
-Interact API models for the Binalyze AIR SDK.
-"""
-
-from typing import List, Optional, Dict, Any
-from datetime import datetime
-from enum import Enum
-from pydantic import Field
-
-from ..base import AIRBaseModel
-
-
-class SendToLocation(str, Enum):
-    """Send to location options."""
-    USER_LOCAL = "user-local"
-    REPOSITORY = "repository"
-    EVIDENCE_REPOSITORY = "evidence-repository"
-
-
-class TaskConfigChoice(str, Enum):
-    """Task configuration choice options."""
-    USE_POLICY = "use-policy"
-    USE_CUSTOM_OPTIONS = "use-custom-options"
-
-
-class SendToConfig(AIRBaseModel):
-    """Send to configuration model."""
-    
-    location: SendToLocation
-    repository_id: Optional[str] = Field(default=None, alias="repositoryId")
-    evidence_repository_id: Optional[str] = Field(default=None, alias="evidenceRepositoryId")
-
-
-class BandwidthConfig(AIRBaseModel):
-    """Bandwidth configuration model."""
-    
-    limit: Optional[int] = None
-
-
-class DiskSpaceConfig(AIRBaseModel):
-    """Disk space configuration model."""
-    
-    reserve: Optional[int] = None
-
-
-class TaskConfig(AIRBaseModel):
-    """Task configuration model."""
-    
-    choice: TaskConfigChoice
-    send_to: Optional[SendToConfig] = Field(default=None, alias="sendTo")
-    bandwidth: Optional[BandwidthConfig] = None
-    disk_space: Optional[DiskSpaceConfig] = Field(default=None, alias="diskSpace")
-
-
-class AssignInteractiveShellTaskRequest(AIRBaseModel):
-    """Request model for assigning interactive shell task."""
-    
-    asset_id: str = Field(alias="assetId")
-    case_id: str = Field(alias="caseId")
-    task_config: TaskConfig = Field(alias="taskConfig")
-
-
-class InteractiveShellTaskResponse(AIRBaseModel):
-    """Response model for interactive shell task assignment."""
-    
-    session_id: str = Field(alias="sessionId")
-    idle_timeout: int = Field(alias="idleTimeout")
-    config: TaskConfig
-
-
-# Legacy models for backward compatibility (deprecated)
-class InteractionType(str, Enum):
-    """Interaction types."""
-    SHELL = "shell"
-    POWERSHELL = "powershell"
-    CMD = "cmd"
-    BASH = "bash"
-
-
-class InteractionStatus(str, Enum):
-    """Interaction status."""
-    PENDING = "pending"
-    RUNNING = "running"
-    COMPLETED = "completed"
-    FAILED = "failed"
-    CANCELLED = "cancelled"
-    TIMEOUT = "timeout"
-
-
-class ShellInteraction(AIRBaseModel):
-    """Shell interaction model (legacy)."""
-    
-    id: str
-    task_id: str
-    endpoint_id: str
-    endpoint_name: str
-    interaction_type: InteractionType
-    command: str
-    output: Optional[str] = None
-    error_output: Optional[str] = None
-    exit_code: Optional[int] = None
-    status: InteractionStatus = InteractionStatus.PENDING
-    timeout: int = 300  # seconds
-    organization_id: int
-    created_by: Optional[str] = None
-    created_at: Optional[datetime] = None
-    started_at: Optional[datetime] = None
-    completed_at: Optional[datetime] = None
-    duration: Optional[int] = None  # seconds
-    environment_variables: Optional[Dict[str, str]] = None
-    working_directory: Optional[str] = None
-
-
-class AssignShellTaskRequest(AIRBaseModel):
-    """Request model for assigning shell interaction tasks (legacy)."""
-    
-    endpoint_ids: List[str]
-    command: str
-    interaction_type: InteractionType = InteractionType.SHELL
-    timeout: Optional[int] = 300
-    organization_ids: Optional[List[int]] = None
-    case_id: Optional[str] = None
-    environment_variables: Optional[Dict[str, str]] = None
-    working_directory: Optional[str] = None
-    description: Optional[str] = None
-
-
-class ShellTaskResponse(AIRBaseModel):
-    """Response model for shell task assignment (legacy)."""
-    
-    task_id: str
-    endpoint_interactions: List[ShellInteraction]
-    success_count: int
-    failure_count: int
-    total_count: int
+"""
+Interact API models for the Binalyze AIR SDK.
+"""
+
+from typing import List, Optional, Dict, Any
+from datetime import datetime
+from enum import Enum
+from pydantic import Field
+
+from ..base import AIRBaseModel
+
+
+class SendToLocation(str, Enum):
+    """Send to location options."""
+    USER_LOCAL = "user-local"
+    REPOSITORY = "repository"
+    EVIDENCE_REPOSITORY = "evidence-repository"
+
+
+class TaskConfigChoice(str, Enum):
+    """Task configuration choice options."""
+    USE_POLICY = "use-policy"
+    USE_CUSTOM_OPTIONS = "use-custom-options"
+
+
+class SendToConfig(AIRBaseModel):
+    """Send to configuration model."""
+    
+    location: SendToLocation
+    repository_id: Optional[str] = Field(default=None, alias="repositoryId")
+    evidence_repository_id: Optional[str] = Field(default=None, alias="evidenceRepositoryId")
+
+
+class BandwidthConfig(AIRBaseModel):
+    """Bandwidth configuration model."""
+    
+    limit: Optional[int] = None
+
+
+class DiskSpaceConfig(AIRBaseModel):
+    """Disk space configuration model."""
+    
+    reserve: Optional[int] = None
+
+
+class TaskConfig(AIRBaseModel):
+    """Task configuration model."""
+    
+    choice: TaskConfigChoice
+    send_to: Optional[SendToConfig] = Field(default=None, alias="sendTo")
+    bandwidth: Optional[BandwidthConfig] = None
+    disk_space: Optional[DiskSpaceConfig] = Field(default=None, alias="diskSpace")
+
+
+class AssignInteractiveShellTaskRequest(AIRBaseModel):
+    """Request model for assigning interactive shell task."""
+    
+    asset_id: str = Field(alias="assetId")
+    case_id: str = Field(alias="caseId")
+    task_config: TaskConfig = Field(alias="taskConfig")
+
+
+class InteractiveShellTaskResponse(AIRBaseModel):
+    """Response model for interactive shell task assignment."""
+    
+    session_id: str = Field(alias="sessionId")
+    idle_timeout: int = Field(alias="idleTimeout")
+    config: TaskConfig
+
+
+# NEW MODELS FOR INTERACT API
+
+class LibraryFile(AIRBaseModel):
+    """Library file model."""
+    
+    id: str = Field(alias="_id")
+    name: str
+    organization_ids: List[int] = Field(alias="organizationIds")
+    size: int
+    sha256: str
+    uploaded_by: str = Field(alias="uploadedBy")
+    uploaded_at: datetime = Field(alias="uploadedAt")
+    last_used_at: Optional[datetime] = Field(default=None, alias="lastUsedAt")
+    last_used_by: Optional[str] = Field(default=None, alias="lastUsedBy")
+
+
+class LibraryFileFilter(AIRBaseModel):
+    """Library file filter model."""
+    
+    search_term: Optional[str] = Field(default=None, alias="searchTerm")
+    organization_ids: Optional[List[int]] = Field(default=None, alias="organizationIds")
+    name: Optional[str] = None
+    uploaded_by: Optional[str] = Field(default=None, alias="uploadedBy")
+    uploaded_at: Optional[str] = Field(default=None, alias="uploadedAt")
+    
+    def to_params(self) -> Dict[str, Any]:
+        """Convert filter to query parameters."""
+        params = {}
+        if self.search_term:
+            params["filter[searchTerm]"] = self.search_term
+        if self.organization_ids:
+            params["filter[organizationIds]"] = ",".join([str(x) for x in self.organization_ids])
+        if self.name:
+            params["filter[name]"] = self.name
+        if self.uploaded_by:
+            params["filter[uploadedBy]"] = self.uploaded_by
+        if self.uploaded_at:
+            params["filter[uploadedAt]"] = self.uploaded_at
+        return params
+
+
+class ExecuteCommandRequest(AIRBaseModel):
+    """Request model for executing interact command."""
+    
+    command: str
+    accept: Optional[str] = "json"  # json or text
+
+
+class ExecuteCommandResponse(AIRBaseModel):
+    """Response model for execute command."""
+    
+    message_id: str = Field(alias="messageId")
+    session: Dict[str, str]
+    body: str
+    cwd: str
+    exit_code: int = Field(alias="exitCode")
+
+
+class InteractCommandOption(AIRBaseModel):
+    """Interact command option model."""
+    
+    name: str
+    abbreviations: List[str]
+    effect: str
+    required: bool
+
+
+class InteractCommand(AIRBaseModel):
+    """Interact command model."""
+    
+    command: List[str]
+    description: str
+    options: List[InteractCommandOption]
+    privilege: str
+    usage: List[str]
+
+
+class SessionInfo(AIRBaseModel):
+    """Session information model."""
+    
+    id: str
+
+
+class CommandMessage(AIRBaseModel):
+    """Command message model."""
+    
+    message_id: str = Field(alias="messageId")
+    session: SessionInfo
+    body: str
+    cwd: str
+    exit_code: int = Field(alias="exitCode")
+
+
+class InterruptCommandRequest(AIRBaseModel):
+    """Request model for interrupting command."""
+    pass  # Empty request body
+
+
+class CloseSessionRequest(AIRBaseModel):
+    """Request model for closing session."""
+    pass  # Empty request body
+
+
+class UploadFileRequest(AIRBaseModel):
+    """Request model for uploading file to library."""
+    
+    file_content: bytes
+    filename: str
+    organization_ids: Optional[List[int]] = Field(default=None, alias="organizationIds")
+
+
+class DeleteFileRequest(AIRBaseModel):
+    """Request model for deleting file from library."""
+    
+    file_id: str = Field(alias="fileId")
+
+
+class CheckFileExistsRequest(AIRBaseModel):
+    """Request model for checking if file exists."""
+    
+    filename: str
+    sha256: Optional[str] = None
+
+
+class FileExistsResponse(AIRBaseModel):
+    """Response model for file exists check."""
+    
+    exists: bool
+    file_id: Optional[str] = Field(default=None, alias="fileId")
+
+
+# Legacy models for backward compatibility (deprecated)
+class InteractionType(str, Enum):
+    """Interaction types."""
+    SHELL = "shell"
+    POWERSHELL = "powershell"
+    CMD = "cmd"
+    BASH = "bash"
+
+
+class InteractionStatus(str, Enum):
+    """Interaction status."""
+    PENDING = "pending"
+    RUNNING = "running"
+    COMPLETED = "completed"
+    FAILED = "failed"
+    CANCELLED = "cancelled"
+    TIMEOUT = "timeout"
+
+
+class ShellInteraction(AIRBaseModel):
+    """Shell interaction model (legacy)."""
+    
+    id: str
+    task_id: str
+    endpoint_id: str
+    endpoint_name: str
+    interaction_type: InteractionType
+    command: str
+    output: Optional[str] = None
+    error_output: Optional[str] = None
+    exit_code: Optional[int] = None
+    status: InteractionStatus = InteractionStatus.PENDING
+    timeout: int = 300  # seconds
+    organization_id: int
+    created_by: Optional[str] = None
+    created_at: Optional[datetime] = None
+    started_at: Optional[datetime] = None
+    completed_at: Optional[datetime] = None
+    duration: Optional[int] = None  # seconds
+    environment_variables: Optional[Dict[str, str]] = None
+    working_directory: Optional[str] = None
+
+
+class AssignShellTaskRequest(AIRBaseModel):
+    """Request model for assigning shell interaction tasks (legacy)."""
+    
+    endpoint_ids: List[str]
+    command: str
+    interaction_type: InteractionType = InteractionType.SHELL
+    timeout: Optional[int] = 300
+    organization_ids: Optional[List[int]] = None
+    case_id: Optional[str] = None
+    environment_variables: Optional[Dict[str, str]] = None
+    working_directory: Optional[str] = None
+    description: Optional[str] = None
+
+
+class ShellTaskResponse(AIRBaseModel):
+    """Response model for shell task assignment (legacy)."""
+    
+    task_id: str
+    endpoint_interactions: List[ShellInteraction]
+    success_count: int
+    failure_count: int
+    total_count: int
     errors: Optional[List[Dict[str, Any]]] = None 

binalyze_air/models/investigation_hub.py

@@ -0,0 +1,265 @@
+"""
+Investigation Hub models for the Binalyze AIR SDK.
+"""
+
+from typing import Optional, List, Dict, Any
+from datetime import datetime
+from pydantic import Field
+
+from ..base import AIRBaseModel
+
+
+class InvestigationMeta(AIRBaseModel):
+    """Investigation metadata model."""
+    
+    case_id: Optional[str] = Field(default=None, alias="caseId")
+    disk_usage_in_bytes: Optional[int] = Field(default=None, alias="diskUsageInBytes")
+
+
+class Investigation(AIRBaseModel):
+    """Investigation model."""
+    
+    uid: str
+    investigation_schema: str = Field(alias="schema")
+    type: str
+    meta: Optional[InvestigationMeta] = None
+    timezone: str = "UTC"
+    organization_id: int = Field(alias="organizationId")
+    last_accessed_at: Optional[datetime] = Field(default=None, alias="lastAccessedAt")
+    status: str
+    migrations_completed: bool = Field(default=True, alias="migrationsCompleted")
+
+
+class UpdateInvestigationRequest(AIRBaseModel):
+    """Request model for updating investigations."""
+    
+    timezone: Optional[str] = None
+    meta: Optional[Dict[str, Any]] = None
+
+
+class InvestigationAsset(AIRBaseModel):
+    """Investigation asset model."""
+    
+    id: str = Field(alias="_id")
+    name: str
+    platform: str
+    ip_address: Optional[str] = Field(default=None, alias="ipAddress")
+    organization_id: int = Field(alias="organizationId")
+
+
+class FlagSummary(AIRBaseModel):
+    """Flag summary model."""
+    
+    id: int
+    name: str
+    color: str
+    count: int
+
+
+class EvidenceSection(AIRBaseModel):
+    """Evidence section model."""
+    
+    name: str
+    display_name: str = Field(alias="displayName")
+    table_count: int = Field(alias="tableCount")
+    record_count: int = Field(alias="recordCount")
+
+
+class EvidenceStructure(AIRBaseModel):
+    """Evidence structure model."""
+    
+    name: str
+    type: str
+    nullable: bool = Field(default=True)
+    primary_key: bool = Field(default=False, alias="primaryKey")
+
+
+class SQLQueryRequest(AIRBaseModel):
+    """SQL query request model."""
+    
+    query: str
+    page_size: int = Field(default=10, alias="pageSize")
+    page_number: int = Field(default=1, alias="pageNumber")
+
+
+class SQLQueryResult(AIRBaseModel):
+    """SQL query result model."""
+    
+    entities: List[Dict[str, Any]]
+    total_entity_count: int = Field(alias="totalEntityCount")
+    current_page: int = Field(alias="currentPage")
+    page_size: int = Field(alias="pageSize")
+    next_page: int = Field(alias="nextPage")
+    total_page_count: int = Field(alias="totalPageCount")
+    previous_page: int = Field(alias="previousPage")
+
+
+class FlagEvidenceRequest(AIRBaseModel):
+    """Flag evidence request model."""
+    
+    flag_id: int = Field(alias="flagId")
+    records: List[Dict[str, Any]]
+    section: str
+
+
+class EvidenceNote(AIRBaseModel):
+    """Evidence note model."""
+    
+    id: str
+    content: str
+    created_at: datetime = Field(alias="createdAt")
+    created_by: str = Field(alias="createdBy")
+
+
+class FindingsSummary(AIRBaseModel):
+    """Findings summary model."""
+    
+    total_findings: int = Field(alias="totalFindings")
+    critical_count: int = Field(alias="criticalCount")
+    high_count: int = Field(alias="highCount")
+    medium_count: int = Field(alias="mediumCount")
+    low_count: int = Field(alias="lowCount")
+    info_count: int = Field(alias="infoCount")
+
+
+class FindingsStructure(AIRBaseModel):
+    """Findings structure model."""
+    
+    flagged_by_list: List[str] = Field(alias="flaggedByList")
+    tactics: List[str]
+    tactic_ids: List[str] = Field(alias="tacticIds")
+    techniques: List[str]
+    technique_ids: List[str] = Field(alias="techniqueIds")
+    evidence_categories: List[Dict[str, str]] = Field(alias="evidenceCategories")
+    extra_columns: List[Dict[str, Any]] = Field(alias="extraColumns")
+    columns_structure: List[Dict[str, Any]] = Field(alias="columnsStructure")
+
+
+class FindingsFilter(AIRBaseModel):
+    """Findings filter model."""
+    
+    assignment_ids: Optional[List[str]] = Field(default=None, alias="assignmentIds")
+    flag_ids: Optional[List[int]] = Field(default=None, alias="flagIds")
+    verdict_scores: Optional[List[str]] = Field(default=None, alias="verdictScores")
+    created_by: Optional[List[str]] = Field(default=None, alias="createdBy")
+    mitre_technique_ids: Optional[List[str]] = Field(default=None, alias="mitreTechniqueIds")
+    mitre_tactic_ids: Optional[List[str]] = Field(default=None, alias="mitreTacticIds")
+
+
+class FindingsRequest(AIRBaseModel):
+    """Findings request model."""
+    
+    take: int = Field(default=50)
+    skip: int = Field(default=0)
+    filter: Optional[List[Dict[str, Any]]] = Field(default=None)
+    global_filter: Optional[FindingsFilter] = Field(default=None, alias="globalFilter")
+    sort: Optional[List[Dict[str, str]]] = Field(default=None)
+
+
+class FindingsResult(AIRBaseModel):
+    """Findings result model."""
+    
+    entities: List[Dict[str, Any]]
+    total_count: int = Field(alias="totalCount")
+    total_count_with_no_filter: int = Field(alias="totalCountWithNoFilter")
+
+
+class MitreMatch(AIRBaseModel):
+    """MITRE ATT&CK match model."""
+    
+    technique_id: str = Field(alias="techniqueId")
+    technique_name: str = Field(alias="techniqueName")
+    tactic: str
+    confidence: float
+
+
+class InvestigationComment(AIRBaseModel):
+    """Investigation comment model."""
+    
+    id: str
+    content: str
+    evidence_id: Optional[str] = Field(default=None, alias="evidenceId")
+    created_at: datetime = Field(alias="createdAt")
+    created_by: str = Field(alias="createdBy")
+
+
+class EvidenceItem(AIRBaseModel):
+    """Evidence item for comments."""
+    
+    evidence: str
+    task_assignment_id: str = Field(alias="taskAssignmentId")
+    object_id: int = Field(alias="objectId")
+
+
+class AddNoteRequest(AIRBaseModel):
+    """Add note request model."""
+    
+    items: List[EvidenceItem]
+    note: str
+    source: str = "EVIDENCE"
+
+
+class CreateCommentRequest(AIRBaseModel):
+    """Create comment request model."""
+    
+    items: List[EvidenceItem]
+    content: str
+    mentioned_usernames: List[str] = Field(default_factory=list, alias="mentionedUsernames")
+    source: str = "EVIDENCE"
+
+
+class InvestigationActivity(AIRBaseModel):
+    """Investigation activity model."""
+    
+    id: str
+    type: str
+    description: str
+    user_id: str = Field(alias="userId")
+    created_at: datetime = Field(alias="createdAt")
+    read: bool = Field(default=False)
+
+
+class MarkActivityAsReadRequest(AIRBaseModel):
+    """Mark activity as read request model."""
+    
+    mark_all: bool = Field(default=False, alias="markAll")
+    activity_id: Optional[int] = Field(default=None, alias="activityId")
+
+
+class AdvancedFilter(AIRBaseModel):
+    """Advanced filter model."""
+    
+    id: int  # API returns integer ID
+    name: str
+    organization_id: int = Field(alias="organizationId")
+    filter: Optional[Dict[str, Any]] = Field(default=None)  # Match API response
+    table_name: Optional[str] = Field(default=None, alias="tableName")
+    created_at: datetime = Field(alias="createdAt")
+    created_by: Optional[str] = Field(default=None, alias="createdBy")
+    updated_at: Optional[datetime] = Field(default=None, alias="updatedAt")
+
+
+class CreateAdvancedFilterRequest(AIRBaseModel):
+    """Create advanced filter request model."""
+    
+    name: str
+    filter: Dict[str, Any]  # Contains main and sections
+    table_name: str = Field(alias="tableName")
+    organization_id: int = Field(alias="organizationId")
+
+
+class UpdateAdvancedFilterRequest(AIRBaseModel):
+    """Update advanced filter request model."""
+    
+    name: Optional[str] = Field(default=None)
+    filter: Optional[Dict[str, Any]] = Field(default=None)  # Contains main and sections
+    table_name: Optional[str] = Field(default=None, alias="tableName")
+    organization_id: Optional[int] = Field(default=None, alias="organizationId")
+
+
+class ExportRequest(AIRBaseModel):
+    """Export request model."""
+    
+    format: str = "csv"  # csv, json, xlsx
+    filters: Optional[Dict[str, Any]] = None
+    columns: Optional[List[str]] = None