class Tunnel(object): _TCP_INITIAL_DATA = 0 _TCP_FIN_DATA = 1 _TCP_CLOSED_DATA = 2 _UDP_INITIAL_DATA = 3 _UDP_CLOSED_DATA = 4 _TUN_INITIAL_DATA = 5 _PAYLOAD = 10 _HEARTBEAT = 100 _static_handlers = { _HEARTBEAT: (lambda _, __, ___: None) } @staticmethod def set_tcp_initial_handler(handler): Tunnel._static_handlers[Tunnel._TCP_INITIAL_DATA] = handler @staticmethod def set_tcp_fin_received_handler(handler): Tunnel._static_handlers[Tunnel._TCP_FIN_DATA] = handler @staticmethod def set_tcp_closed_handler(handler): Tunnel._static_handlers[Tunnel._TCP_CLOSED_DATA] = handler @staticmethod def set_udp_initial_handler(handler): Tunnel._static_handlers[Tunnel._UDP_INITIAL_DATA] = handler @staticmethod def set_udp_closed_handler(handler): Tunnel._static_handlers[Tunnel._UDP_CLOSED_DATA] = handler @staticmethod def set_tun_initial_handler(handler): Tunnel._static_handlers[Tunnel._TUN_INITIAL_DATA] = handler def __init__(self, connection=None, connect_to=None): self._stream = connection self._connect_to = connect_to self._on_initial_data = None self._on_payload = None self._on_stream_closed = None self._handlers = self._static_handlers.copy() self._handlers.update({ Tunnel._PAYLOAD: lambda _, id_, data: self._on_payload(self, id_, data) }) self._on_ready_to_send = None self._on_send_buffer_full = None self._hb_event = None self.connections = {} def __hash__(self): return hash(self._stream) def __eq__(self, other): if not isinstance(other, Tunnel): return False return self._stream == other._stream def __str__(self): return str(self._stream) def _send_heartbeat(self): self._send_content(Tunnel._HEARTBEAT, None, None) self._enable_heartbeat() def _enable_heartbeat(self): self._hb_event = Event.add_timer(HEARTBEAT_INTERVAL) self._hb_event.set_handler(lambda ev: self._send_heartbeat()) def _disable_heartbeat(self): if self._hb_event is not None: self._hb_event.del_timer() self._hb_event = None def _on_fin_received(self): self._disable_heartbeat() self._stream.close() def initialize(self): if self._stream is None: self._stream = Stream(prefix='TUNNEL') # self._stream.set_buffer_size(BUFF_SIZE) self._stream.set_tcp_no_delay() self._stream.append_send_handler(obscure.pack_data) self._stream.append_send_handler(obscure.random_padding) # self._stream.append_send_handler(obscure.gen_aes_encrypt()) self._stream.append_send_handler(obscure.gen_xor_encrypt()) # self._stream.append_send_handler(obscure.base64_encode) self._stream.append_send_handler(obscure.gen_http_encode(self._connect_to is not None)) self._stream.append_receive_handler(obscure.gen_http_decode(self._connect_to is not None)) # self._stream.append_receive_handler(obscure.base64_decode) self._stream.append_receive_handler(obscure.gen_xor_decrypt()) # self._stream.append_receive_handler(obscure.gen_aes_decrypt()) self._stream.append_receive_handler(obscure.unpad_random) self._stream.append_receive_handler(obscure.unpack_data) self._stream.set_on_ready_to_send(lambda _: self._on_tunnel_ready_to_send()) self._stream.set_on_send_buffer_full(lambda _: self._on_tunnel_send_buffer_full()) self._stream.set_on_received(lambda _, data, addr: self._on_received(data, addr)) self._stream.set_on_fin_received(lambda _: self._on_fin_received()) self._stream.set_on_closed(lambda _: self._on_closed()) if self._connect_to is not None: self._stream.connect(*self._connect_to) else: self._stream.set_on_decode_error(lambda _, received: self._on_decode_error(received)) self._stream.start_receiving() self._enable_heartbeat() def register(self, key, conn): _logger.debug('%s, register: %s(%s)', str(self), str(key), str(conn)) assert(key not in self.connections) self.connections[key] = conn def deregister(self, key): _logger.debug('%s, deregister(%s)', str(self), str(key)) if key in self.connections: del self.connections[key] def get_connection(self, key): if key in self.connections: return self.connections[key] else: _logger.debug('no such connection: %s', str(key)) return None def clear_connections(self): _logger.debug('%s, clear_connections (%d)', str(self), len(self.connections)) self.connections.clear() def is_ready_to_send(self): return self._stream.is_ready_to_send() def _send_content(self, type_, id_, content): if id_ is None: to_send = struct.pack('!HI', type_, 0) + '\x00' * 16 else: to_send = struct.pack('!HI', type_, len(content)) + id_.get_bytes() + content self._stream.send(to_send) def _on_tunnel_ready_to_send(self): if self._on_ready_to_send is not None: self._on_ready_to_send(self) def _on_tunnel_send_buffer_full(self): if self._on_send_buffer_full is not None: self._on_send_buffer_full(self) def _on_received(self, data, _addr): _logger.debug("tunnel %s received %d bytes" % (str(self), len(data))) if len(data) < 6 + 16: raise Exception('corrupted data') type_, content_length = struct.unpack('!HI', data[: 6]) id_ = uuid.UUID(bytes=data[6: 6 + 16]) if type_ not in self._handlers or self._handlers[type_] is None: _logger.warning("tunnel message type %d can not be handled", type_) if len(data) - 6 - 16 != content_length: raise Exception('corrupted data') self._handlers[type_](self, id_, data[6 + 16:]) return True def send_tcp_initial_data(self, id_, data): self._send_content(Tunnel._TCP_INITIAL_DATA, id_, data) def send_tcp_fin_data(self, id_, data=''): self._send_content(Tunnel._TCP_FIN_DATA, id_, data) def send_tcp_closed_data(self, id_, data=''): self._send_content(Tunnel._TCP_CLOSED_DATA, id_, data) def send_udp_initial_data(self, id_, data): self._send_content(Tunnel._UDP_INITIAL_DATA, id_, data) def send_udp_closed_data(self, id_, data=''): self._send_content(Tunnel._UDP_CLOSED_DATA, id_, data) def send_tun_initial_data(self, id_, data): self._send_content(Tunnel._TUN_INITIAL_DATA, id_, data) def send_payload(self, id_, payload): self._send_content(Tunnel._PAYLOAD, id_, payload) def set_on_payload(self, handler): self._on_payload = handler def set_on_ready_to_send(self, handler): self._on_ready_to_send = handler def set_on_send_buffer_full(self, handler): self._on_send_buffer_full = handler def _on_closed(self): self._disable_heartbeat() if self._on_stream_closed is not None: self._on_stream_closed(self) def set_on_closed(self, handler): self._on_stream_closed = handler def close(self): self._stream.close() def is_closed(self): return self._stream.is_closed() def _on_decode_error(self, received): self._disable_heartbeat() self._stream._encoders = [] backend = Stream(prefix="SIMPLE") def tunnel_ready_to_send(_): backend.start_receiving() def tunnel_send_buffer_full(_): backend.stop_receiving() def tunnel_received(_, data, _addr): backend.send(data) return backend.is_ready_to_send() def tunnel_closed(_): backend.close() def backend_received(_, data, _addr): self._stream.send(data) return self._stream.is_ready_to_send() def backend_closed(_self): self._stream.close() self._stream.set_on_ready_to_send(tunnel_ready_to_send) self._stream.set_on_send_buffer_full(tunnel_send_buffer_full) self._stream.set_on_received(tunnel_received) self._stream.set_on_closed(tunnel_closed) backend.set_on_received(backend_received) backend.set_on_closed(backend_closed) if received is not None and len(received) > 0: backend.send(received) backend.connect(UNKNOWN_CONN_ADDR, UNKNOWN_CONN_PORT)