Converted NetworkConnection schema/response.

platypush/message/response/system/__init__.py

@@ -28,39 +28,6 @@ class SensorResponse(SystemResponse):
     pass
 
 
-class NetworkConnectionResponse(NetworkResponse):
-    # noinspection PyShadowingBuiltins
-    def __init__(
-        self,
-        fd: int,
-        family: str,
-        type: str,
-        local_address: str,
-        local_port: int,
-        remote_address: str,
-        remote_port: int,
-        status: str,
-        pid: int,
-        *args,
-        **kwargs
-    ):
-        super().__init__(
-            *args,
-            output={
-                'fd': fd,
-                'family': family,
-                'type': type,
-                'local_address': local_address,
-                'local_port': local_port,
-                'remote_address': remote_address,
-                'remote_port': remote_port,
-                'status': status,
-                'pid': pid,
-            },
-            **kwargs
-        )
-
-
 class NetworkAddressResponse(NetworkResponse):
     def __init__(
         self,

platypush/plugins/system/__init__.py

@@ -23,7 +23,6 @@ from platypush.entities.system import (
 )
 from platypush.message.response.system import (
     NetworkResponseList,
-    NetworkConnectionResponse,
     NetworkAddressResponse,
     NetworkInterfaceStatsResponse,
     SensorTemperatureResponse,
@@ -38,6 +37,7 @@ from platypush.message.response.system import (
 from platypush.plugins import action
 from platypush.plugins.sensor import SensorPlugin
 from platypush.schemas.system import (
+    ConnectionSchema,
     CpuFrequency,
     CpuFrequencySchema,
     CpuInfo,
@@ -299,50 +299,36 @@ class SystemPlugin(SensorPlugin, EntityManager):
         return NetworkInterfaceSchema().dump(self._net_io_counters_avg())
 
     @action
-    def net_connections(
+    def net_connections(self, type: str = 'inet') -> List[dict]:
-        self, type: Optional[str] = None
-    ) -> Union[NetworkConnectionResponse, NetworkResponseList]:
         """
         Get the list of active network connections.
-        On macOS this function requires root privileges.
+        On MacOS this function requires root privileges.
 
-        :param type: Connection type to filter. Supported types:
+        :param type: Connection type to filter (default: ``inet``). Supported
+            types:
 
             +------------+----------------------------------------------------+
-            | Kind Value | Connections using                                  |
+            | ``type``   | Description                                        |
             +------------+----------------------------------------------------+
-            | inet       | IPv4 and IPv6                                      |
+            | ``inet``   | IPv4 and IPv6                                      |
-            | inet4      | IPv4                                               |
+            | ``inet4``  | IPv4                                               |
-            | inet6      | IPv6                                               |
+            | ``inet6``  | IPv6                                               |
-            | tcp        | TCP                                                |
+            | ``tcp``    | TCP                                                |
-            | tcp4       | TCP over IPv4                                      |
+            | ``tcp4``   | TCP over IPv4                                      |
-            | tcp6       | TCP over IPv6                                      |
+            | ``tcp6``   | TCP over IPv6                                      |
-            | udp        | UDP                                                |
+            | ``udp``    | UDP                                                |
-            | udp4       | UDP over IPv4                                      |
+            | ``udp4``   | UDP over IPv4                                      |
-            | udp6       | UDP over IPv6                                      |
+            | ``udp6``   | UDP over IPv6                                      |
-            | unix       | UNIX socket (both UDP and TCP protocols)           |
+            | ``unix``   | UNIX socket (both UDP and TCP protocols)           |
-            | all        | the sum of all the possible families and protocols |
+            | ``all``    | Any families and protocols                         |
             +------------+----------------------------------------------------+
 
-        :return: List of :class:`platypush.message.response.system.NetworkConnectionResponse`.
+        :return: .. schema:: system.ConnectionSchema(many=True)
         """
-        conns = psutil.net_connections(kind=type)
+        schema = ConnectionSchema()
-
+        return schema.dump(
-        return NetworkResponseList(
+            schema.load(psutil.net_connections(kind=type), many=True),  # type: ignore
-            [
+            many=True,
-                NetworkConnectionResponse(
-                    fd=conn.fd,
-                    family=conn.family.name,
-                    type=conn.type.name,
-                    local_address=conn.laddr[0] if conn.laddr else None,
-                    local_port=conn.laddr[1] if len(conn.laddr) > 1 else None,
-                    remote_address=conn.raddr[0] if conn.raddr else None,
-                    remote_port=conn.raddr[1] if len(conn.raddr) > 1 else None,
-                    status=conn.status,
-                    pid=conn.pid,
-                )
-                for conn in conns
-            ]
         )
 
     @action

platypush/schemas/system/__init__.py

@@ -1,3 +1,4 @@
+from ._connection import Connection, ConnectionSchema
 from ._cpu import (
     Cpu,
     CpuFrequency,
@@ -17,6 +18,8 @@ from ._schemas import SystemInfoSchema
 
 
 __all__ = [
+    "Connection",
+    "ConnectionSchema",
     "Cpu",
     "CpuFrequency",
     "CpuFrequencySchema",

platypush/schemas/system/_connection/__init__.py

@@ -0,0 +1,4 @@
+from ._model import Connection
+from ._schemas import ConnectionSchema
+
+__all__ = ['Connection', 'ConnectionSchema']

platypush/schemas/system/_connection/_base.py

@@ -0,0 +1,41 @@
+from marshmallow import pre_load
+
+from platypush.schemas.dataclasses import DataClassSchema
+
+
+class ConnectionBaseSchema(DataClassSchema):
+    """
+    Base schema for connections.
+    """
+
+    @pre_load
+    def pre_load(self, data, **_) -> dict:
+        if hasattr(data, '_asdict'):
+            data = data._asdict()
+
+        addr_mapping = {
+            'laddr': ('local_address', 'local_port'),
+            'raddr': ('remote_address', 'remote_port'),
+        }
+
+        # Parse laddr/raddr attributes
+        for ext_attr, (addr_attr, port_attr) in addr_mapping.items():
+            value = data.pop(ext_attr, None)
+            if not value:
+                data[addr_attr] = data[port_attr] = None
+            elif isinstance(value, tuple):
+                data[addr_attr], data[port_attr] = value
+            elif isinstance(value, str):
+                data[addr_attr] = value
+                data[port_attr] = None
+
+        # Handle enum values
+        for attr in ['type', 'family']:
+            value = data.pop(attr, None)
+            if value is not None:
+                data[attr] = value.name
+
+        if data.get('status') == 'NONE':
+            data['status'] = None
+
+        return data

platypush/schemas/system/_connection/_model.py

@@ -0,0 +1,96 @@
+from dataclasses import dataclass, field
+from socket import AddressFamily, SocketKind
+from typing import Optional
+
+
+@dataclass
+class Connection:
+    """
+    Network/UNIX socket data class.
+    """
+
+    fd: int = field(
+        metadata={
+            'metadata': {
+                'description': 'File descriptor',
+                'example': 3,
+            },
+        }
+    )
+
+    family: AddressFamily = field(
+        metadata={
+            'metadata': {
+                'description': 'Socket family',
+                'example': AddressFamily.AF_INET.name,
+            }
+        }
+    )
+
+    type: SocketKind = field(
+        metadata={
+            'metadata': {
+                'description': 'Socket type',
+                'example': SocketKind.SOCK_STREAM.name,
+            }
+        }
+    )
+
+    local_address: str = field(
+        metadata={
+            'metadata': {
+                'description': 'Local address, as an IP address for network '
+                'connections and a socket path for a UNIX socket',
+                'example': '192.168.1.2',
+            }
+        }
+    )
+
+    local_port: Optional[int] = field(
+        metadata={
+            'metadata': {
+                'description': 'Local port, if this is a TCP/UDP connection, '
+                'otherwise null',
+                'example': 12345,
+            }
+        }
+    )
+
+    remote_address: Optional[str] = field(
+        metadata={
+            'metadata': {
+                'description': 'Remote address, if this is a network '
+                'connection, otherwise null',
+                'example': '192.168.1.1',
+            }
+        }
+    )
+
+    remote_port: Optional[int] = field(
+        metadata={
+            'metadata': {
+                'description': 'Local port, if this is a TCP/UDP connection, '
+                'otherwise null',
+                'example': 443,
+            }
+        }
+    )
+
+    status: Optional[str] = field(
+        metadata={
+            'metadata': {
+                'description': 'Connection status, if this is a network '
+                'connection, otherise null',
+                'example': 'ESTABLISHED',
+            }
+        }
+    )
+
+    pid: Optional[int] = field(
+        metadata={
+            'metadata': {
+                'description': 'ID of the process that owns the connection',
+                'example': 4321,
+            }
+        }
+    )

platypush/schemas/system/_connection/_schemas.py

@@ -0,0 +1,7 @@
+from marshmallow_dataclass import class_schema
+
+from ._base import ConnectionBaseSchema
+from ._model import Connection
+
+
+ConnectionSchema = class_schema(Connection, base_schema=ConnectionBaseSchema)