iotsploit-core 0.0.6__py3-none-any.whl

iotsploit_core/core/device_registry.py

@@ -0,0 +1,118 @@
+from __future__ import annotations
+
+import logging
+from typing import Dict, List, Optional
+
+from iotsploit_core.core.device_config import DeviceConfigManager
+from iotsploit_core.core.device_scanner import CompositeDeviceScanner, PluginDeviceScanner
+from iotsploit_core.core.device_store import DeviceStore
+from iotsploit_core.domain.device import Device
+
+logger = logging.getLogger(__name__)
+
+class DeviceRegistry:
+    _instance = None
+    
+    def __new__(cls, *args, **kwargs):
+        if cls._instance is None:
+            cls._instance = super(DeviceRegistry, cls).__new__(cls)
+            cls._instance._initialized = False
+        return cls._instance
+    
+    def __init__(
+        self,
+        *,
+        driver_manager,
+        device_store: DeviceStore | None = None,
+        config_manager: DeviceConfigManager | None = None,
+        scanner: CompositeDeviceScanner | None = None,
+    ):
+        if not self._initialized:
+            if driver_manager is None:
+                raise TypeError("DeviceRegistry requires 'driver_manager' (inject via adapter factory).")
+            self.device_store = device_store or DeviceStore()
+            self.config_manager = config_manager or DeviceConfigManager()
+            self.driver_manager = driver_manager
+            self.scanner = scanner or CompositeDeviceScanner(self.device_store)
+            self._initialized = True
+
+    @classmethod
+    def get_instance(
+        cls,
+        *,
+        driver_manager,
+        device_store: DeviceStore | None = None,
+        config_manager: DeviceConfigManager | None = None,
+        scanner: CompositeDeviceScanner | None = None,
+    ) -> "DeviceRegistry":
+        """Get singleton instance (requires injection on first initialization)."""
+
+        inst = cls(
+            driver_manager=driver_manager,
+            device_store=device_store,
+            config_manager=config_manager,
+            scanner=scanner,
+        )
+        return inst
+            
+    def initialize(self):
+        """初始化设备注册表"""
+        logger.debug("Initializing device registry...")
+        
+        # 清理现有的扫描器
+        self.scanner.clear_scanners()
+        
+        # 为每个驱动创建扫描器
+        for driver_name, driver in self.driver_manager.drivers.items():
+            scanner = PluginDeviceScanner(self.device_store, driver)
+            self.scanner.add_scanner(scanner)
+            logger.debug(f"Added scanner for driver: {driver_name}")
+            
+        # 加载配置的设备
+        try:
+            stored_configs = self.config_manager.load_configs()
+            if stored_configs.get('devices'):
+                for device_config in stored_configs['devices']:
+                    try:
+                        device = self.device_store._dict_to_device(device_config)
+                        self.device_store.register_device(device, source="static")
+                        logger.debug(f"Loaded static device: {device.device_id}")
+                    except Exception as e:
+                        logger.error(f"Error loading device from config: {e}")
+        except Exception as e:
+            logger.error(f"Error loading stored configurations: {e}")
+            
+    def scan_devices(self) -> List[Device]:
+        """扫描所有可用设备"""
+        try:
+            discovered_devices = self.scanner.scan()
+            # 注册发现的设备
+            for device in discovered_devices:
+                try:
+                    self.device_store.register_device(device, source="dynamic")
+                except Exception as e:
+                    logger.error(f"Error registering discovered device: {e}")
+            return discovered_devices
+        except Exception as e:
+            logger.error(f"Error during device scan: {e}")
+            return []
+        
+    def merge_devices(self, stored_configs: Dict, discovered_devices: List[Device]):
+        """合并存储的配置和发现的设备"""
+        # 处理发现的设备
+        for device in discovered_devices:
+            try:
+                self.device_store.register_device(device, source="dynamic")
+            except Exception as e:
+                logger.error(f"Error registering dynamic device: {e}")
+            
+        # 处理存储的配置
+        if stored_configs.get('devices'):
+            for device_config in stored_configs['devices']:
+                try:
+                    device_id = device_config.get('device_id')
+                    if device_id and device_id not in self.device_store.devices:
+                        device = self.device_store._dict_to_device(device_config)
+                        self.device_store.register_device(device, source="static")
+                except Exception as e:
+                    logger.error(f"Error registering static device: {e}") 

