python实现清华大学联网助手(三)——linux下带参数的自定义命令/import机制

上两篇已经大致满足了我起初设计的需求。

最后，对代码进行了一些优化，并且修改代码使一些函数方法可以重用。最终实现了如下功能：登录校园网、登出校园网、查询当前网络状态、查询流量、帐户余额、当前用户组等基本信息、查询本账号当前在线IP信息、查询每日流量使用明细并生成直方图。

本篇主要探讨两个问题：

1 如何将.py脚本制作成linux下带参数的自定义命令；

2 import机制及使用。

1 如何将.py脚本制作成linux下带参数的自定义命令：

具体要求：

咱们有1个py脚本文件，它可以接受很多参数。但现在要执行它，必须先cd到.py文件的目录下，然后python pytunet.py -i -a --check。当然可传入的参数有很多，比如-i -o -a -h --check等等，每一次执行都可能是这些参数的任意组合。但操作这么麻烦，违背了初衷。要是这样还不如打开浏览器登录我的网费查询账户呢。因此我们迫切需要一个方法，让它成为一个自定义命令。

解决办法：

a) 写一个shell脚本，将参数传给pyhton脚本，然后将shell脚本放在自己的bin下。

显然这个方法可以实现，但比较麻烦，难道没有更好的办法了吗？

b) 修改linux的默认python解释器，#!/usr/bin/env python，给主脚本修改755执行权限，将alias写到.bashrc中去。

可能这么一说我们并不懂，我们一条一条来看。

①修改linux的默认python解释器。

因为一般linux都是预装python2的版本，而且系统内一些例行程序都需要python2，因此我们并不能删除。首先我们要安装python3.4。再删除原来/usr/bin下的python符号链接，再将python3.4设置为新的python符号链接。

sudo rm /usr/bin/python

sudo ln -s /usr/bin/python3.4 /usr/bin/python

这样我们的python3.4就是默认的解释器了。

②将#!/usr/bin/env python作为你脚本文件中的第一句。

这样的话，我们就不必为脚本指定python3.4作为运行该脚本的解释器。因为开头包含#符号，所以一般情况下被视作注释。在执行脚本时，如果指定了解释器则忽视这段代码。如果没指定，就默认使用该代码所指定的解释器。

因此python ./pytunet.py -q也就可以写作./pytunet.py -q了。

这样离我们的目标更近了一些。

③给主脚本设置755权限。

一般文件的权限是rwxr--r--类似这样的形式。分别表示root用户, 同组用户和其他用户的权限。r为读权限，w为写权限，x为执行权限，-表示没有这种权限。同样为了更为简单的表示，我们用二进制数来代表响应的权限，如rwx为111，r-x为101，对应的十进制数rwx为7，r-x为5。因此所谓的755权限是rwxr-xr-x。一般而言，在我们普通的私人电脑中，设置755权限是没有问题的。因此可以放心。

那么我们可以这样设置主脚本的权限：sudo chmod 755 pytunet.py。

④将alias写到.bashrc。

sudo权限打开用户路径下$HOME/.bashrc隐藏文件。选一个合适的位置加上如下alias：

alias tunet=‘/home/coderhuhy/mycommands/pytunet.py’，这里是你自己保存py脚本的位置。tunet为其别名。

保存并关闭.bashrc。注意需要重启或注销电脑，或者source ~/.bashrc。

由此一来你就可以使用tunet作为代号，完成所有工作。我们再也不用登录网费查询系统了。只需alt+ctrl+T调出terminal，输入tunet -q，就查询所有网费情况、流量使用明细等。是不是很炫酷。

2 import机制及使用：

因为有了多个.py文件，这就涉及到各个py文件之间的调用。

import是载入模块的关键字。我们可以import另一个脚本，并带全名地使用里面的函数。也可以from pytunet_query import function。区别在于前者需要全名，比如使用pytunet_query中的function方法，就需要pytunet_query.function()；但后者可以直接用function()。

个人建议，尽量使用import而不要使用from...import，原因在于前者能够区分是本脚本的方法还是本脚本外的方法，这样有利于理清逻辑关系。

最后，将我的最终v1.2版本的清华大学联网助手python版实现代码放上来。

pytunet_query.py：

import sys, time, os
import urllib.request, hashlib, http.cookiejar
import codecs, re

