import socket try: import ssl except ImportError: _TLS_SUPPORT = False else: _TLS_SUPPORT = True from .SimpleService import SimpleService class SocketService(SimpleService): def __init__(self): self._sock = None self._keep_alive = False # self.host = 'localhost' # Using lo instead of localhost incase /etc/hosts is wrong self.host = '127.0.0.1' 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() self.get_initial = True SimpleService.__init__(self) 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)) pass else: if self.__socket_config is not None: af, sock_type, proto, canon_name, 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)) pass 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, canon_name, 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._sock = ssl.wrap_socket(self._sock, keyfile=self.key, certfile=self.cert, server_side=False, cert_reqs=ssl.CERT_NONE) except (socket.error, ssl.SSLError) as error: # self.error('Failed to wrap socket.') self._disconnect() self.__socket_config = None return False try: # self.debug('connecting socket to "{address}", port {port}'.format(address=sa[0], port=sa[1])) # print('connecting socket to "{address}", port {port}'.format(address=sa[0], port=sa[1])) 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])) # print('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.connect(self.unix_socket) # self.debug('connected DGRAM unix socket "{0}"'.format(self.unix_socket)) return True except socket.error as error: pass # 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.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: self._sock = None self.__socket_config = None if self._sock is not None: self._sock.setblocking(0) self._sock.settimeout(5) # self.debug('set socket timeout to: {0}'.format(self._sock.gettimeout())) # print('set socket timeout to: {0}'.format(self._sock.gettimeout())) 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: pass self._sock = None def _send(self): """ Send request. :return: boolean """ # Send request if it is needed if self.request != self.__empty_request: try: # self.debug('sending request: {0}'.format(self.request)) self._sock.send(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') # print('receiving response') try: 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') pass 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('final response: {0}'.format(data)) # print('final response: {0}'.format(data)) return data def _get_raw_data(self, raw=False): """ 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 # Get initial connection message data = self._receive(raw) if self.get_initial else '' # Send request and get post request data if self._send(): 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 check(self): return SimpleService.check(self)