iotsploit-cli 0.0.6__py3-none-any.whl

iotsploit_cli/commands/django_commands.py

@@ -0,0 +1,254 @@
+#!/usr/bin/env python
+
+import cmd2
+from cmd2 import ansi
+import sys
+import subprocess
+import time
+from .base_commands import BaseCommands
+from iotsploit_core.utils import iots_logger
+import os
+import signal
+import socket
+from typing import Tuple
+try:
+    import redis  # type: ignore
+except Exception:  # pragma: no cover
+    redis = None
+
+from django.conf import settings
+
+logger = iots_logger.get_logger(__name__)
+
+
+class DjangoCommands(BaseCommands):
+    """Django-related commands for the SAT Shell"""
+
+    def _check_redis_available(self) -> Tuple[bool, str]:
+        """
+        Preflight check for Redis reachability using Django settings.
+
+        Returns (ok, message). If ok is False, message contains human-friendly guidance.
+        """
+        redis_host = getattr(settings, 'REDIS_HOST', '127.0.0.1')
+        redis_port = int(getattr(settings, 'REDIS_PORT', 6379))
+        redis_db = int(getattr(settings, 'REDIS_DB', 0))
+
+        # Quick TCP reachability check first to avoid long client timeouts
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+            sock.settimeout(0.5)
+            try:
+                sock.connect((redis_host, redis_port))
+            except Exception:
+                guide = (
+                    f"Redis is not reachable at {redis_host}:{redis_port}.\n"
+                    "Celery and WebSocket features require Redis.\n\n"
+                    "How to start Redis:\n"
+                    "  - If you use Docker: docker run --name redis -p 6379:6379 -d redis:latest\n"
+                    "  - On Ubuntu/Debian: sudo apt install redis-server && sudo systemctl start redis-server\n"
+                    "  - Or inside project (docker-compose): enable the redis service\n\n"
+                    "Please start Redis and try again."
+                )
+                return False, guide
+
+        # If redis-py is available, ping for certainty
+        if redis is not None:
+            try:
+                client = redis.Redis(host=redis_host, port=redis_port, db=redis_db)
+                client.ping()
+            except Exception as e:  # pragma: no cover
+                return False, (
+                    f"Connected to {redis_host}:{redis_port} but Redis ping failed: {e}\n"
+                    "Please ensure Redis is healthy and try again."
+                )
+
+        return True, "Redis is available"
+
+    @cmd2.with_category('Django Commands')
+    def do_runserver(self, arg):
+        'Start Django development server, Daphne WebSocket server, MCP WebSocket bridge, and Celery worker in the background'
+        if self.django_server_process or self.daphne_server_process:
+            self.poutput("Servers are already running.")
+            return
+
+        try:
+            logger.info("Attempting to start Django, Daphne, MCP WebSocket bridge, and Celery servers in background...")
+            
+            # Preflight: check Redis and fail fast if unavailable
+            redis_ok, redis_msg = self._check_redis_available()
+            if not redis_ok:
+                self.poutput(ansi.style("\n[ERROR] Redis is unavailable. Cannot start services that depend on it.", fg=ansi.Fg.RED, bold=True))
+                self.poutput(redis_msg)
+                return False
+            
+            # Prepare the commands
+            django_cmd = [sys.executable, '-m', 'django', 'runserver', '--noreload', '0.0.0.0:8888']
+            daphne_cmd = [
+                sys.executable, 
+                '-m', 
+                'daphne', 
+                '-b', 
+                '0.0.0.0', 
+                '-p', 
+                '9999', 
+                'iotsploit_django.asgi:application'
+            ]
+            mcp_bridge_cmd = [
+                sys.executable,
+                '-m',
+                'iotsploit_mcp.cli',
+                'ws',
+                '--host',
+                '0.0.0.0',
+                '--port',
+                '9998',
+            ]
+            celery_cmd = [
+                sys.executable,
+                '-m',
+                'celery',
+                '-A',
+                'iotsploit_django.tasks.celery_app:app',
+                'worker',
+                '--loglevel=info'
+            ]
+            
+            logger.info(f"Running Django command: {' '.join(django_cmd)}")
+            logger.info(f"Running Daphne command: {' '.join(daphne_cmd)}")
+            logger.info(f"Running MCP WebSocket Bridge command: {' '.join(mcp_bridge_cmd)}")
+            logger.info(f"Running Celery command: {' '.join(celery_cmd)}")
+            
+            # Start the processes with direct output to stdout/stderr
+            self.django_server_process = subprocess.Popen(
+                django_cmd, 
+                stdout=sys.stdout,  # 直接输出到控制台
+                stderr=sys.stderr,
+                universal_newlines=True
+            )
+            
+            self.daphne_server_process = subprocess.Popen(
+                daphne_cmd, 
+                stdout=sys.stdout,  # 直接输出到控制台
+                stderr=sys.stderr,
+                universal_newlines=True
+            )
+            
+            # Start the WebSocket bridge in its own process group so that we can
+            # later terminate the entire group (bridge + sat_fastmcp child)
+            # Set up environment variables for MCP bridge (Django API URL)
+            mcp_env = os.environ.copy()
+            mcp_env.setdefault('IOTSPLOIT_DJANGO_API_BASE_URL', 'http://127.0.0.1:8888')
+            
+            self.mcp_bridge_process = subprocess.Popen(
+                mcp_bridge_cmd,
+                stdout=sys.stdout,  # 直接输出到控制台
+                stderr=sys.stderr,
+                universal_newlines=True,
+                start_new_session=True,  # create new session = new PGID on POSIX
+                env=mcp_env
+            )
+            
+            self.celery_worker_process = subprocess.Popen(
+                celery_cmd,
+                stdout=sys.stdout,  # 直接输出到控制台
+                stderr=sys.stderr,
+                universal_newlines=True
+            )
+            
+            logger.info("All servers started successfully in the background.")
+            logger.info("Services running on:")
+            logger.info("  - Django HTTP API: http://localhost:8888")
+            logger.info("  - Daphne WebSocket: ws://localhost:9999")
+            logger.info("  - MCP WebSocket Bridge: ws://localhost:9998")
+            logger.info("  - Celery Worker: background task processing")
+            
+            # Wait for HTTP server to be available and initialize devices
+            import requests
+            max_retries = 30
+            retry_interval = 1
+            
+            logger.info("Waiting for HTTP server to be available...")
+            for i in range(max_retries):
+                try:
+                    # Try to initialize devices using the HTTP endpoint (GET method)
+                    response = requests.get('http://127.0.0.1:8888/api/initialize_devices/')
+                    if response.status_code == 200:
+                        logger.info("Devices initialized successfully via HTTP API")
+                        break
+                    else:
+                        logger.error(f"Failed to initialize devices: {response.text}")
+                        break
+                except requests.exceptions.ConnectionError:
+                    if i < max_retries - 1:
+                        time.sleep(retry_interval)
+                    else:
+                        logger.error("HTTP server did not become available in time")
+                    continue
+                except Exception as e:
+                    logger.error(f"Error initializing devices: {str(e)}")
+                    break
+            
+            return None
+        
+        except Exception as e:
+            logger.error(f"Failed to start servers: {str(e)}")
+            logger.debug("Detailed traceback:", exc_info=True)
+            return False
+
+    @cmd2.with_category('Django Commands')
+    def do_stop_server(self, arg):
+        'Stop Django development server, Daphne WebSocket server, MCP WebSocket bridge, and Celery worker'
+        try:
+            # Cleanup devices using HTTP endpoint (GET method)
+            # Only attempt HTTP cleanup if the Django server process is still alive
+            if self.django_server_process and self.django_server_process.poll() is None:
+                import requests
+                try:
+                    response = requests.get('http://127.0.0.1:8888/api/cleanup_devices/')
+                    if response.status_code == 200:
+                        logger.info("Devices cleaned up successfully via HTTP API")
+                    else:
+                        logger.error(f"Failed to cleanup devices: {response.text}")
+                except requests.exceptions.ConnectionError:
+                    logger.warning("Could not reach HTTP server for device cleanup")
+                except Exception as e:
+                    logger.error(f"Error during device cleanup: {str(e)}")
+            else:
+                logger.debug("Django server not running, skipping HTTP device cleanup")
+            
+            # Stop the servers
+            if self.django_server_process:
+                self.django_server_process.terminate()
+                self.django_server_process = None
+            
+            if self.daphne_server_process:
+                self.daphne_server_process.terminate()
+                self.daphne_server_process = None
+            
+            if self.mcp_bridge_process:
+                try:
+                    # Terminate the entire process group so that child processes
+                    # such as sat_fastmcp are also stopped.
+                    os.killpg(os.getpgid(self.mcp_bridge_process.pid), signal.SIGTERM)
+                except ProcessLookupError:
+                    # Process already gone
+                    pass
+                except Exception as e:
+                    logger.error(f"Failed to terminate MCP bridge process group: {e}")
+                finally:
+                    self.mcp_bridge_process = None
+            
+            if self.celery_worker_process:
+                self.celery_worker_process.terminate()
+                self.celery_worker_process = None
+            
+            if not any([self.django_server_process, self.daphne_server_process, 
+                        getattr(self, 'mcp_bridge_process', None),
+                        getattr(self, 'celery_worker_process', None)]):
+                logger.info("All servers stopped.")
+            else:
+                logger.error("No servers were running.")
+                
+        except Exception as e:
+            logger.error(f"Error stopping servers: {str(e)}")
+            logger.debug("Detailed error:", exc_info=True)

