客户端分IP上行流量统计

ip_flow.py

import os, json, subprocess, time

server_port = 8388
chain_name = 'ip_flow'

last_connected_clients = []
client_info = []

def get_client_list_from_file():
	if not os.path.isfile('clients.json'):
		return []
	f = open('clients.json', 'r')
	client_list = json.loads(f.read())
	f.close()
	return client_list

def get_connected_client_ips():
	command = "netstat -antup | grep :" + str(server_port) + " | awk '{print $5}' | cut -d ':' -f 1 | sort | uniq"
	output = subprocess.Popen(['bash', '-c', command], stdout=subprocess.PIPE).communicate()[0]
	return [ip.strip() for ip in output.split('\n') if ip.strip() != '' and ip.strip() != '0.0.0.0']

def add_ip_log(client_ip):
	os.system('iptables -t filter -A ' + chain_name + ' -d ' + client_ip)

def del_ip_log(client_ip):
	os.system('iptables -t filter -D ' + chain_name + ' -d ' + client_ip)

def get_client_flow(client_ip):
	command = 'iptables -t filter -L ' + chain_name + ' -n -v -x | grep ' + client_ip + " | awk '{print $2}'"
	output = subprocess.Popen(['bash', '-c', command], stdout=subprocess.PIPE).communicate()[0]
	return int(output if output.strip() != '' else 0)

def get_ip_geo_location(ip):
	command = "geoiplookup " + ip + " | grep 'City Edition' | cut -d ':' -f 2 | cut -d ',' -f 1-4"
	output = subprocess.Popen(['bash', '-c', command], stdout=subprocess.PIPE).communicate()[0]
	return output.strip().replace('\n', '')

def log_client_flow(client_ip, flow):
	client_in_list = False
	for client in client_info:
		if client['ip'] == client_ip:
			client_in_list = True
			selected_client = client
			break
	if client_in_list == False:
		selected_client = {
			"ip": client_ip,
			"geo_location": get_ip_geo_location(client_ip),
			"flow": 0,
		}
		client_info.append(selected_client)
	if client_ip not in last_connected_clients:
		selected_client['__base_flow'] = selected_client['flow']
	selected_client['rate'] = str((selected_client['__base_flow'] + flow - selected_client['flow']) * 1.0 / 1024) + 'kb/s'
	selected_client['flow'] = selected_client['__base_flow'] + flow

def once():
	global last_connected_clients
	current_connected_clients = get_connected_client_ips()
	for current_ip in current_connected_clients:
		if current_ip not in last_connected_clients:
			add_ip_log(current_ip)
			print('New connection from: ' + current_ip)
	for last_ip in last_connected_clients:
		if last_ip not in current_connected_clients:
			del_ip_log(last_ip)
			print('End connection from: ' + last_ip)

	for client_ip in current_connected_clients:
		flow = 	get_client_flow(client_ip)
		log_client_flow(client_ip, flow)

	f = open('clients.json', 'w+')
	f.write(json.dumps(client_info, indent=4, separators=(',', ': ')))
	f.close()

	last_connected_clients = current_connected_clients

def main():
	global client_info
	os.system('iptables -t filter -F ' + chain_name)
	client_info = get_client_list_from_file()
	while True:
		once()
		time.sleep(1)

if __name__ == '__main__':
	main()