# -*- coding: utf-8 -*- # vi:si:et:sw=4:sts=4:ts=4 from __future__ import division import json import socket import struct import thread from threading import Thread import time from utils import valid, get_public_ipv6, get_local_ipv4, get_interface from settings import preferences, server, USER_ID, sk import state import db import user.models import logging logger = logging.getLogger('oml.localnodes') def can_connect(data): try: if ':' in data['host']: s = socket.socket(socket.AF_INET6, socket.SOCK_DGRAM) else: s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) s.settimeout(1) s.connect((data['host'], data['port'])) s.close() return True except: pass logger.debug('can_connect failed') return False 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_packet(self): self.host = self.get_ip() if self.host: message = json.dumps({ 'username': preferences.get('username', 'anonymous'), 'host': self.host, 'port': server['node_port'], 'cert': server['cert'] }) sig = sk.sign(message, encoding='base64') packet = json.dumps([sig, USER_ID, message]) else: packet = None return packet def get_socket(self): pass def send(self): pass 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: if self._active: now = time.mktime(time.localtime()) if now - last > 60: last = now thread.start_new_thread(self.send, ()) except: if self._active: logger.debug('receive failed. restart later', exc_info=1) time.sleep(10) def verify(self, data): try: packet = json.loads(data) 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', 'username', 'host', 'port', 'cert']: if key not in message: return None return message def update_node(self, data): #fixme use local link address #print addr 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 get(self, user_id): if user_id in self._nodes: if can_connect(self._nodes[user_id]): return self._nodes[user_id] def new_node(self, data): logger.debug('NEW NODE %s', data) if can_connect(data): self._nodes[data['id']] = data with db.session(): u = user.models.User.get_or_create(data['id']) u.info['username'] = data['username'] u.info['local'] = data u.update_name() u.save() state.nodes.queue('add', u.id) self.send() def get_ip(self): pass def run(self): self.send() self.receive() 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) class LocalNodes4(LocalNodesBase): _BROADCAST = "239.255.255.250" _TTL = 1 def send(self): packet = self.get_packet() if packet: logger.debug('send4 %s', packet) 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 + '\0', sockaddr) except: logger.debug('LocalNodes4.send failed', exc_info=1) s.close() 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 get_ip(self): return get_local_ipv4() class LocalNodes6(LocalNodesBase): _BROADCAST = "ff02::1" 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 + '\0', sockaddr) except: logger.debug('LocalNodes6.send failed', exc_info=1) s.close() 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) + '\0'*4 s.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_JOIN_GROUP, group_bin) self._socket = s return s def get_ip(self): return get_public_ipv6() 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 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 join(self): self._active = False if self._nodes4: self._nodes4.join() if self._nodes6: self._nodes6.join()