iotsploit_core/core/device_scanner.py

@@ -0,0 +1,74 @@
+from typing import List, Dict, Type
+import logging
+from iotsploit_core.domain.device import Device
+from iotsploit_core.core.device_store import DeviceStore
+from iotsploit_core.core.base_plugin import BaseDeviceDriver
+
+logger = logging.getLogger(__name__)
+
+class DeviceScanner:
+    """设备扫描器基类"""
+    def __init__(self, device_store: DeviceStore):
+        self.device_store = device_store
+        
+    def scan(self) -> List[Device]:
+        """执行设备扫描"""
+        return self._perform_scan()
+        
+    def _perform_scan(self) -> List[Device]:
+        """实际的扫描实现"""
+        raise NotImplementedError("Subclasses must implement _perform_scan()")
+
+class PluginDeviceScanner(DeviceScanner):
+    """使用插件系统的设备扫描器"""
+    def __init__(self, device_store: DeviceStore, driver: BaseDeviceDriver):
+        super().__init__(device_store)
+        self.driver = driver
+        
+    def _perform_scan(self) -> List[Device]:
+        """使用驱动执行扫描"""
+        try:
+            logger.info(f"Scanning devices using driver: {self.driver.__class__.__name__}")
+            result = self.driver.scan()
+            
+            if not result:
+                logger.info("No devices found")
+                return []
+                
+            devices = result if isinstance(result, list) else [result]
+            for device in devices:
+                self.device_store.register_device(device, source="dynamic")
+                
+            return devices
+                
+        except Exception as e:
+            logger.error(f"Error scanning devices: {str(e)}")
+            return []
+
+class CompositeDeviceScanner(DeviceScanner):
+    """组合多个扫描器的复合扫描器"""
+    def __init__(self, device_store: DeviceStore):
+        super().__init__(device_store)
+        self.scanners: List[DeviceScanner] = []
+        
+    def clear_scanners(self):
+        """清理所有已注册的扫描器"""
+        self.scanners = []
+        
+    def add_scanner(self, scanner: DeviceScanner):
+        """添加一个扫描器"""
+        self.scanners.append(scanner)
+        
+    def _perform_scan(self) -> List[Device]:
+        """执行所有注册的扫描器"""
+        all_devices = []
+        for scanner in self.scanners:
+            try:
+                # 直接调用 _perform_scan 避免重复注册
+                devices = scanner._perform_scan()
+                all_devices.extend(devices)
+            except Exception as e:
+                logger.error(f"Error in scanner {scanner.__class__.__name__}: {str(e)}")
+                continue
+                
+        return all_devices 

iotsploit_core/core/device_spec.py

@@ -0,0 +1,105 @@
+import pluggy
+import warnings
+from iotsploit_core.domain.device import Device
+from typing import List, Optional
+from enum import Enum
+
+hookspec = pluggy.HookspecMarker("device_mgr")
+
+"""
+Scan -> Init -> Connect -> Commands -> Close
+"""
+
+
+class DeviceState(Enum):
+    """Device lifecycle states"""
+    UNKNOWN = "unknown"
+    DISCOVERED = "discovered"  # After scan
+    INITIALIZED = "initialized"  # After init
+    CONNECTED = "connected"    # After connect
+    ACTIVE = "active"         # During command execution
+    DISCONNECTED = "disconnected"  # After close
+    ERROR = "error"           # Error state
+
+class DevicePluginSpec:
+    """Specification for device plugin lifecycle management"""
+    
+    @hookspec
+    def scan(self) -> List[Device]:
+        """Scan for available devices.
+        First step in device lifecycle.
+        
+        Returns:
+            list[Device]: A list of discovered devices matching the configuration.
+        """
+
+    @hookspec
+    def initialize(self, device: Device) -> bool:
+        """Initialize the plugin for a specific device.
+        Second step in device lifecycle.
+        
+        Args:
+            device: Device to initialize
+            
+        Returns:
+            bool: True if initialization successful
+        """
+
+    @hookspec 
+    def connect(self, device: Device) -> bool:
+        """Connect to the device.
+        Third step in device lifecycle.
+        
+        Args:
+            device: Device to connect to
+            
+        Returns:
+            bool: True if connection successful
+        """
+
+    @hookspec
+    def command(self, device: Device, command: str) -> Optional[str]:
+        """Send a command to the device.
+        Only available in CONNECTED state.
+        
+        Args:
+            device: Target device
+            command: Command string to execute
+            
+        Returns:
+            Optional[str]: Command response if any
+        """
+
+    @hookspec
+    def reset(self, device: Device) -> bool:
+        """Reset the device to its initial state.
+        Can be called from any state except UNKNOWN.
+        
+        Args:
+            device: Device to reset
+            
+        Returns:
+            bool: True if reset successful
+        """
+
+    @hookspec
+    def close(self, device: Device) -> bool:
+        """Close the device connection.
+        Final step in device lifecycle.
+        
+        Args:
+            device: Device to close
+            
+        Returns:
+            bool: True if closed successfully
+        """
+    
+    @hookspec
+    def execute(self, device: Device, target: str):
+        """Execute an action on the target using the device."""
+        warnings.warn(
+            "The 'execute' method is deprecated and will be removed in a future release. "
+            "Please use the 'command' method instead.",
+            DeprecationWarning,
+            stacklevel=2
+        )