iotsploit_cli/commands/firmware_commands.py

@@ -0,0 +1,213 @@
+#!/usr/bin/env python
+
+import cmd2
+from cmd2 import ansi
+from .base_commands import BaseCommands
+from iotsploit_django.tools.input_mgr import Input_Mgr
+from iotsploit_core.core.tool_service import get_firmware_service
+from iotsploit_core.utils import iots_logger
+
+logger = iots_logger.get_logger(__name__)
+
+
+class FirmwareCommands(BaseCommands):
+    """Firmware-related commands for the SAT Shell"""
+
+    @cmd2.with_category('Firmware Commands')
+    def do_list_firmware(self, arg):
+        'List all available firmware'
+        try:
+            firmware_service = get_firmware_service()
+            firmware_list = firmware_service.list_firmware()
+            
+            if not firmware_list:
+                self.poutput(ansi.style("No firmware found in registry.", fg=ansi.Fg.YELLOW))
+                return
+            
+            self.poutput(ansi.style(f"\nFound {len(firmware_list)} firmware(s):", fg=ansi.Fg.GREEN))
+            self.poutput("-" * 80)
+            
+            for firmware in firmware_list:
+                name = firmware['name']
+                device_type = firmware.get('device_type', 'unknown')
+                version = firmware.get('version', 'unknown')
+                path = firmware.get('path', 'unknown')
+                
+                # Check if file exists
+                from pathlib import Path
+                file_exists = Path(path).exists() if path != 'unknown' else False
+                status_color = ansi.Fg.GREEN if file_exists else ansi.Fg.RED
+                status_text = "✓" if file_exists else "✗"
+                
+                self.poutput(ansi.style(f"  {name}:", fg=ansi.Fg.CYAN, bold=True))
+                self.poutput(f"    Device Type: {device_type}")
+                self.poutput(f"    Version: {version}")
+                self.poutput(f"    Path: {path}")
+                self.poutput(f"    Status: {ansi.style(status_text, fg=status_color)} {'(File exists)' if file_exists else '(File missing)'}")
+                
+                # Show flash options if available
+                flash_options = firmware.get('flash_options', {})
+                if flash_options:
+                    self.poutput(f"    Flash Options: {flash_options}")
+                
+                self.poutput("")
+
+        except Exception as e:
+            logger.error(f"Error listing firmware: {str(e)}")
+            self.poutput(ansi.style(f"Error listing firmware: {str(e)}", fg=ansi.Fg.RED))
+
+    do_lsfw = do_list_firmware
+
+    @cmd2.with_category('Firmware Commands')
+    def do_add_firmware(self, arg):
+        'Add new firmware to the system'
+        try:
+            # Parse arguments
+            args = arg.split()
+            if len(args) < 4:
+                self.poutput(ansi.style("Usage: add_firmware <name> <path> <device_type> <version> [flash_options_json]", fg=ansi.Fg.YELLOW))
+                self.poutput("Example: add_firmware my_esp32 /path/to/firmware.bin esp32 1.0.0")
+                return
+            
+            name = args[0]
+            path = args[1]
+            device_type = args[2]
+            version = args[3]
+            
+            # Parse optional flash options
+            flash_options = None
+            if len(args) > 4:
+                import json
+                try:
+                    flash_options = json.loads(' '.join(args[4:]))
+                except json.JSONDecodeError:
+                    self.poutput(ansi.style("Invalid JSON format for flash options", fg=ansi.Fg.RED))
+                    return
+            
+            firmware_service = get_firmware_service()
+            success = firmware_service.add_firmware(name, path, device_type, version, flash_options)
+            
+            if success:
+                self.poutput(ansi.style(f"Successfully added firmware: {name}", fg=ansi.Fg.GREEN))
+            else:
+                self.poutput(ansi.style(f"Failed to add firmware: {name}", fg=ansi.Fg.RED))
+                
+        except Exception as e:
+            logger.error(f"Error adding firmware: {str(e)}")
+            self.poutput(ansi.style(f"Error adding firmware: {str(e)}", fg=ansi.Fg.RED))
+
+    do_addfw = do_add_firmware
+
+    @cmd2.with_category('Firmware Commands')
+    def do_flash_firmware(self, arg):
+        'Flash firmware to a device'
+        try:
+            # Parse arguments
+            args = arg.split()
+            if len(args) < 2:
+                self.poutput(ansi.style("Usage: flash_firmware <firmware_name> <device_name> [options_json]", fg=ansi.Fg.YELLOW))
+                self.poutput("Example: flash_firmware esp32s3_wifi_penetration_tool esp32")
+                return
+            
+            firmware_name = args[0]
+            device_name = args[1]
+            
+            # Parse optional options
+            options = None
+            if len(args) > 2:
+                import json
+                try:
+                    options = json.loads(' '.join(args[2:]))
+                except json.JSONDecodeError:
+                    self.poutput(ansi.style("Invalid JSON format for options", fg=ansi.Fg.RED))
+                    return
+            
+            firmware_service = get_firmware_service()
+            
+            # Check if firmware exists
+            firmware_info = firmware_service.get_firmware_info(firmware_name)
+            if not firmware_info:
+                self.poutput(ansi.style(f"Firmware '{firmware_name}' not found", fg=ansi.Fg.RED))
+                return
+            
+            self.poutput(ansi.style(f"Flashing firmware '{firmware_name}' to device '{device_name}'...", fg=ansi.Fg.YELLOW))
+            
+            # Flash the firmware
+            success = firmware_service.flash_registered_firmware(firmware_name, options)
+            
+            if success:
+                self.poutput(ansi.style(f"Successfully flashed firmware: {firmware_name}", fg=ansi.Fg.GREEN))
+            else:
+                self.poutput(ansi.style(f"Failed to flash firmware: {firmware_name}", fg=ansi.Fg.RED))
+                
+        except Exception as e:
+            logger.error(f"Error flashing firmware: {str(e)}")
+            self.poutput(ansi.style(f"Error flashing firmware: {str(e)}", fg=ansi.Fg.RED))
+
+    do_flashfw = do_flash_firmware
+
+    @cmd2.with_category('Firmware Commands')
+    def do_remove_firmware(self, arg):
+        'Remove firmware from the system'
+        try:
+            if not arg.strip():
+                self.poutput(ansi.style("Usage: remove_firmware <firmware_name>", fg=ansi.Fg.YELLOW))
+                self.poutput("Example: remove_firmware esp32_blink")
+                return
+            
+            firmware_name = arg.strip()
+            firmware_service = get_firmware_service()
+            
+            # Check if firmware exists
+            firmware_info = firmware_service.get_firmware_info(firmware_name)
+            if not firmware_info:
+                self.poutput(ansi.style(f"Firmware '{firmware_name}' not found", fg=ansi.Fg.RED))
+                return
+            
+            # Confirm removal
+            confirm = input(f"Are you sure you want to remove firmware '{firmware_name}'? (y/N): ")
+            if confirm.lower() != 'y':
+                self.poutput(ansi.style("Firmware removal cancelled", fg=ansi.Fg.YELLOW))
+                return
+            
+            success = firmware_service.remove_firmware(firmware_name)
+            
+            if success:
+                self.poutput(ansi.style(f"Successfully removed firmware: {firmware_name}", fg=ansi.Fg.GREEN))
+            else:
+                self.poutput(ansi.style(f"Failed to remove firmware: {firmware_name}", fg=ansi.Fg.RED))
+                
+        except Exception as e:
+            logger.error(f"Error removing firmware: {str(e)}")
+            self.poutput(ansi.style(f"Error removing firmware: {str(e)}", fg=ansi.Fg.RED))
+
+    do_rmfw = do_remove_firmware
+
+    @cmd2.with_category('Firmware Commands')
+    def do_download_firmware(self, arg):
+        'Download firmware from URL'
+        try:
+            if not arg.strip():
+                self.poutput(ansi.style("Usage: download_firmware <url> [output_path]", fg=ansi.Fg.YELLOW))
+                self.poutput("Example: download_firmware https://example.com/firmware.bin")
+                return
+            
+            args = arg.split()
+            url = args[0]
+            output_path = args[1] if len(args) > 1 else None
+            
+            firmware_service = get_firmware_service()
+            self.poutput(ansi.style(f"Downloading firmware from: {url}", fg=ansi.Fg.YELLOW))
+            
+            result_path = firmware_service.download_firmware(url, output_path)
+            
+            if result_path:
+                self.poutput(ansi.style(f"Successfully downloaded firmware to: {result_path}", fg=ansi.Fg.GREEN))
+            else:
+                self.poutput(ansi.style("Failed to download firmware", fg=ansi.Fg.RED))
+                
+        except Exception as e:
+            logger.error(f"Error downloading firmware: {str(e)}")
+            self.poutput(ansi.style(f"Error downloading firmware: {str(e)}", fg=ansi.Fg.RED))
+
+    do_dlfw = do_download_firmware