query_login_url  = 'https://usereg.tsinghua.edu.cn/do.php'
user_info_url    = 'https://usereg.tsinghua.edu.cn/user_info.php'
online_state_url = 'https://usereg.tsinghua.edu.cn/online_user_ipv4.php'
query_logout_url = 'https://usereg.tsinghua.edu.cn/do.php'

info_header = [ '#' * 89,
				'#\t\t\t\tUser Flux Detail Infomation\t\t\t\t#',
				'#' * 89]

#########################################################
#					File I/O Modules					#
#########################################################

def create_path(relative_path):
 	base_path = os.path.abspath(os.path.dirname(sys.argv[0]))
 	return os.path.join(base_path, relative_path)

def write_inline(file_handler, contents):
	for line in contents:
		file_handler.write(line + '\r\n')

def save_query(contents):
	relative_path = 'USER_DETAIL_INFOMATION.LOG'
	file_handler = open(create_path(relative_path), 'w')
	write_inline(file_handler, info_header)
	file_handler.write('\t\t\t\tDatetime: ' + time.strftime('%Y-%m-%d %H:%M:%S') + '\r\n' + '-' * 89 + '\r\n')
	write_inline(file_handler, contents)
	file_handler.write('\r\n')
	file_handler.close()

#########################################################
#					Connection Modules					#
#########################################################

def create_opener():
	cookie = http.cookiejar.CookieJar()
	cookie_proc = urllib.request.HTTPCookieProcessor(cookie)
	return urllib.request.build_opener(cookie_proc)

def response_login(login_data):
	request_url = urllib.request.Request(query_login_url, login_data.encode())
	response_url = urllib.request.urlopen(request_url)
	return response_url.read().decode()

#########################################################
#				Main Login/Logout Modules				#
#########################################################

def query_login(username, password):
	hashcd_md5 = hashlib.md5()
	hashcd_md5.update(password.encode())
	tr_password = hashcd_md5.hexdigest()
	login_data = 'user_login_name=' + username + '&user_password=' + tr_password + '&action=login'
	urllib.request.install_opener(create_opener())
	answer = response_login(login_data)
	if answer == 'ok':
		return True
	else:
		return False

def query_logout():
	logout_data = 'action=logout'
	request_url = urllib.request.Request(query_logout_url, logout_data.encode())
	response_url = urllib.request.urlopen(request_url)
	print ('Your flux details and other infomations are saved in USER_DETAIL_INFOMATION.LOG under the SAME directory')

#########################################################
#				Data Post-Process Modules				#
#########################################################

def post_process(info):
	end_time = time.strftime('%Y-%m-%d')
	start_time = end_time[:8:] + '01'
	flux_detail_url = 'https://usereg.tsinghua.edu.cn/user_detail_list.php?action=balance2&user_login_name=&user_real_name=&desc=&order=&start_time=' + start_time + '&end_time=' + end_time + '&user_ip=&user_mac=&nas_ip=&nas_port=&is_ipv6=0&page=1&offset=200'	

	response_usr = urllib.request.urlopen(user_info_url)
	response_state = urllib.request.urlopen(online_state_url)
	response_details = urllib.request.urlopen(flux_detail_url)

	info = flux_account_query(info, response_usr)
	info = online_state_query(info, response_state)
	info = flux_detail_query(info, response_details)
	
	return info

#########################################################
#				Integrated Query Modules				#
#########################################################

flux_account_keys = ('用户名', '用户组', '姓名', '证件号', '当前计费组', '使用时长(IPV4)', '使用流量(IPV4)',
					 '使用时长(IPV6)', '使用流量(IPV6)', '帐户余额')
online_state_keys = ()
flux_detail_keys  = ()

#Auxiliary Function
def turn_key(key):
	if key[-5:-1] == 'byte':
		flux, unit = key.split('(')
		flux = float(flux) / 1024 / 1024
		new_key = '-->' + str(int(flux)) + '(MB)'
		key += new_key
	return key


def get_days(year, month):
	month_length = (31,28,31,30,31,30,31,31,30,31,30,31)
	month_length_leap = (31,29,31,30,31,30,31,31,30,31,30,31)
	if year % 400 == 0 or year % 100 != 0 and year % 4 == 0:
		return month_length_leap[month-1]
	else:
		return month_length[month-1]

def solve_flux(flux):
	unit = flux[-1]
	val = float(flux[:len(flux)-1:])

	if unit == 'B':
		val /= 1024 * 1024
	elif unit == 'K':
		val /= 1024
	elif unit == 'G':
		val *= 1024

	return int(val)

