Ejemplo n.º 1
0
class UDPmcSocket(socket.socket):
    '''
    The UdpSocket class enables the send and receive UDP messages to a
    multicast group and unicast address. The unicast socket is only created if
    'send_mcast' and 'listen_mcast' parameter are set to False or a specific interface is defined.
    '''
    def __init__(self,
                 port,
                 mgroup,
                 router=None,
                 ttl=16,
                 interface='',
                 logger_name='udp_mc',
                 send_buffer=0,
                 recv_buffer=0,
                 queue_length=0,
                 loglevel='info'):
        '''
        Creates a socket, bind it to a given port and join to a given multicast
        group. IPv4 and IPv6 are supported.

        :param int port: the port to bind the socket
        :param str mgroup: the multicast group to join
        :param router: class which provides `route_udp_msg(fkie_iop_node_manager.message.Message)` method. If `None` receive will be disabled.
        :param type: fkie_iop_node_manager.queue
        :param int ttl: time to leave (Default: 20)
        :param str interface: IP of interface to bind (Default: '').
        '''
        self.logger = NMLogger(
            '%s[%s:%d]' % (logger_name, mgroup.replace('.', '_'), port),
            loglevel)
        self.port = port
        self.mgroup = mgroup
        self._lock = threading.RLock()
        self._closed = False
        self._recv_buffer = recv_buffer
        self._locals = [ip for _ifname, ip in localifs()]
        self._locals.append('localhost')
        self._sender_endpoints = {}
        self.sock_5_error_printed = []
        self.SOKET_ERRORS_NEEDS_RECONNECT = False
        self.interface = interface
        # get the AF_INET information for group to ensure that the address family
        # of group is the same as for interface
        addrinfo = getaddrinfo(self.mgroup)
        self.interface_ip = ''
        if self.interface:
            addrinfo = getaddrinfo(self.interface, addrinfo[0])
            if addrinfo is not None:
                self.interface_ip = addrinfo[4][0]
        self.logger.debug("destination: %s" % self.mgroup)
        self.logger.debug("interface : %s (detected ip: %s)" %
                          (self.interface, self.interface_ip))
        self.logger.debug("inet: %s" % str(addrinfo))

        socket.socket.__init__(self, addrinfo[0], socket.SOCK_DGRAM,
                               socket.IPPROTO_UDP)
        self.logger.info("Create multicast socket @('%s', %d)" %
                         (self.mgroup, port))
        # initialize multicast socket
        # Allow multiple copies of this program on one machine
        if hasattr(socket, "SO_REUSEPORT"):
            try:
                self.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
            except Exception:
                self.logger.warning(
                    "SO_REUSEPORT failed: Protocol not available, some functions are not available."
                )
        # Set Time-to-live (optional) and loop count
        ttl_bin = struct.pack('@i', ttl)
        if addrinfo[0] == socket.AF_INET:  # IPv4
            self.setsockopt(socket.IPPROTO_IP, socket.IP_MULTICAST_TTL,
                            ttl_bin)
            self.setsockopt(socket.IPPROTO_IP, socket.IP_MULTICAST_LOOP, 1)
        else:  # IPv6
            self.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_MULTICAST_HOPS,
                            ttl_bin)
            self.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_MULTICAST_LOOP, 1)

        try:
            if addrinfo[0] == socket.AF_INET:  # IPv4
                # Create group_bin for de-register later
                # Set socket options for multicast specific interface or general
                if not self.interface_ip:
                    self.group_bin = socket.inet_pton(
                        socket.AF_INET, self.mgroup) + struct.pack(
                            '=I', socket.INADDR_ANY)
                    self.setsockopt(socket.IPPROTO_IP,
                                    socket.IP_ADD_MEMBERSHIP, self.group_bin)
                else:
                    self.group_bin = socket.inet_aton(
                        self.mgroup) + socket.inet_aton(self.interface_ip)
                    self.setsockopt(socket.IPPROTO_IP, socket.IP_MULTICAST_IF,
                                    socket.inet_aton(self.interface_ip))
                    self.setsockopt(socket.IPPROTO_IP,
                                    socket.IP_ADD_MEMBERSHIP, self.group_bin)
            else:  # IPv6
                # Create group_bin for de-register later
                # Set socket options for multicast
                self.group_bin = socket.inet_pton(addrinfo[0],
                                                  self.mgroup) + struct.pack(
                                                      '@I', socket.INADDR_ANY)
                self.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_JOIN_GROUP,
                                self.group_bin)
        except socket.error as errobj:
            msg = str(errobj)
            if errobj.errno in [errno.ENODEV]:
                msg = "socket.error[%d]: %s,\nis multicast route set? e.g. sudo route add -net 224.0.0.0 netmask 224.0.0.0 eth0" % (
                    errobj.errno, msg)
            raise Exception(msg)

        # set buffer size if configured
        if send_buffer:
            old_bufsize = self.getsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF)
            if old_bufsize != send_buffer:
                self.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF,
                                send_buffer)
                bufsize = self.getsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF)
                self.logger.debug("Changed buffer size from %d to %d" %
                                  (old_bufsize, bufsize))

        # Bind to the port
        try:
            # bind to default interfaces if not unicast socket was created
            self.bind((self.interface_ip, port))
        except socket.error as errobj:
            msg = str(errobj)
            self.logger.critical(
                "Unable to bind multicast to interface: %s, check that it exists: %s"
                % (self.mgroup, msg))
            raise
        self._router = router
        self._queue_send = queue.PQueue(queue_length,
                                        'queue_udp_send',
                                        loglevel=loglevel)
        self._parser_mcast = MessageParser(None, loglevel=loglevel)
        self.addrinfo = addrinfo
        # create a thread to handle the received multicast messages
        if self._router is not None:
            self._thread_recv = threading.Thread(target=self._loop_recv)
            self._thread_recv.start()
        self._thread_send = threading.Thread(target=self._loop_send)
        self._thread_send.start()

    def close(self):
        '''
        Unregister from the multicast group and close the socket.
        '''
        self._closed = True
        self.logger.info("Close multicast socket")
        try:
            # shutdown to cancel recvfrom()
            socket.socket.shutdown(self, socket.SHUT_RD)
        except socket.error:
            pass
        # Use the stored group_bin to de-register
        if self.addrinfo[0] == socket.AF_INET:  # IPv4
            self.setsockopt(socket.IPPROTO_IP, socket.IP_DROP_MEMBERSHIP,
                            self.group_bin)
        else:  # IPv6
            self.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_LEAVE_GROUP,
                            self.group_bin)
        socket.socket.close(self)
        self._queue_send.clear()

    def send_queued(self, msg):
        try:
            self._queue_send.put(msg)
        except queue.Full as full:
            print(traceback.format_exc())
            self.logger.warning("Can't send message: %s" % full)
        except Exception as e:
            self.logger.warning("Error while put message into queue: %s" % e)

    def _loop_send(self):
        while not self._closed:
            # Wait for next available Message. This method cancel waiting on clear() of PQueue and return None.
            msg = self._queue_send.get()
            if msg is not None:
                dst = msg.tinfo_dst
                if dst is None:  # or msg.dst_id.has_wildcards():
                    dst = AddressBook.Endpoint(AddressBook.Endpoint.UDP,
                                               self.mgroup,
                                               self.getsockname()[1])
                if dst is not None:
                    self._sendto(msg, dst)
                else:
                    self.logger.warning(
                        "Can't send message to %s, destination not found!" %
                        (dst))
                # TODO: add retry mechanism

    def _sendto(self, msg, endpoint):
        # send to given addresses
        try:
            # self.logger.debug("Send to %s:%d" % (endpoint.address, endpoint.port))
            val = self.sendto(msg.bytes(), (endpoint.address, endpoint.port))
            if val != msg.raw_size:
                raise Exception("not complete send %d of %d" %
                                (val, msg.raw_size))
            if endpoint.address in SEND_ERRORS:
                del SEND_ERRORS[endpoint.address]
        except socket.error as errobj:
            erro_msg = "Error while send to '%s': %s" % (endpoint.address,
                                                         errobj)
            SEND_ERRORS[endpoint.address] = erro_msg
            # -2: Name or service not known
            if errobj.errno in [-5, -2]:
                if endpoint.address not in self.sock_5_error_printed:
                    self.logger.warning(erro_msg)
                    self.sock_5_error_printed.append(endpoint.address)
            else:
                self.logger.warning(erro_msg)
            if errobj.errno in [
                    errno.ENETDOWN, errno.ENETUNREACH, errno.ENETRESET, errno
            ]:
                self.SOKET_ERRORS_NEEDS_RECONNECT = True
        except Exception as e:
            erro_msg = "Send to host '%s' failed: %s" % (endpoint.address, e)
            self.logger.warning(erro_msg)
            SEND_ERRORS[endpoint.address] = erro_msg

    def _loop_recv(self):
        '''
        This method handles the received multicast messages.
        '''
        while not self._closed:
            try:
                (data, address) = self.recvfrom(self._recv_buffer)
                if data and not self._closed and address[
                        0] not in self._locals:  # skip messages received from self
                    msgs = self._parser_mcast.unpack(data)
                    for msg in msgs:
                        try:
                            msg.tinfo_src = self._sender_endpoints[address]
                        except KeyError:
                            endpoint = AddressBook.Endpoint(
                                AddressBook.Endpoint.UDP, address[0],
                                address[1])
                            msg.tinfo_src = endpoint
                            self._sender_endpoints[address] = endpoint
                        # self.logger.debug("Received from %s" % (msg.tinfo_src))
                        self._router.route_udp_msg(msg)
            except socket.timeout:
                pass
            except queue.Full as full_error:
                self.logger.warning(
                    "Error while process received multicast message: %s" %
                    full_error)
            except socket.error:
                if not self._closed:
                    self.logger.warning("socket error: %s" %
                                        traceback.format_exc())