iotsploit_cli/commands/linux_commands.py

@@ -0,0 +1,35 @@
+#!/usr/bin/env python
+
+import cmd2
+from cmd2 import ansi
+import subprocess
+from .base_commands import BaseCommands
+from iotsploit_core.utils import iots_logger
+
+logger = iots_logger.get_logger(__name__)
+
+
+class LinuxCommands(BaseCommands):
+    """Linux system commands for the SAT Shell"""
+
+    @cmd2.with_category('Linux Commands')
+    def do_ls(self, arg):
+        'List directory contents'
+        try:
+            result = subprocess.run(['ls'] + arg.split(), capture_output=True, text=True)
+            self.poutput(result.stdout)
+            if result.stderr:
+                self.poutput(result.stderr)
+        except Exception as e:
+            logger.error(f"Error executing ls: {str(e)}")
+
+    @cmd2.with_category('Linux Commands')
+    def do_lsusb(self, arg):
+        'List USB devices'
+        try:
+            result = subprocess.run(['lsusb'] + arg.split(), capture_output=True, text=True)
+            self.poutput(result.stdout)
+            if result.stderr:
+                self.poutput(result.stderr)
+        except Exception as e:
+            logger.error(f"Error executing lsusb: {str(e)}")

iotsploit_cli/commands/network_commands.py