iotsploit_core/core/device_store.py

@@ -0,0 +1,86 @@
+from typing import Dict, Optional
+import logging
+from iotsploit_core.domain.device import Device, DeviceType, SerialDevice, SocketCANDevice, USBDevice
+from iotsploit_core.core.device_config import DeviceConfigManager
+
+logger = logging.getLogger(__name__)
+
+class DeviceStore:
+    def __init__(self):
+        self.devices = {}
+        self.device_sources = {}
+        self.config_manager = DeviceConfigManager()
+        
+    def register_device(self, device: Device, source: str = "dynamic"):
+        """注册设备并保存其配置"""
+        self.devices[device.device_id] = device
+        self.device_sources[device.device_id] = source
+        
+    def get_device(self, device_id: str) -> Optional[Device]:
+        device = self.devices.get(device_id)
+        if not device:
+            config = self.config_manager.get_device_config(device_id)
+            if config:
+                return self._dict_to_device(config)
+        return device
+        
+    def _device_to_dict(self, device: Device) -> Dict:
+        """将设备对象转换为字典"""
+        base_dict = {
+            "device_id": device.device_id,
+            "name": device.name,
+            "device_type": device.device_type.value,
+            "attributes": device.attributes
+        }
+        
+        if isinstance(device, SerialDevice):
+            base_dict.update({
+                "port": device.port,
+                "baud_rate": device.baud_rate
+            })
+        elif isinstance(device, USBDevice):
+            base_dict.update({
+                "vendor_id": device.vendor_id,
+                "product_id": device.product_id
+            })
+        elif isinstance(device, SocketCANDevice):
+            base_dict.update({
+                "interface": device.interface
+            })
+            
+        return base_dict
+        
+    def _dict_to_device(self, data: Dict) -> Device:
+        """从字典创建设备对象"""
+        device_type = DeviceType(data["device_type"])
+        
+        if device_type == DeviceType.Serial:
+            return SerialDevice(
+                device_id=data["device_id"],
+                name=data["name"],
+                port=data.get("port", ""),
+                baud_rate=data.get("baud_rate", 115200),
+                attributes=data.get("attributes", {})
+            )
+        elif device_type == DeviceType.USB:
+            return USBDevice(
+                device_id=data["device_id"],
+                name=data["name"],
+                vendor_id=data.get("vendor_id", ""),
+                product_id=data.get("product_id", ""),
+                attributes=data.get("attributes", {})
+            )
+        elif device_type == DeviceType.CAN:
+            return SocketCANDevice(
+                device_id=data["device_id"],
+                name=data["name"],
+                interface=data.get("interface", "can0"),
+                attributes=data.get("attributes", {})
+            )
+        else:
+            return Device(
+                device_id=data["device_id"],
+                name=data["name"],
+                device_type=device_type,
+                attributes=data.get("attributes", {})
+            )