add control manager
This commit is contained in:
clowwindy
parent
58df1d82d0
commit
956199efcd
|
|
@ -0,0 +1,152 @@
|
|||
#!/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 shadowsocks import common, eventloop, tcprelay, udprelay, asyncdns, shell
|
||||
|
||||
|
||||
BUF_SIZE = 2048
|
||||
|
||||
|
||||
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._control_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM,
|
||||
socket.IPPROTO_UDP)
|
||||
self._statistics = collections.defaultdict(int)
|
||||
self._control_client_addr = None
|
||||
try:
|
||||
self._control_socket.bind(('127.0.0.1',
|
||||
int(config['manager_port'])))
|
||||
self._control_socket.setblocking(False)
|
||||
except (OSError, IOError) as e:
|
||||
logging.error(e)
|
||||
logging.error('can not bind to manager port')
|
||||
exit(1)
|
||||
self._loop.add(self._control_socket,
|
||||
eventloop.POLL_IN, self)
|
||||
|
||||
port_password = config['port_password']
|
||||
del config['port_password']
|
||||
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)
|
||||
u = udprelay.UDPRelay(config, self._dns_resolver, False)
|
||||
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:
|
||||
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)
|
||||
elif command == 'remove':
|
||||
self.remove_port(a_config)
|
||||
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 = json.loads(config_json)
|
||||
return command, config
|
||||
except Exception as e:
|
||||
logging.error(e)
|
||||
return None
|
||||
|
||||
def handle_periodic(self):
|
||||
# TODO send statistics
|
||||
pass
|
||||
|
||||
def _send_control_data(self, data):
|
||||
if self._control_client_addr:
|
||||
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()
|
||||
|
|
@ -24,7 +24,8 @@ import logging
|
|||
import signal
|
||||
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../'))
|
||||
from shadowsocks import shell, daemon, eventloop, tcprelay, udprelay, asyncdns
|
||||
from shadowsocks import shell, daemon, eventloop, tcprelay, udprelay, \
|
||||
asyncdns, manager
|
||||
|
||||
|
||||
def main():
|
||||
|
|
@ -48,10 +49,17 @@ def main():
|
|||
else:
|
||||
config['port_password'][str(server_port)] = config['password']
|
||||
|
||||
if config['manager_port']:
|
||||
logging.info('entering manager mode')
|
||||
manager.run(config)
|
||||
return
|
||||
|
||||
tcp_servers = []
|
||||
udp_servers = []
|
||||
dns_resolver = asyncdns.DNSResolver()
|
||||
for port, password in config['port_password'].items():
|
||||
port_password = config['port_password']
|
||||
del config['port_password']
|
||||
for port, password in port_password.items():
|
||||
a_config = config.copy()
|
||||
a_config['server_port'] = int(port)
|
||||
a_config['password'] = password
|
||||
|
|
|
|||
|
|
@ -136,7 +136,7 @@ def get_config(is_local):
|
|||
else:
|
||||
shortopts = 'hd:s:p:k:m:c:t:vq'
|
||||
longopts = ['help', 'fast-open', 'pid-file=', 'log-file=', 'workers=',
|
||||
'forbidden-ip=', 'user=', 'version']
|
||||
'forbidden-ip=', 'user=', 'manager-port=', 'version']
|
||||
try:
|
||||
config_path = find_config()
|
||||
optlist, args = getopt.getopt(sys.argv[1:], shortopts, longopts)
|
||||
|
|
@ -181,6 +181,8 @@ def get_config(is_local):
|
|||
config['fast_open'] = True
|
||||
elif key == '--workers':
|
||||
config['workers'] = int(value)
|
||||
elif key == '--manager-port':
|
||||
config['manager_port'] = int(value)
|
||||
elif key == '--user':
|
||||
config['user'] = to_str(value)
|
||||
elif key == '--forbidden-ip':
|
||||
|
|
@ -317,6 +319,7 @@ Proxy options:
|
|||
--fast-open use TCP_FASTOPEN, requires Linux 3.7+
|
||||
--workers WORKERS number of workers, available on Unix/Linux
|
||||
--forbidden-ip IPLIST comma seperated IP list forbidden to connect
|
||||
--manager-port PORT optional server manager UDP port
|
||||
|
||||
General options:
|
||||
-h, --help show this help message and exit
|
||||
|
|
|
|||
|
|
@ -708,3 +708,5 @@ class TCPRelay(object):
|
|||
self._eventloop.remove_periodic(self.handle_periodic)
|
||||
self._eventloop.remove(self._server_socket)
|
||||
self._server_socket.close()
|
||||
for handler in list(self._fd_to_handlers.values()):
|
||||
handler.destroy()
|
||||
|
|
|
|||
|
|
@ -290,3 +290,5 @@ class UDPRelay(object):
|
|||
self._eventloop.remove_periodic(self.handle_periodic)
|
||||
self._eventloop.remove(self._server_socket)
|
||||
self._server_socket.close()
|
||||
for client in list(self._cache.values()):
|
||||
client.close()
|
||||
|
|
|
|||
Loading…
Reference in New Issue