Skip to content

Architecture Overview

This document provides a comprehensive overview of AI Shell's architecture, design patterns, and component interactions.

🏗️ High-Level Architecture

┌─────────────────────────────────────────────────────┐
│                   User Interface                    │
│                     (ui.py)                        │
└─────────────────────┬───────────────────────────────┘
                      │
┌─────────────────────▼───────────────────────────────┐
│                  Main Application                   │
│                   (main.py)                        │
│  ┌─────────────┬─────────────┬─────────────────────┐│
│  │ Mode        │ Provider    │ Security            ││
│  │ Selection   │ Management  │ Validation          ││
│  └─────────────┴─────────────┴─────────────────────┘│
└─────────────────────┬───────────────────────────────┘
                      │
        ┌─────────────┼─────────────┐
        │             │             │
        ▼             ▼             ▼
┌─────────────┐ ┌──────────────┐ ┌──────────────┐
│Configuration│ │ LLM Provider │ │ Command      │
│ Management  │ │ (llm.py)     │ │ Executor     │
│ (config.py) │ │              │ │ (executor.py)│
└─────────────┘ └──────────────┘ └──────────────┘
        │             │             │
        └─────────────┼─────────────┘
                      │
        ┌─────────────▼─────────────┐
        │     External Services     │
        │                          │
        │ ┌─────────┐ ┌───────────┐ │
        │ │ Gemini  │ │ Ollama    │ │
        │ │ API     │ │ Local LLM │ │
        │ └─────────┘ └───────────┘ │
        │                          │
        │ ┌─────────┐ ┌───────────┐ │
        │ │MSFConsole│ │  System   │ │
        │ │   PTY    │ │ Commands  │ │
        │ └─────────┘ └───────────┘ │
        └──────────────────────────┘

🧩 Core Components

1. Main Application (main.py)

Responsibilities: - Application entry point and CLI argument parsing - Mode selection and workflow orchestration - Provider initialization and management - Interactive terminal handling with PTY support

Key Functions: - main(): Primary application entry point - setup_logging(): Logging configuration - parse_arguments(): CLI argument handling - pty_loop_base(): Generic PTY-based tool integration loop - metasploit_loop() / wapiti_loop(): Mode-specific async PTY wrappers

Design Patterns: - Command Pattern: Mode selection and execution - Factory Pattern: Provider instantiation - Observer Pattern: Real-time output streaming

2. Configuration Management (config.py)

Responsibilities: - YAML-based configuration loading and validation - Environment variable integration - Default value management - Dynamic configuration updates

Key Classes: 

class Config:
    def __init__(self, config_data=None, config_file=None)
    def get(self, key, default=None)
    def set(self, key, value)
    def save(self, filename)

Configuration Structure: 

llm:
  provider: gemini|local
  gemini:
    api_key: string
    model: string
  local:
    host: string
    port: integer
    model: string

security:
  require_confirmation: boolean
  dangerous_commands: list

logging:
  level: DEBUG|INFO|WARNING|ERROR
  file: string
  format: string

training:
  dataset_file: string
  auto_log: boolean

3. LLM Provider System (llm.py)

Architecture: 

# Base Provider Interface
class LLMProvider:
    def generate_response(
        self,
        prompt: str,
        mode: str,
        system_prompt: str = ASSISTANT_SYSTEM_PROMPT,
        chat_session: Any = None,
    ) -> Tuple[Optional[str], Any]:
        raise NotImplementedError

# Concrete Implementations
class GeminiProvider(LLMProvider)      # Google Gemini via google-generativeai
class LocalLLMProvider(LLMProvider)    # Local models via Ollama REST API

Provider Selection Logic: 

def get_llm_provider() -> LLMProvider:
    config = get_config()
    provider_type = config.get("llm.provider", "gemini")
    if provider_type == "gemini":
        return GeminiProvider(config.get("llm.gemini", {}))
    elif provider_type == "local":
        return LocalLLMProvider(config.get("llm.local", {}))

System Prompts: - TRANSLATOR_SYSTEM_PROMPT: Command translation mode - ASSISTANT_SYSTEM_PROMPT: Conversational assistance mode - METASPLOIT_SYSTEM_PROMPT: Security testing mode - WAPITI_SYSTEM_PROMPT: Web application scanning mode

4. Command Execution (executor.py)

Security Architecture: 

class SecurityChecker:
    def is_dangerous_command(self, command: str) -> bool
    def validate_command(self, command: str) -> Tuple[bool, Optional[str]]

class CommandExecutor:
    def execute_command(self, command: str, user_prompt: str) -> bool

class TrainingDataLogger:
    def log_training_pair(self, prompt: str, command: str, feedback: str) -> None

def get_executor() -> CommandExecutor   # Returns global singleton instance