def trans_content(response):
	raw = response.read().decode('gb2312')
	raw = re.sub('<[^>]+>| |[\n\t]+|-->',' ',raw)
	raw = re.sub(' +', ' ', raw)
	return raw

def push_front(figure, line):
	tf = []
	tf.append(line)
	return tf + figure

def display_fluxAccount_onlineState(info):
	print()
	for line in info:
		if line[0] != '-':
			print(line)
		else:
			print()

def display_flux_detail(fluxin, year, month, day):
	maxflux = 0
	divide = 10
	figure = []
	for flux in fluxin:
		if flux > maxflux:
			maxflux = flux

	top = str(int(maxflux)) + 'MB|'
	length = len(top)
	mid = str(int(maxflux / 2)) + 'MB|'
	mid = ' ' * (length - len(mid)) + mid
	bottom = '0MB|'
	bottom = ' ' * (length - len(bottom)) + bottom
	unit = maxflux / divide

	for i in range(day):
		fluxin[i] = int(fluxin[i] / unit)

	for i in range(divide):
		line = ''
		if i == divide - 1:
			line = top
		elif i == int((divide - 1) / 2):
			line = mid
		elif i == 0:
			line = bottom
		else:
			line = ' ' * (length - 1) + '|'
		for j in range(day):
			if fluxin[j] > 0:
				line += '**'
				fluxin[j] -= 1
			else:
				line += '  '

		figure = push_front(figure, line)

	figure = push_front(figure, '**每日流量使用统计列表**')
	figure.append(' ' * length + '--' * day)
	date_front = str(year) + '-' + str(month) +'-' + '1'
	date_rear  = str(year) + '-' + str(month) +'-' + str(day)
	date_mid   = str(year) + '-' + str(month) +'-' + '15'
	figure.append(' %s\t\t\t%s\t\t\t%s' %(date_front, date_mid, date_rear))

	for line in figure:
		print(line)

	print()

#Integrated Query
def flux_account_query(info, response):
	info.append('**用户基本信息**')
	done = trans_content(response)
	match = re.search('用户名.*?(元) ', done)
	done = match.group()
	tlist = done.split(' ')
	line = ''

	for key in tlist:
		if line != '':
			key = turn_key(key)
			line = line + '\t: ' + key
			info.append(line)
			line = ''
		elif key in flux_account_keys:
			line = key

	info.append('-' * 89)
	return info

def online_state_query(info, response):
	info.append('**用户在线状态**')
	
	done = trans_content(response)
	match = re.search('\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}.*\d{2}:\d{2}', done)

	if match == None:
		info.append('当前没有任何IP在线')
	else:
		info.append('在线IP地址\t登陆日期\t登陆时间')
		done = match.group()
		tlist = done.split(' ')
		line = ''
		count = 0

		for key in tlist:
			if count == 0:
				if re.search('\d{1,3}\.\d{1,3}\.\d{1,3}.\d{1,3}', key) != None:
					line = key
					count += 1
			elif count == 1:
				if re.search('\d{4}-\d{2}-\d{2}', key) != None:
					line += '\t' + key
					count += 1
			elif count == 2:
				if re.search('\d{2}:\d{2}:\d{2}', key) != None:
					line += '\t' + key
					info.append(line)
					line = ''
					count = 0

	info.append('-' * 89)
	display_fluxAccount_onlineState(info);
	return info

def flux_detail_query(info, response):
	info.append('**每日流量使用统计列表**')
	info.append('登出日期\t入流量\t出流量')
	done = trans_content(response)
	year, month = time.strftime('%Y %m').split(' ')
	days = get_days(int(year), int(month))
	tlist = done.split(' ')

	fluxin_perday = [0 for i in range(days)]
	fluxout_perday = [0 for i in range(days)]
	offline_date = True
	count = 0

	for key in tlist:
		if re.search('\d{4}-\d{2}-\d{2}', key) and offline_date:
			offline_date = False
		elif re.search('\d{4}-\d{2}-\d{2}', key) and not offline_date:
			offline_date = True
			year, month, day = key.split('-')
			iday = int(day)
		elif re.search('\d+[.]\d*[BKMG]', key) and count == 0:
			fluxin_perday[iday-1] += solve_flux(key)
			count += 1
		elif re.search('\d+[.]\d*[BKMG]', key) and count == 1:
			fluxout_perday[iday-1] += solve_flux(key)
			count += 1
		elif re.search('\d+[.]\d*[BKMG]', key) and count == 2:
			count = 0
	
	for i in range(days):
		if i + 1 < 10:
			d = '0' + str(i + 1)
		else:
			d = str(i + 1)
		info.append('%s\t%s\t%s' %(time.strftime('%Y-%m-') + d, str(fluxin_perday[i]) + 'MB', str(fluxout_perday[i]) + 'MB'))

	display_flux_detail(fluxin_perday, int(year), int(month), days)
	return info

