src/collectors/python.d.plugin/python_modules/bases/FrameworkServices/SocketService.py
# -*- coding: utf-8 -*-
# Description:
# Author: Pawel Krupa (paulfantom)
# Author: Ilya Mashchenko (ilyam8)
# SPDX-License-Identifier: GPL-3.0-or-later
import errno
import socket
try:
import ssl
except ImportError:
_TLS_SUPPORT = False
else:
_TLS_SUPPORT = True
if _TLS_SUPPORT:
try:
PROTOCOL_TLS = ssl.PROTOCOL_TLS
except AttributeError:
PROTOCOL_TLS = ssl.PROTOCOL_SSLv23
from bases.FrameworkServices.SimpleService import SimpleService
DEFAULT_CONNECT_TIMEOUT = 2.0
DEFAULT_READ_TIMEOUT = 2.0
DEFAULT_WRITE_TIMEOUT = 2.0
class SocketService(SimpleService):
def __init__(self, configuration=None, name=None):
self._sock = None
self._keep_alive = False
self.host = 'localhost'
self.port = None
self.unix_socket = None
self.dgram_socket = False
self.request = ''
self.tls = False
self.cert = None
self.key = None
self.__socket_config = None
self.__empty_request = "".encode()
SimpleService.__init__(self, configuration=configuration, name=name)
self.connect_timeout = configuration.get('connect_timeout', DEFAULT_CONNECT_TIMEOUT)
self.read_timeout = configuration.get('read_timeout', DEFAULT_READ_TIMEOUT)
self.write_timeout = configuration.get('write_timeout', DEFAULT_WRITE_TIMEOUT)
def _socket_error(self, message=None):
if self.unix_socket is not None:
self.error('unix socket "{socket}": {message}'.format(socket=self.unix_socket,
message=message))
else:
if self.__socket_config is not None:
_, _, _, _, sa = self.__socket_config
self.error('socket to "{address}" port {port}: {message}'.format(address=sa[0],
port=sa[1],
message=message))
else:
self.error('unknown socket: {0}'.format(message))
def _connect2socket(self, res=None):
"""
Connect to a socket, passing the result of getaddrinfo()
:return: boolean
"""
if res is None:
res = self.__socket_config
if res is None:
self.error("Cannot create socket to 'None':")
return False
af, sock_type, proto, _, sa = res
try:
self.debug('Creating socket to "{address}", port {port}'.format(address=sa[0], port=sa[1]))
self._sock = socket.socket(af, sock_type, proto)
except socket.error as error:
self.error('Failed to create socket "{address}", port {port}, error: {error}'.format(address=sa[0],
port=sa[1],
error=error))
self._sock = None
self.__socket_config = None
return False
if self.tls:
try:
self.debug('Encapsulating socket with TLS')
self.debug('Using keyfile: {0}, certfile: {1}, cert_reqs: {2}, ssl_version: {3}'.format(
self.key, self.cert, ssl.CERT_NONE, PROTOCOL_TLS
))
self._sock = ssl.wrap_socket(self._sock,
keyfile=self.key,
certfile=self.cert,
server_side=False,
cert_reqs=ssl.CERT_NONE,
ssl_version=PROTOCOL_TLS,
)
except (socket.error, ssl.SSLError, IOError, OSError) as error:
self.error('failed to wrap socket : {0}'.format(repr(error)))
self._disconnect()
self.__socket_config = None
return False
try:
self.debug('connecting socket to "{address}", port {port}'.format(address=sa[0], port=sa[1]))
self._sock.settimeout(self.connect_timeout)
self.debug('set socket connect timeout to: {0}'.format(self._sock.gettimeout()))
self._sock.connect(sa)
except (socket.error, ssl.SSLError) as error:
self.error('Failed to connect to "{address}", port {port}, error: {error}'.format(address=sa[0],
port=sa[1],
error=error))
self._disconnect()
self.__socket_config = None
return False
self.debug('connected to "{address}", port {port}'.format(address=sa[0], port=sa[1]))
self.__socket_config = res
return True
def _connect2unixsocket(self):
"""
Connect to a unix socket, given its filename
:return: boolean
"""
if self.unix_socket is None:
self.error("cannot connect to unix socket 'None'")
return False
try:
self.debug('attempting DGRAM unix socket "{0}"'.format(self.unix_socket))
self._sock = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)
self._sock.settimeout(self.connect_timeout)
self.debug('set socket connect timeout to: {0}'.format(self._sock.gettimeout()))
self._sock.connect(self.unix_socket)
self.debug('connected DGRAM unix socket "{0}"'.format(self.unix_socket))
return True
except socket.error as error:
self.debug('Failed to connect DGRAM unix socket "{socket}": {error}'.format(socket=self.unix_socket,
error=error))
try:
self.debug('attempting STREAM unix socket "{0}"'.format(self.unix_socket))
self._sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
self._sock.settimeout(self.connect_timeout)
self.debug('set socket connect timeout to: {0}'.format(self._sock.gettimeout()))
self._sock.connect(self.unix_socket)
self.debug('connected STREAM unix socket "{0}"'.format(self.unix_socket))
return True
except socket.error as error:
self.debug('Failed to connect STREAM unix socket "{socket}": {error}'.format(socket=self.unix_socket,
error=error))
self._sock = None
return False
def _connect(self):
"""
Recreate socket and connect to it since sockets cannot be reused after closing
Available configurations are IPv6, IPv4 or UNIX socket
:return:
"""
try:
if self.unix_socket is not None:
self._connect2unixsocket()
else:
if self.__socket_config is not None:
self._connect2socket()
else:
if self.dgram_socket:
sock_type = socket.SOCK_DGRAM
else:
sock_type = socket.SOCK_STREAM
for res in socket.getaddrinfo(self.host, self.port, socket.AF_UNSPEC, sock_type):
if self._connect2socket(res):
break
except Exception as error:
self.error('unhandled exception during connect : {0}'.format(repr(error)))
self._sock = None
self.__socket_config = None
def _disconnect(self):
"""
Close socket connection
:return:
"""
if self._sock is not None:
try:
self.debug('closing socket')
self._sock.shutdown(2) # 0 - read, 1 - write, 2 - all
self._sock.close()
except Exception as error:
if not (hasattr(error, 'errno') and error.errno == errno.ENOTCONN):
self.error(error)
self._sock = None
def _send(self, request=None):
"""
Send request.
:return: boolean
"""
# Send request if it is needed
if self.request != self.__empty_request:
try:
self.debug('set socket write timeout to: {0}'.format(self._sock.gettimeout()))
self._sock.settimeout(self.write_timeout)
self.debug('sending request: {0}'.format(request or self.request))
self._sock.send(request or self.request)
except Exception as error:
self._socket_error('error sending request: {0}'.format(error))
self._disconnect()
return False
return True
def _receive(self, raw=False):
"""
Receive data from socket
:param raw: set `True` to return bytes
:type raw: bool
:return: decoded str or raw bytes
:rtype: str/bytes
"""
data = "" if not raw else b""
while True:
self.debug('receiving response')
try:
self.debug('set socket read timeout to: {0}'.format(self._sock.gettimeout()))
self._sock.settimeout(self.read_timeout)
buf = self._sock.recv(4096)
except Exception as error:
self._socket_error('failed to receive response: {0}'.format(error))
self._disconnect()
break
if buf is None or len(buf) == 0: # handle server disconnect
if data == "" or data == b"":
self._socket_error('unexpectedly disconnected')
else:
self.debug('server closed the connection')
self._disconnect()
break
self.debug('received data')
data += buf.decode('utf-8', 'ignore') if not raw else buf
if self._check_raw_data(data):
break
self.debug(u'final response: {0}'.format(data if not raw else u'binary data'))
return data
def _get_raw_data(self, raw=False, request=None):
"""
Get raw data with low-level "socket" module.
:param raw: set `True` to return bytes
:type raw: bool
:return: decoded data (str) or raw data (bytes)
:rtype: str/bytes
"""
if self._sock is None:
self._connect()
if self._sock is None:
return None
# Send request if it is needed
if not self._send(request):
return None
data = self._receive(raw)
if not self._keep_alive:
self._disconnect()
return data
@staticmethod
def _check_raw_data(data):
"""
Check if all data has been gathered from socket
:param data: str
:return: boolean
"""
return bool(data)
def _parse_config(self):
"""
Parse configuration data
:return: boolean
"""
try:
self.unix_socket = str(self.configuration['socket'])
except (KeyError, TypeError):
self.debug('No unix socket specified. Trying TCP/IP socket.')
self.unix_socket = None
try:
self.host = str(self.configuration['host'])
except (KeyError, TypeError):
self.debug('No host specified. Using: "{0}"'.format(self.host))
try:
self.port = int(self.configuration['port'])
except (KeyError, TypeError):
self.debug('No port specified. Using: "{0}"'.format(self.port))
self.tls = bool(self.configuration.get('tls', self.tls))
if self.tls and not _TLS_SUPPORT:
self.warning('TLS requested but no TLS module found, disabling TLS support.')
self.tls = False
if _TLS_SUPPORT and not self.tls:
self.debug('No TLS preference specified, not using TLS.')
if self.tls and _TLS_SUPPORT:
self.key = self.configuration.get('tls_key_file')
self.cert = self.configuration.get('tls_cert_file')
if not self.cert:
# If there's not a valid certificate, clear the key too.
self.debug('No valid TLS client certificate configuration found.')
self.key = None
self.cert = None
elif not self.key:
# If a key isn't listed, the config may still be
# valid, because there may be a key attached to the
# certificate.
self.info('No TLS client key specified, assuming it\'s attached to the certificate.')
self.key = None
try:
self.request = str(self.configuration['request'])
except (KeyError, TypeError):
self.debug('No request specified. Using: "{0}"'.format(self.request))
self.request = self.request.encode()
def check(self):
self._parse_config()
return SimpleService.check(self)