Security Layers: 1. Input Validation: Command syntax and structure validation 2. Dangerous Command Detection: Pattern matching against known dangerous commands 3. User Confirmation: Interactive approval for command execution 4. Output Sanitization: Safe handling of command output

Execution Modes: - Interactive: Real-time output streaming with PTY - Batch: Standard subprocess execution - Dry Run: Validation without execution

5. User Interface (ui.py)

Component Categories: 

# Color and Formatting
colors = {
    'primary': '\033[96m',      # Cyan
    'secondary': '\033[95m',     # Magenta
    'success': '\033[92m',       # Green
    'warning': '\033[93m',       # Yellow
    'error': '\033[91m',         # Red
    'info': '\033[94m',          # Blue
    'reset': '\033[0m'           # Reset
}

# UI Components
def print_banner()
def print_mode_selection()
def print_provider_selection()
def format_command_output(output: str) -> str

Responsive Design: - Terminal width detection - Dynamic content wrapping - Cross-platform color support

🔄 Data Flow

Command Translation Flow

graph TD
    A[User Input] --> B[Input Validation]
    B --> C[LLM Provider Selection]
    C --> D[Generate Response]
    D --> E[Extract Command]
    E --> F[Security Validation]
    F --> G{User Confirmation}
    G -->|Yes| H[Execute Command]
    G -->|No| I[Return to Input]
    H --> J[Stream Output]
    J --> K[Log Training Data]

Configuration Loading Flow

graph TD
    A[Application Start] --> B{Config File Exists?}
    B -->|Yes| C[Load YAML Config]
    B -->|No| D[Use Default Config]
    C --> E[Environment Variable Override]
    D --> E
    E --> F[Validate Configuration]
    F --> G[Initialize Components]

PTY Session Flow

graph TD
    A[Start PTY Session] --> B[Fork Process]
    B --> C[Setup Master/Slave PTY]
    C --> D[Launch Target Application]
    D --> E[Real-time I/O Handling]
    E --> F{User Input or Tool Output?}
    F -->|User Input| G[Process with LLM if needed]
    F -->|Tool Output| H[Display to User]
    G --> I[Send to Tool]
    H --> E
    I --> E

🔒 Security Architecture

Defense in Depth

  1. Input Layer
  2. Command injection prevention
  3. Input sanitization

  4. Length and format validation

  5. Validation Layer

  6. Dangerous command detection
  7. Pattern matching

  8. Whitelist/blacklist enforcement

  9. Execution Layer

  10. User confirmation requirements
  11. Process isolation

  12. Resource limiting

  13. Output Layer

  14. Output sanitization
  15. Sensitive data filtering
  16. Logging and auditing

Threat Model

Threats Addressed: - Command injection attacks - Malicious LLM responses - Privilege escalation - Data exfiltration - Denial of service

Mitigations: - Input validation and sanitization - Command whitelisting/blacklisting - User confirmation workflows - Process sandboxing - Resource monitoring

🧪 Testing Architecture

Test Categories

  1. Unit Tests
  2. Individual component testing
  3. Mock external dependencies

  4. Edge case validation

  5. Integration Tests

  6. Component interaction testing
  7. Configuration validation

  8. Provider integration

  9. Security Tests

  10. Command injection prevention
  11. Dangerous command detection

  12. Input validation

  13. End-to-End Tests

  14. Full workflow testing
  15. User interaction simulation
  16. Real provider integration

Test Structure

tests/
├── unit/
│   ├── test_config.py
│   ├── test_llm.py
│   ├── test_executor.py
│   └── test_ui.py
├── integration/
│   ├── test_provider_integration.py
│   └── test_workflow_integration.py
├── security/
│   ├── test_command_validation.py
│   └── test_injection_prevention.py
└── e2e/
    └── test_full_workflow.py

📊 Performance Considerations

Optimization Strategies

  1. Async Operations
  2. Non-blocking LLM API calls
  3. Concurrent request processing

  4. Real-time response streaming

  5. Caching

  6. Configuration caching
  7. Provider initialization caching

  8. Response pattern caching

  9. Resource Management

  10. Connection pooling for API calls
  11. Memory-efficient output streaming
  12. Process lifecycle management

Scalability

  • Horizontal Scaling: Multiple provider instances
  • Vertical Scaling: Resource optimization
  • Load Balancing: Provider selection strategies

🔮 Future Architecture

Planned Enhancements

  1. Plugin System
  2. Modular tool integration
  3. Third-party provider support

  4. Custom command handlers

  5. Distributed Architecture

  6. Remote LLM provider support
  7. Distributed configuration management

  8. Multi-user support

  9. Advanced Security

  10. Role-based access control
  11. Audit logging

  12. Compliance frameworks

  13. Machine Learning

  14. Local model fine-tuning
  15. Usage pattern analysis
  16. Predictive command suggestions

This architecture enables AI Shell to be maintainable, extensible, and secure while providing a rich user experience for command-line operations.