peerlink/peerlink/localnodes.py

# -*- coding: utf-8 -*-
# vi:si:et:sw=4:sts=4:ts=4


from threading import Thread
import json
import socket
import struct
import _thread
import time

from .link import can_connect
from .settings import server, USER_ID, sk, ENCODING
from .utils import valid, get_public_ipv6, get_local_ipv4, get_interface

import logging
logger = logging.getLogger('localnodes')


class LocalNodesBase(Thread):

    _PORT = 9851 
    _TTL = 1

    def __init__(self, nodes):
        self._socket = None
        self._active = True
        self._nodes = nodes
        Thread.__init__(self)
        if not server['localnode_discovery']:
            return
        self.daemon = True
        self.start()

    def get(self, user_id):
        if user_id in self._nodes:
            if can_connect(self._nodes[user_id]):
                return self._nodes[user_id]

    def get_ip(self):
        pass

    def get_packet(self):
        self.host = self.get_ip()
        if self.host:
            message = json.dumps({
                'host': self.host,
                'port': server['node_port'],
                'cert': server['cert']
            })
            sig = sk.sign(message.encode(), encoding=ENCODING).decode()
            packet = json.dumps([sig, USER_ID, message]).encode()
        else:
            packet = None
        return packet

    def get_socket(self):
        pass

    def join(self):
        self._active = False
        '''
        if self._socket:
            try:
                self._socket.shutdown(socket.SHUT_RDWR)
            except:
                pass
            self._socket.close()
        '''
        return Thread.join(self)

    def new_node(self, data):
        logger.debug('new node %s', data)
        if can_connect(data):
            self._nodes[data['id']] = data
            self.send()

    def receive(self):
        last = time.mktime(time.localtime())
        while self._active:
            try:
                s = self.get_socket()
                s.settimeout(2)
                s.bind(('', self._PORT))
                while self._active:
                    data, addr = s.recvfrom(1024)
                    if self._active:
                        while data[-1] == 0:
                            data = data[:-1] # Strip trailing \0's
                        data = self.verify(data)
                        if data:
                            self.update_node(data)
            except socket.timeout:
                now = time.mktime(time.localtime())
                if now - last > 60:
                    last = now
                    _thread.start_new_thread(self.send, ())
            except:
                logger.debug('receive failed. restart later', exc_info=1)
                time.sleep(10)

    def run(self):
        self.send()
        self.receive()

    def send(self):
        pass

    def update_node(self, data):
        #logger.debug('update node %s', data)
        if data['id'] != USER_ID:
            if data['id'] not in self._nodes:
                _thread.start_new_thread(self.new_node, (data, ))
            elif can_connect(data):
                self._nodes[data['id']] = data

    def verify(self, data):

        try:
            packet = json.loads(data.decode())
        except:
            return None
        if len(packet) == 3:
            sig, user_id, data = packet
            if valid(user_id, data, sig):
                message = json.loads(data)
                message['id'] = user_id
                for key in ['id', 'host', 'port', 'cert']:
                    if key not in message:
                        return None
                return message


class LocalNodes4(LocalNodesBase):

    _BROADCAST = "239.255.255.250"
    _TTL = 1

    def get_ip(self):
        return get_local_ipv4()

    def get_socket(self):
        s = socket.socket (socket.AF_INET, socket.SOCK_DGRAM)
        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
        mreq = struct.pack("=4sl", socket.inet_aton(self._BROADCAST), socket.INADDR_ANY)
        s.setsockopt(socket.IPPROTO_IP, socket.IP_ADD_MEMBERSHIP, mreq)
        self._socket = s
        return s

    def send(self):
        packet = self.get_packet()
        if packet:
            #logger.debug('send4')
            sockaddr = (self._BROADCAST, self._PORT)
            s = socket.socket (socket.AF_INET, socket.SOCK_DGRAM)
            s.setsockopt (socket.IPPROTO_IP, socket.IP_MULTICAST_TTL, self._TTL)
            try:
                s.sendto(packet + b'\0', sockaddr)
            except:
                logger.debug('LocalNodes4.send failed', exc_info=1)
            s.close()


class LocalNodes6(LocalNodesBase):

    _BROADCAST = "ff02::1"

    def get_ip(self):
        return get_public_ipv6()

    def get_socket(self):
        s = socket.socket(socket.AF_INET6, socket.SOCK_DGRAM)
        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
        group_bin = socket.inet_pton(socket.AF_INET6, self._BROADCAST) + b'\0'*4
        s.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_JOIN_GROUP, group_bin)
        self._socket = s
        return s

    def send(self):
        packet = self.get_packet()
        if packet:
            logger.debug('send6 %s', packet)
            ttl = struct.pack('@i', self._TTL)
            address = self._BROADCAST + get_interface()
            addrs = socket.getaddrinfo(address, self._PORT, socket.AF_INET6, socket.SOCK_DGRAM)
            addr = addrs[0]
            (family, socktype, proto, canonname, sockaddr) = addr
            s = socket.socket(family, socktype, proto)
            s.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_MULTICAST_HOPS, ttl)
            try:
                s.sendto(packet + b'\0', sockaddr)
            except:
                logger.debug('LocalNodes6.send failed', exc_info=1)
            s.close()

class LocalNodes(object):

    _active = True
    _nodes4 = None
    _nodes6 = None

    def __init__(self):
        self._nodes = {}
        if not server['localnode_discovery']:
            return
        self._nodes4 = LocalNodes4(self._nodes)
        self._nodes6 = LocalNodes6(self._nodes)

    def cleanup(self):
        if self._active:
            for id in list(self._nodes.keys()):
                if not can_connect(self._nodes[id]):
                    del self._nodes[id]
                if not self._active:
                    break

    def get(self, user_id):
        if user_id in self._nodes:
            if can_connect(self._nodes[user_id]):
                return self._nodes[user_id]

    def info(self):
        return list(self._nodes.keys())

    def join(self):
        self._active = False
        if self._nodes4:
            self._nodes4.join()
        if self._nodes6:
            self._nodes6.join()