Ejemplo n.º 2
0
class UDPucSocket(socket.socket):
    def __init__(self,
                 port=0,
                 router=None,
                 interface='',
                 logger_name='udp',
                 default_dst=None,
                 send_buffer=0,
                 recv_buffer=0,
                 queue_length=0,
                 loglevel='info'):
        '''
        Creates a socket, bind it to a given interface+port for unicast send/receive.
        IPv4 and IPv6 are supported.

        :param int port: the port to bind the socket. If zero an empty one will be used.
        :param router: class which provides `route_udp_msg(fkie_iop_node_manager.message.Message)` method. If `None` receive will be disabled.
        :param str interface: The interface to bind to. If empty, it binds to all interfaces
        :param tuple(str,int) default_dst: used for loopback to send messages to predefined destination.
        '''
        self._closed = False
        self.logger = NMLogger('%s[%s:%d]' % (logger_name, interface, port),
                               loglevel)
        self.interface = interface
        self.port = port
        self._router = router
        self._default_dst = default_dst
        self._recv_buffer = recv_buffer
        self._sender_endpoints = {}
        self.sock_5_error_printed = []
        # If interface isn't specified, try to find an non localhost interface to
        # get some info for binding. Otherwise use localhost
        # if not self.interface:
        #     ifaces = localifs()
        #     for iface in ifaces:
        #         if not (iface[1].startswith('127') or iface[1].startswith('::1')):
        #             self.interface = iface[1]
        #             break
        self.logger.info("+ Bind to unicast socket @(%s:%s)" %
                         (self.interface, port))
        socket_type = socket.AF_INET
        bind_ip = self.interface
        if self.interface:
            addrinfo = getaddrinfo(self.interface)
            socket_type = addrinfo[0]
            bind_ip = addrinfo[4][0]
            # Configure socket type
        socket.socket.__init__(self, socket_type, socket.SOCK_DGRAM,
                               socket.IPPROTO_UDP)
        # Bind to the port
        try:
            self.logger.debug("Ucast bind to: (%s:%s)" % (bind_ip, port))
            self.bind((bind_ip, port))
        except socket.error as errobj:
            msg = str(errobj)
            self.logger.critical(
                "Unable to bind unicast to interface: %s, check that it exists: %s"
                % (bind_ip, msg))
            raise
        if self.port == 0:
            self.port = self.getsockname()[1]
        if send_buffer:
            # update buffer size
            old_bufsize = self.getsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF)
            if old_bufsize != send_buffer:
                self.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF,
                                send_buffer)
                #                self.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, buffersize)
                bufsize = self.getsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF)
                self.logger.debug("Changed buffer size from %d to %d" %
                                  (old_bufsize, bufsize))
        self._parser_ucast = MessageParser(None, loglevel=loglevel)
        self._queue_send = queue.PQueue(queue_length,
                                        'queue_%s_send' % logger_name,
                                        loglevel=loglevel)
        # create a thread to handle the received unicast messages
        if self._router is not None:
            self._thread_recv = threading.Thread(target=self._loop_recv)
            self._thread_recv.start()
        self._thread_send = threading.Thread(target=self._loop_send)
        self._thread_send.start()

    def close(self):
        """ Cleanup and close the socket"""
        self._closed = True
        self.logger.info("Close unicast socket")
        try:
            # shutdown to cancel recvfrom()
            socket.socket.shutdown(self, socket.SHUT_RD)
        except socket.error:
            pass
        socket.socket.close(self)
        self._queue_send.clear()

    def send_queued(self, msg):
        try:
            self._queue_send.put(msg)
        except queue.Full as full:
            print(traceback.format_exc())
            self.logger.warning("Can't send message: %s" % full)
        except Exception as e:
            self.logger.warning("Error while put message into queue: %s" % e)

    def _loop_send(self):
        while not self._closed:
            # Waits for next available Message. This method cancel waiting on clear() of PQueue and return None.
            msg = self._queue_send.get()
            if msg is not None:
                dst = msg.tinfo_dst
                if self._default_dst is not None:
                    # it is a loopback socket, send to fictive debug destination
                    dst = AddressBook.Endpoint(AddressBook.Endpoint.UDP,
                                               self._default_dst[0],
                                               self._default_dst[1])
                if dst is not None:
                    # send to given addresses
                    self._sendto(msg.bytes(), dst.address, dst.port)
            # TODO: add retry mechanism?

    def _sendto(self, msg, addr, port):
        '''
        Sends the given message to the joined multicast group. Some errors on send
        will be ignored (``ENETRESET``, ``ENETDOWN``, ``ENETUNREACH``)

        :param str msg: message to send
        :param str addr: IPv4 or IPv6 address
        :param int port: destination port
        '''
        try:
            self.logger.debug("Send to %s:%d" % (addr, port))
            self.sendto(msg, (addr, port))
        except socket.error as errobj:
            msg = str(errobj)
            if errobj.errno in [-5]:
                if addr not in self.sock_5_error_printed:
                    self.logger.warning("socket.error[%d]: %s, addr: %s" %
                                        (errobj.errno, msg, addr))
                    self.sock_5_error_printed.append(addr)
            elif errobj.errno in [errno.EINVAL, -2]:
                raise
            elif errobj.errno not in [
                    errno.ENETDOWN, errno.ENETUNREACH, errno.ENETRESET
            ]:
                raise

    def _loop_recv(self):
        '''
        This method handles the received unicast messages.
        '''
        while not self._closed:
            try:
                (data, address) = self.recvfrom(self._recv_buffer)
                if data and not self._closed:
                    msgs = self._parser_ucast.unpack(data)
                    for msg in msgs:
                        try:
                            msg.tinfo_src = self._sender_endpoints[address]
                        except KeyError:
                            endpoint = AddressBook.Endpoint(
                                AddressBook.Endpoint.UDP, address[0],
                                address[1])
                            msg.tinfo_src = endpoint
                            self._sender_endpoints[address] = endpoint
                        self.logger.debug("Received from %s" % (msg.tinfo_src))
                        self._router.route_udp_msg(msg)
            except queue.Full as full_error:
                self.logger.warning(
                    "Error while process received unicast message: %s" %
                    full_error)
            except socket.error:
                if not self._closed:
                    self.logger.warning("unicast socket error: %s" %
                                        traceback.format_exc())