#########################################################
#						Main Part						#
#########################################################

def tunet_query(username, password):
	print('FETCHING DATA FROM http://usereg.tsinghua.edu.cn, PLEASE WAIT FOR A MOMENT...')
	is_login = query_login(username, password)
	if is_login:
		info = []
		info = post_process(info)
		save_query(info)
		query_logout()
	else:
		print ('CAN\'T CAPTURE YOUR FLUX DATA, PLEASE TRY AGAIN LATER')

def pytunet_query():
	username = 'hhy14'
	password = '123456'
	tunet_query(username, password)

if __name__ == '__main__':
	pytunet_query()

pytunet_connect.py：

import time
import urllib.request, hashlib
import codecs

login_url  = 'http://net.tsinghua.edu.cn/cgi-bin/do_login'
logout_url = 'http://net.tsinghua.edu.cn/cgi-bin/do_logout'
check_url  = 'http://net.tsinghua.edu.cn/cgi-bin/do_login'
query_url  = 'https://usereg.tsinghua.edu.cn/login.php'

times_cnt = {1: 'FIRST', 2: 'SECOND', 3: 'THIRD', 4: 'FORTH', 5: 'FIFTH'}
ret_type  = {'logout_ok'       : 'LOGOUT SUCCESS',
			'not_online_error' : 'NOT ONLINE',
			'ip_exist_error'   : 'IP ALREADY EXISTS',
			'user_tab_error'   : 'THE CERTIFICATION PROGRAM WAS NOT STARTED',
			'username_error'   : 'WRONG USERNAME',
			'user_group_error' : 'ACCOUNT INFOMATION INCORRECT',
			'password_error'   : 'WRONG PASSWORD',
			'status_error'     : 'ACCOUNT OVERDUE, PLEASE RECHARGE',
			'available_error'  : 'ACCOUNT HAS BEEN SUSPENDED',
			'delete_error'     : 'ACCOUNT HAS BEEN DELETED',
			'usernum_error'    : 'USERS NUMBER LIMITED',
			'online_num_error' : 'USERS NUMBER LIMITED',
			'mode_error'       : 'DISABLE WEB REGISTRY',
			'time_policy_error': 'CURRENT TIME IS NOT ALLOWED TO CONNECT',
			'flux_error'       : 'FLUX OVER',
			'ip_error'         : 'IP NOT VALID',
			'mac_error'        : 'MAC NOT VALID',
			'sync_error'       : 'YOUR INFOMATION HAS BEEN MODIFIED, PLEASE TRY AGAIN AFTER 2 MINUTES',
			'ip_alloc'         : 'THE IP HAS BEEN ASSIGNED TO OTHER USER'
			}

version  = '1.2'
sleep_time = 8

#########################################################
#				Main Login/Logout Modules				#
#########################################################

def trans_content(response):
	content = response.read().decode()
	ret = ''
	for ch in content:
		if ch.isalpha() or ch == '_':
			ret += ch
	return ret

def tunet_login(username, password):
	hashcd_md5 = hashlib.md5()
	hashcd_md5.update(password.encode())
	tr_password = hashcd_md5.hexdigest()
	login_data = 'username=' + username + '&password=' + tr_password + '&drop=0&type=1&n=100'
	login_data = login_data.encode()
	request_url = urllib.request.Request(login_url, login_data)
	response_url = urllib.request.urlopen(request_url)
	ret = trans_content(response_url)
	print (ret_type.get(ret, 'CONNECTED'))
	return ret

def tunet_logout():
	response_url = urllib.request.urlopen(logout_url)
	ret = trans_content(response_url)
	print (ret_type.get(ret, 'CONNECTED'))
	return ret

