src/models/device.py - master - devsec/anothertap

src/models/device.py
108 строк4 KB
devsec
clean code gitignore
24 авг 2025, 21:27
24 авг 2025, 21:2775a9527
"""
Device model for representing measurement devices with SCPI capabilities
"""

from dataclasses import dataclass, field
from typing import Dict, Any, Optional
from dataclasses_json import dataclass_json
from enum import Enum
import uuid


class ConnectionType(Enum):
    """Supported connection types for devices"""
    TCPIP = "TCPIP"
    USB = "USB"
    SERIAL = "SERIAL"
    GPIB = "GPIB"
    SIMULATION = "SIMULATION"  # For simulated devices


class DeviceStatus(Enum):
    """Device connection status"""
    DISCONNECTED = "DISCONNECTED"
    CONNECTED = "CONNECTED"
    ERROR = "ERROR"


@dataclass_json
@dataclass
class Device:
    """
    Represents a measurement device with SCPI remote command capabilities
    """
    id: str = field(default_factory=lambda: str(uuid.uuid4()))
    name: str = ""
    manufacturer: str = ""
    model: str = ""
    serial_number: str = ""
    connection_type: ConnectionType = ConnectionType.TCPIP
    connection_string: str = ""  # e.g., "TCPIP::192.168.1.100::INSTR"
    description: str = ""
    capabilities: Dict[str, Any] = field(default_factory=dict)
    custom_properties: Dict[str, Any] = field(default_factory=dict)
    status: DeviceStatus = DeviceStatus.DISCONNECTED
    last_error: Optional[str] = None
    simulation_mode: bool = False  # Whether device is simulated
    simulation_config: Dict[str, Any] = field(default_factory=dict)  # Simulation configuration
    
    def __post_init__(self):
        """Post-initialization processing"""
        if isinstance(self.connection_type, str):
            self.connection_type = ConnectionType(self.connection_type)
        if isinstance(self.status, str):
            self.status = DeviceStatus(self.status)
    
    @property
    def display_name(self) -> str:
        """Get a user-friendly display name for the device"""
        if self.name:
            return self.name
        elif self.manufacturer and self.model:
            return f"{self.manufacturer} {self.model}"
        else:
            return f"Device {self.id[:8]}"
    
    def is_connected(self) -> bool:
        """Check if device is currently connected"""
        return self.status == DeviceStatus.CONNECTED
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert device to dictionary for JSON serialization"""
        return {
            'id': self.id,
            'name': self.name,
            'manufacturer': self.manufacturer,
            'model': self.model,
            'serial_number': self.serial_number,
            'connection_type': self.connection_type.value,
            'connection_string': self.connection_string,
            'description': self.description,
            'capabilities': self.capabilities,
            'custom_properties': self.custom_properties,
            'status': self.status.value,
            'last_error': self.last_error,
            'simulation_mode': self.simulation_mode,
            'simulation_config': self.simulation_config
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'Device':
        """Create device from dictionary"""
        device = cls(
            id=data.get('id', str(uuid.uuid4())),
            name=data.get('name', ''),
            manufacturer=data.get('manufacturer', ''),
            model=data.get('model', ''),
            serial_number=data.get('serial_number', ''),
            connection_type=ConnectionType(data.get('connection_type', ConnectionType.TCPIP.value)),
            connection_string=data.get('connection_string', ''),
            description=data.get('description', ''),
            capabilities=data.get('capabilities', {}),
            custom_properties=data.get('custom_properties', {}),
            status=DeviceStatus(data.get('status', DeviceStatus.DISCONNECTED.value)),
            last_error=data.get('last_error'),
            simulation_mode=data.get('simulation_mode', False),
            simulation_config=data.get('simulation_config', {})
        )
        return device