mujson_mgr.py

#!/usr/bin/python
# -*- coding: UTF-8 -*-

import traceback
from shadowsocks import shell, common
from configloader import load_config, get_config
import random
import getopt
import sys
import json
import base64


class MuJsonLoader(object):
	def __init__(self):
		self.json = None

	def load(self, path):
		l = "[]"
		try:
			with open(path, 'rb+') as f:
				l = f.read().decode('utf8')
		except:
			pass
		self.json = json.loads(l)

	def save(self, path):
		if self.json is not None:
			output = json.dumps(self.json, sort_keys=True, indent=4, separators=(',', ': '))
			with open(path, 'a'):
				pass
			with open(path, 'rb+') as f:
				f.write(output.encode('utf8'))
				f.truncate()


class MuMgr(object):
	def __init__(self):
		self.config_path = get_config().MUDB_FILE
		try:
			self.server_addr = get_config().SERVER_PUB_ADDR
		except:
			self.server_addr = '127.0.0.1'
		self.data = MuJsonLoader()

		if self.server_addr == '127.0.0.1':
			self.server_addr = self.getipaddr()

	def getipaddr(self, ifname='eth0'):
		import socket
		import struct
		ret = '127.0.0.1'
		try:
			ret = socket.gethostbyname(socket.getfqdn(socket.gethostname()))
		except:
			pass
		if ret == '127.0.0.1':
			try:
				import fcntl
				s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
				ret = socket.inet_ntoa(fcntl.ioctl(s.fileno(), 0x8915, struct.pack('256s', ifname[:15]))[20:24])
			except:
				pass
		return ret

	def ssrlink(self, user, encode, muid):
		protocol = user.get('protocol', '')
		obfs = user.get('obfs', '')
		protocol = protocol.replace("_compatible", "")
		obfs = obfs.replace("_compatible", "")
		protocol_param = ''
		if muid is not None:
			protocol_param_ = user.get('protocol_param', '')
			param = protocol_param_.split('#')
			if len(param) == 2:
				for row in self.data.json:
					if int(row['port']) == muid:
						param = str(muid) + ':' + row['passwd']
						protocol_param = '/?protoparam=' + common.to_str(base64.urlsafe_b64encode(common.to_bytes(param))).replace("=", "")
						break
		link = ("%s:%s:%s:%s:%s:%s" % (self.server_addr, user['port'], protocol, user['method'], obfs, common.to_str(base64.urlsafe_b64encode(common.to_bytes(user['passwd']))).replace("=", ""))) + protocol_param
		return "ssr://" + (encode and common.to_str(base64.urlsafe_b64encode(common.to_bytes(link))).replace("=", "") or link)

	def userinfo(self, user, muid = None):
		ret = ""
		key_list = ['user', 'port', 'method', 'passwd', 'protocol', 'protocol_param', 'obfs', 'obfs_param', 'transfer_enable', 'u', 'd']
		for key in sorted(user):
			if key not in key_list:
				key_list.append(key)
		for key in key_list:
			if key in ['enable'] or key not in user:
				continue
			ret += '\n'
			if (muid is not None) and (key in ['protocol_param']):
				for row in self.data.json:
					if int(row['port']) == muid:
						ret += "    %s : %s" % (key, str(muid) + ':' + row['passwd'])
						break
			elif key in ['transfer_enable', 'u', 'd']:
				if muid is not None:
					for row in self.data.json:
						if int(row['port']) == muid:
							val = row[key]
							break
				else:
					val = user[key]
				if val / 1024 < 4:
					ret += "    %s : %s" % (key, val)
				elif val / 1024 ** 2 < 4:
					val /= float(1024)
					ret += "    %s : %s  K Bytes" % (key, val)
				elif val / 1024 ** 3 < 4:
					val /= float(1024 ** 2)
					ret += "    %s : %s  M Bytes" % (key, val)
				else:
					val /= float(1024 ** 3)
					ret += "    %s : %s  G Bytes" % (key, val)
			else:
				ret += "    %s : %s" % (key, user[key])
		ret += "\n    " + self.ssrlink(user, False, muid)
		ret += "\n    " + self.ssrlink(user, True, muid)
		return ret

	def rand_pass(self):
		return ''.join([random.choice('''ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789~-_=+(){}[]^&%$@''') for i in range(8)])

	def add(self, user):
		up = {'enable': 1, 'u': 0, 'd': 0, 'method': "aes-128-ctr",
		'protocol': "auth_aes128_md5",
		'obfs': "tls1.2_ticket_auth_compatible",
		'transfer_enable': 9007199254740992}
		up['passwd'] = self.rand_pass()
		up.update(user)

		self.data.load(self.config_path)
		for row in self.data.json:
			match = False
			if 'user' in user and row['user'] == user['user']:
				match = True
			if 'port' in user and row['port'] == user['port']:
				match = True
			if match:
				print("user [%s] port [%s] already exist" % (row['user'], row['port']))
				return
		self.data.json.append(up)
		print("### add user info %s" % self.userinfo(up))
		self.data.save(self.config_path)

	def edit(self, user):
		self.data.load(self.config_path)
		for row in self.data.json:
			match = True
			if 'user' in user and row['user'] != user['user']:
				match = False
			if 'port' in user and row['port'] != user['port']:
				match = False
			if match:
				print("edit user [%s]" % (row['user'],))
				row.update(user)
				print("### new user info %s" % self.userinfo(row))
				break
		self.data.save(self.config_path)

	def delete(self, user):
		self.data.load(self.config_path)
		index = 0
		for row in self.data.json:
			match = True
			if 'user' in user and row['user'] != user['user']:
				match = False
			if 'port' in user and row['port'] != user['port']:
				match = False
			if match:
				print("delete user [%s]" % row['user'])
				del self.data.json[index]
				break
			index += 1
		self.data.save(self.config_path)

	def clear_ud(self, user):
		up = {'u': 0, 'd': 0}
		self.data.load(self.config_path)
		for row in self.data.json:
			match = True
			if 'user' in user and row['user'] != user['user']:
				match = False
			if 'port' in user and row['port'] != user['port']:
				match = False
			if match:
				row.update(up)
				print("clear user [%s]" % row['user'])
		self.data.save(self.config_path)

	def list_user(self, user):
		self.data.load(self.config_path)
		if not user:
			for row in self.data.json:
				print("user [%s] port %s" % (row['user'], row['port']))
			return
		for row in self.data.json:
			match = True
			if 'user' in user and row['user'] != user['user']:
				match = False
			if 'port' in user and row['port'] != user['port']:
				match = False
			if match:
				muid = None
				if 'muid' in user:
					muid = user['muid']
				print("### user [%s] info %s" % (row['user'], self.userinfo(row, muid)))


