text: #!/usr/bin/python
  # -*- coding: utf-8 -*-
  #
  # Copyright 2015 clowwindy
  #
  # Licensed under the Apache License, Version 2.0 (the "License"); you may
  # not use this file except in compliance with the License. You may obtain
  # a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
  # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
  # License for the specific language governing permissions and limitations
  # under the License.
  
  from __future__ import absolute_import, division, print_function, \
      with_statement
  
  import errno
  import traceback
  import socket
  import logging
  import json
  import collections
  
  from enuma_elish import common, eventloop, tcprelay, udprelay, asyncdns, shell
  
  
  BUF_SIZE = 1506
  STAT_SEND_LIMIT = 50
  
  
  class Manager(object):
  
      def __init__(self, config):
          self._config = config
          self._relays = {}  # (tcprelay, udprelay)
          self._loop = eventloop.EventLoop()
          self._dns_resolver = asyncdns.DNSResolver()
          self._dns_resolver.add_to_loop(self._loop)
  
          self._statistics = collections.defaultdict(int)
          self._control_client_addr = None
          try:
              manager_address = config['manager_address']
              if ':' in manager_address:
                  addr = manager_address.rsplit(':', 1)
                  addr = addr[0], int(addr[1])
                  addrs = socket.getaddrinfo(addr[0], addr[1])
                  if addrs:
                      family = addrs[0][0]
                  else:
                      logging.error('invalid address: %s', manager_address)
                      exit(1)
              else:
                  addr = manager_address
                  family = socket.AF_UNIX
              self._control_socket = socket.socket(family,
                                                   socket.SOCK_DGRAM)
              self._control_socket.bind(addr)
              self._control_socket.setblocking(False)
          except (OSError, IOError) as e:
              logging.error(e)
              logging.error('can not bind to manager address')
              exit(1)
          self._loop.add(self._control_socket,
                         eventloop.POLL_IN, self)
          self._loop.add_periodic(self.handle_periodic)
  
          port_password = config['port_password']
          del config['port_password']
          config['crypto_path'] = config.get('crypto_path', dict())
          for port, password in port_password.items():
              a_config = config.copy()
              a_config['server_port'] = int(port)
              a_config['password'] = password
              self.add_port(a_config)
  
      def add_port(self, config):
          port = int(config['server_port'])
          servers = self._relays.get(port, None)
          if servers:
              logging.error("server already exists at %s:%d" % (config['server'],
                                                                port))
              return
          logging.info("adding server at %s:%d" % (config['server'], port))
          t = tcprelay.TCPRelay(config, self._dns_resolver, False,
                                self.stat_callback)
          u = udprelay.UDPRelay(config, self._dns_resolver, False,
                                self.stat_callback)
          t.add_to_loop(self._loop)
          u.add_to_loop(self._loop)
          self._relays[port] = (t, u)
  
      def remove_port(self, config):
          port = int(config['server_port'])
          servers = self._relays.get(port, None)
          if servers:
              logging.info("removing server at %s:%d" % (config['server'], port))
              t, u = servers
              t.close(next_tick=False)
              u.close(next_tick=False)
              del self._relays[port]
          else:
              logging.error("server not exist at %s:%d" % (config['server'],
                                                           port))
  
      def handle_event(self, sock, fd, event):
          if sock == self._control_socket and event == eventloop.POLL_IN:
              data, self._control_client_addr = sock.recvfrom(BUF_SIZE)
              parsed = self._parse_command(data)
              if parsed:
                  command, config = parsed
                  a_config = self._config.copy()
                  if config:
                      # let the command override the configuration file
                      a_config.update(config)
                  if 'server_port' not in a_config:
                      logging.error('can not find server_port in config')
                  else:
                      if command == 'add':
                          self.add_port(a_config)
                          self._send_control_data(b'ok')
                      elif command == 'remove':
                          self.remove_port(a_config)
                          self._send_control_data(b'ok')
                      elif command == 'ping':
                          self._send_control_data(b'pong')
                      else:
                          logging.error('unknown command %s', command)
  
      def _parse_command(self, data):
          # commands:
          # add: {"server_port": 8000, "password": "foobar"}
          # remove: {"server_port": 8000"}
          data = common.to_str(data)
          parts = data.split(':', 1)
          if len(parts) < 2:
              return data, None
          command, config_json = parts
          try:
              config = shell.parse_json_in_str(config_json)
              if 'method' in config:
                  config['method'] = common.to_str(config['method'])
              return command, config
          except Exception as e:
              logging.error(e)
              return None
  
      def stat_callback(self, port, data_len):
          self._statistics[port] += data_len
  
      def handle_periodic(self):
          r = {}
          i = 0
  
          def send_data(data_dict):
              if data_dict:
                  # use compact JSON format (without space)
                  data = common.to_bytes(json.dumps(data_dict,
                                                    separators=(',', ':')))
                  self._send_control_data(b'stat: ' + data)
  
          for k, v in self._statistics.items():
              r[k] = v
              i += 1
              # split the data into segments that fit in UDP packets
              if i >= STAT_SEND_LIMIT:
                  send_data(r)
                  r.clear()
                  i = 0
          if len(r) > 0:
              send_data(r)
          self._statistics.clear()
  
      def _send_control_data(self, data):
          if not self._control_client_addr:
              return
  
          try:
              self._control_socket.sendto(data, self._control_client_addr)
          except (socket.error, OSError, IOError) as e:
              error_no = eventloop.errno_from_exception(e)
              if error_no in (errno.EAGAIN, errno.EINPROGRESS,
                              errno.EWOULDBLOCK):
                  return
              else:
                  shell.print_exception(e)
                  if self._config['verbose']:
                      traceback.print_exc()
  
      def run(self):
          self._loop.run()
  
  
  def run(config):
      Manager(config).run()
  
  
  def test():
      import time
      import threading
      import struct
      from enuma_elish import cryptor
  
      logging.basicConfig(level=5,
                          format='%(asctime)s %(levelname)-8s %(message)s',
                          datefmt='%Y-%m-%d %H:%M:%S')
      enc = []
      eventloop.TIMEOUT_PRECISION = 1
  
      def run_server():
          config = {
              'server': '127.0.0.1',
              'local_port': 1081,
              'port_password': {
                  '8381': 'foobar1',
                  '8382': 'foobar2'
              },
              'method': 'aes-256-cfb',
              'manager_address': '127.0.0.1:6001',
              'timeout': 60,
              'fast_open': False,
              'verbose': 2
          }
          manager = Manager(config)
          enc.append(manager)
          manager.run()
  
      t = threading.Thread(target=run_server)
      t.start()
      time.sleep(1)
      manager = enc[0]
      cli = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
      cli.connect(('127.0.0.1', 6001))
  
      # test add and remove
      time.sleep(1)
      cli.send(b'add: {"server_port":7001, "password":"asdfadsfasdf"}')
      time.sleep(1)
      assert 7001 in manager._relays
      data, addr = cli.recvfrom(1506)
      assert b'ok' in data
  
      cli.send(b'remove: {"server_port":8381}')
      time.sleep(1)
      assert 8381 not in manager._relays
      data, addr = cli.recvfrom(1506)
      assert b'ok' in data
      logging.info('add and remove test passed')
  
      # test statistics for TCP
      header = common.pack_addr(b'google.com') + struct.pack('>H', 80)
      data = cryptor.encrypt_all(b'asdfadsfasdf', 'aes-256-cfb',
                                 header + b'GET /\r\n\r\n')
      tcp_cli = socket.socket()
      tcp_cli.connect(('127.0.0.1', 7001))
      tcp_cli.send(data)
      tcp_cli.recv(4096)
      tcp_cli.close()
  
      data, addr = cli.recvfrom(1506)
      data = common.to_str(data)
      assert data.startswith('stat: ')
      data = data.split('stat:')[1]
      stats = shell.parse_json_in_str(data)
      assert '7001' in stats
      logging.info('TCP statistics test passed')
  
      # test statistics for UDP
      header = common.pack_addr(b'127.0.0.1') + struct.pack('>H', 80)
      data = cryptor.encrypt_all(b'foobar2', 'aes-256-cfb',
                                 header + b'test')
      udp_cli = socket.socket(type=socket.SOCK_DGRAM)
      udp_cli.sendto(data, ('127.0.0.1', 8382))
      tcp_cli.close()
  
      data, addr = cli.recvfrom(1506)
      data = common.to_str(data)
      assert data.startswith('stat: ')
      data = data.split('stat:')[1]
      stats = json.loads(data)
      assert '8382' in stats
      logging.info('UDP statistics test passed')
  
      manager._loop.stop()
      t.join()
  
  
  if __name__ == '__main__':
      test()
  