From b3a089648570052e7f853bd6d4e30966e2cdadd3 Mon Sep 17 00:00:00 2001 From: Fabio Manganiello Date: Sat, 22 Apr 2023 11:17:54 +0200 Subject: [PATCH] Converted `NetworkConnection` schema/response. --- platypush/message/response/system/__init__.py | 33 ------- platypush/plugins/system/__init__.py | 58 +++++------ platypush/schemas/system/__init__.py | 3 + .../schemas/system/_connection/__init__.py | 4 + platypush/schemas/system/_connection/_base.py | 41 ++++++++ .../schemas/system/_connection/_model.py | 96 +++++++++++++++++++ .../schemas/system/_connection/_schemas.py | 7 ++ 7 files changed, 173 insertions(+), 69 deletions(-) create mode 100644 platypush/schemas/system/_connection/__init__.py create mode 100644 platypush/schemas/system/_connection/_base.py create mode 100644 platypush/schemas/system/_connection/_model.py create mode 100644 platypush/schemas/system/_connection/_schemas.py diff --git a/platypush/message/response/system/__init__.py b/platypush/message/response/system/__init__.py index 98374a2ee..4ddf35bd0 100644 --- a/platypush/message/response/system/__init__.py +++ b/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, diff --git a/platypush/plugins/system/__init__.py b/platypush/plugins/system/__init__.py index 4b71e3c4b..059624573 100644 --- a/platypush/plugins/system/__init__.py +++ b/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( - self, type: Optional[str] = None - ) -> Union[NetworkConnectionResponse, NetworkResponseList]: + def net_connections(self, type: str = 'inet') -> List[dict]: """ 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 | - | inet4 | IPv4 | - | inet6 | IPv6 | - | tcp | TCP | - | tcp4 | TCP over IPv4 | - | tcp6 | TCP over IPv6 | - | udp | UDP | - | udp4 | UDP over IPv4 | - | udp6 | UDP over IPv6 | - | unix | UNIX socket (both UDP and TCP protocols) | - | all | the sum of all the possible families and protocols | + | ``inet`` | IPv4 and IPv6 | + | ``inet4`` | IPv4 | + | ``inet6`` | IPv6 | + | ``tcp`` | TCP | + | ``tcp4`` | TCP over IPv4 | + | ``tcp6`` | TCP over IPv6 | + | ``udp`` | UDP | + | ``udp4`` | UDP over IPv4 | + | ``udp6`` | UDP over IPv6 | + | ``unix`` | UNIX socket (both UDP and TCP 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) - - return NetworkResponseList( - [ - 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 - ] + schema = ConnectionSchema() + return schema.dump( + schema.load(psutil.net_connections(kind=type), many=True), # type: ignore + many=True, ) @action diff --git a/platypush/schemas/system/__init__.py b/platypush/schemas/system/__init__.py index 2e45aae89..2b02c0200 100644 --- a/platypush/schemas/system/__init__.py +++ b/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", diff --git a/platypush/schemas/system/_connection/__init__.py b/platypush/schemas/system/_connection/__init__.py new file mode 100644 index 000000000..14e3127ae --- /dev/null +++ b/platypush/schemas/system/_connection/__init__.py @@ -0,0 +1,4 @@ +from ._model import Connection +from ._schemas import ConnectionSchema + +__all__ = ['Connection', 'ConnectionSchema'] diff --git a/platypush/schemas/system/_connection/_base.py b/platypush/schemas/system/_connection/_base.py new file mode 100644 index 000000000..efa798b7f --- /dev/null +++ b/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 diff --git a/platypush/schemas/system/_connection/_model.py b/platypush/schemas/system/_connection/_model.py new file mode 100644 index 000000000..d4c4a760e --- /dev/null +++ b/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, + } + } + ) diff --git a/platypush/schemas/system/_connection/_schemas.py b/platypush/schemas/system/_connection/_schemas.py new file mode 100644 index 000000000..2ad33a7ac --- /dev/null +++ b/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)