def print_server_help():
	print('''usage: python mujson_manage.py -a|-d|-e|-c|-l [OPTION]...

Actions:
  -a ADD               add/edit a user
  -d DELETE            delete a user
  -e EDIT              edit a user
  -c CLEAR             set u/d to zero
  -l LIST              display a user infomation or all users infomation

Options:
  -u USER              the user name
  -p PORT              server port (only this option must be set if add a user)
  -k PASSWORD          password
  -m METHOD            encryption method, default: aes-128-ctr
  -O PROTOCOL          protocol plugin, default: auth_aes128_md5
  -o OBFS              obfs plugin, default: tls1.2_ticket_auth_compatible
  -G PROTOCOL_PARAM    protocol plugin param
  -g OBFS_PARAM        obfs plugin param
  -t TRANSFER          max transfer for G bytes, default: 8388608 (8 PB or 8192 TB)
  -f FORBID            set forbidden ports. Example (ban 1~79 and 81~100): -f "1-79,81-100"
  -i MUID              set sub id to display (only work with -l)

General options:
  -h, --help           show this help message and exit
''')


def main():
	shortopts = 'adeclu:i:p:k:O:o:G:g:m:t:f:hs:S:'
	longopts = ['help']
	action = None
	user = {}
	fast_set_obfs = {'0': 'plain',
			'+1': 'http_simple_compatible',
			'1': 'http_simple',
			'+2': 'tls1.2_ticket_auth_compatible',
			'2': 'tls1.2_ticket_auth'}
	fast_set_protocol = {'0': 'origin',
			's4': 'auth_sha1_v4',
			'+s4': 'auth_sha1_v4_compatible',
			'am': 'auth_aes128_md5',
			'as': 'auth_aes128_sha1',
			'ca': 'auth_chain_a',
			}
	fast_set_method = {'0': 'none',
			'a1c': 'aes-128-cfb',
			'a2c': 'aes-192-cfb',
			'a3c': 'aes-256-cfb',
			'r': 'rc4-md5',
			'r6': 'rc4-md5-6',
			'c': 'chacha20',
			'ci': 'chacha20-ietf',
			's': 'salsa20',
			'a1': 'aes-128-ctr',
			'a2': 'aes-192-ctr',
			'a3': 'aes-256-ctr'}
	try:
		optlist, args = getopt.getopt(sys.argv[1:], shortopts, longopts)
		for key, value in optlist:
			if key == '-a':
				action = 1
			elif key == '-d':
				action = 2
			elif key == '-e':
				action = 3
			elif key == '-l':
				action = 4
			elif key == '-c':
				action = 0
			elif key == '-u':
				user['user'] = value
			elif key == '-i':
				user['muid'] = int(value)
			elif key == '-p':
				user['port'] = int(value)
			elif key == '-k':
				user['passwd'] = value
			elif key == '-o':
				if value in fast_set_obfs:
					user['obfs'] = fast_set_obfs[value]
				else:
					user['obfs'] = value
			elif key == '-O':
				if value in fast_set_protocol:
					user['protocol'] = fast_set_protocol[value]
				else:
					user['protocol'] = value
			elif key == '-g':
				user['obfs_param'] = value
			elif key == '-G':
				user['protocol_param'] = value
			elif key == '-s':
				user['speed_limit_per_con'] = int(value)
			elif key == '-S':
				user['speed_limit_per_user'] = int(value)
			elif key == '-m':
				if value in fast_set_method:
					user['method'] = fast_set_method[value]
				else:
					user['method'] = value
			elif key == '-f':
				user['forbidden_port'] = value
			elif key == '-t':
				val = float(value)
				try:
					val = int(value)
				except:
					pass
				user['transfer_enable'] = int(val * 1024) * (1024 ** 2)
			elif key in ('-h', '--help'):
				print_server_help()
				sys.exit(0)
	except getopt.GetoptError as e:
		print(e)
		sys.exit(2)

	manage = MuMgr()
	if action == 0:
		manage.clear_ud(user)
	elif action == 1:
		if 'user' not in user and 'port' in user:
			user['user'] = str(user['port'])
		if 'user' in user and 'port' in user:
			manage.add(user)
		else:
			print("You have to set the port with -p")
	elif action == 2:
		if 'user' in user or 'port' in user:
			manage.delete(user)
		else:
			print("You have to set the user name or port with -u/-p")
	elif action == 3:
		if 'user' in user or 'port' in user:
			manage.edit(user)
		else:
			print("You have to set the user name or port with -u/-p")
	elif action == 4:
		manage.list_user(user)
	elif action is None:
		print_server_help()

if __name__ == '__main__':
	main()