def tunet_check():
	check_data = 'action=check_online'
	check_data = check_data.encode()
	request_url = urllib.request.Request(check_url, check_data)
	response_url = urllib.request.urlopen(request_url)
	ret = trans_content(response_url)
	if ret == '':
		print ('NOT ONLINE')
	else:
		print (ret_type.get(ret, 'CONNECTED'))
	return ret

#########################################################
#					Help&Version Modules				#
#########################################################

def tunet_help():
	print ('-h, --help   : show all options of Tsinghua University Internet Connector')
	print ('-v, --version: show version of Tsinghua University Internet Connector')
	print ('-u           : input your username after \'-u\'')
	print ('-p           : input your password after \'-p\'')
	print ('-a           : enter username and password later, you can login other campus network account')
	print ('-i, --login  : login operation')
	print ('-o, --logout : logout operation')
	print ('-c, --check  : check the internet')
	print ('-q, --query  : query basic infomation, online state and flux usage details')

def tunet_version():
	print ('Tsinghua University Internet Connector ', version)

def tunet_others():
	print ('UNKNOWN OPTIONS')
	print ('WHICH OPTION DO YOU WANT?')
	tunet_help()
	print ('IF ANY ERROR, PLEASE CONTACT im@huhaoyu.com.')

#########################################################
#						Main Part						#
#########################################################

def tunet_connect(username, password):
	ret = 'ip_exist_error'
	for count in range(5):
		print ('%s attempts to connect...' % times_cnt.get(count + 1))
		if ret != tunet_login(username, password):
			break
		if count == 4:
			print ('please try to reconnect after 1 minute')
			break
		print ('try to reconnect after %s seconds' %sleep_time)
		time.sleep(sleep_time)
		print ()

pytunet.py：

#########################################################
#						Welcome							#
#	Tsinghua University Internet Connector in Python	#
#					 Version: v1.2 						#
#					Date: 2015/04/05					#
#			By: Haoyu hu	Email: im@huhaoyu.com		#
#			Address: Tsinghua University				#
#########################################################

import pytunet_connect
import pytunet_query
import sys, getopt, getpass, os, re

def pytunet():
	relative_path = 'USERNAME_PASSWORD.txt'
	file_handler = open(pytunet_query.create_path(relative_path), 'r')
	lines = file_handler.readlines()
	if len(lines) != 2:
		print('%s IS DAMAGED, PLEASE CHECK THE CONTENT FORMAT IN THE FILE:' %relative_path)
		print('CONTENT FORMAT AS FOLLOWS:')
		print('username=huhy14\npassword=123456')
		exit(1)
	username = re.sub('\s', '', lines[0])
	password = re.sub('\s', '', lines[1])
	_, username = username.split('=')
	_, password = password.split('=')
	file_handler.close()

	try:
		options, args = getopt.getopt(sys.argv[1:], 'achiop:qu:v', ['help', 'login', 'logout', 'check', 'version', 'query'])
	except getopt.GetoptError:
		pytunet_connect.tunet_others()
		sys.exit(1)

	want_login = False
	want_query = False
	flag = False

	name, value = None, None

	for name, value in options:
		if name in ('-h', '--help'):
			pytunet_connect.tunet_help()
			sys.exit(0)
		elif name in ('-v', '--version'):
			pytunet_connect.tunet_version()
			sys.exit(0)
		elif name == '-a':
			flag = True
		elif name == '-u':
			username = value
		elif name == '-p':
			password = value
		elif name in ('-i', '--login'):
			want_login = True
		elif name in ('-o', '--logout'):
			pytunet_connect.tunet_logout()
			sys.exit(0)
		elif name in ('-c', '--check'):
			pytunet_connect.tunet_check()
			sys.exit(0)
		elif name in ('-q', '--query'):
			want_query = True

	if flag:
		username = input('username: ')
		password = getpass.getpass('password: ')
		print (password)

	if want_query:
		pytunet_query.tunet_query(username, password)

	if want_login or not want_query and not want_login:
		pytunet_connect.tunet_connect(username, password)

	# if not want_query and not want_login:
	# 	print ('WARNING: YOU JUST DIDN\'T DO ANYTHING! IF YOU WANT TO CONNECT TO THE CAMPUS NETWORK, THE COMMAND MUST INCLUDE -i OR --login')
	# 	print()
	# 	pytunet_connect.tunet_help()

if __name__ == '__main__':
	pytunet()

实际使用：

如分析或代码有误，请批评指正，谢谢！