@@ -0,0 +1,36 @@
+#!/usr/bin/env python
+
+import cmd2
+from cmd2 import ansi
+import time
+from .base_commands import BaseCommands
+from iotsploit_django.tools.wifi_mgr import WiFi_Mgr
+from iotsploit_django.tools.input_mgr import Input_Mgr
+from iotsploit_core.utils import iots_logger
+
+logger = iots_logger.get_logger(__name__)
+
+
+class NetworkCommands(BaseCommands):
+    """Network-related commands for the SAT Shell"""
+
+    @cmd2.with_category('Network Commands')
+    def do_connect_wifi(self, arg):
+        'Connect to WiFi network by providing SSID and password'
+        logger.info(ansi.style("WiFi Connection Setup", fg=ansi.Fg.CYAN))
+        
+        # Get SSID from user
+        ssid = Input_Mgr.Instance().string_input("Enter WiFi SSID")
+        
+        # Get password from user
+        password = Input_Mgr.Instance().string_input("Enter WiFi password")
+        
+        # Attempt to connect
+        logger.info(ansi.style(f"Attempting to connect to {ssid}...", fg=ansi.Fg.CYAN))
+        WiFi_Mgr.Instance().sta_connect_wifi(ssid, password)
+        
+        # Wait for connection to establish
+        time.sleep(2)
+        
+        # Show connection status
+        WiFi_Mgr.Instance().status()