#! /usr/local/bin/python3
import re
import subprocess
import sys
'''
data_forward=sys.argv[1]
ip_addr=sys.argv[2]
'''
def tc_ip2id(data_forward):
    command='tc'
    arg1='-s'
    arg2='-d'
    arg3='filter'
    arg4='show'
    arg5='dev'
    if data_forward.lower() == 'down' :
        arg6='enp11s0f0'
    else:
        arg6='enp11s0f1'
    result = subprocess.run([command,arg1,arg2,arg3,arg4,arg5,arg6], stdout=subprocess.PIPE, universal_newlines=True )

    result_list = result.stdout.split(sep='filter')

#define regular express about ip and id
    ip_re = re.compile(r'([a-fA-F0-9]{2})([a-fA-F0-9]{2})([a-fA-F0-9]{2})([a-fA-F0-9]{2})(?=/ffffffff)')
    id_re = re.compile(r'\d:\d{4}')

#生成ip和id的对应字典
    global ip_id_dict
    ip_id_dict={}
    for i in result_list :
    #当前十六进制ip
        if ip_re.search(i) != None :
            cur_ip_tuple = ip_re.search(i).groups()
    #转存成十进制Ip
            ip_dec_str=''
            for j in cur_ip_tuple :
                ip_dec_str += str(int(j,16))+'.'
            ip_dec_str=ip_dec_str[:-1]
        #向字典中添加ip和id对应关系
            ip_id_dict[ip_dec_str]=id_re.search(i).group(0)
        #清空之前的值
            ip_dec_str=''
        else:
            continue

#检查单个IP的流量情况
def tc_check_ip(data_forward, ip_addr):
    ip_dec_str=''
#转换要查找到IP到字符串
    if ip_addr != None :
        for i in re.search(r'(\d?\d?\d).(\d?\d?\d).(\d?\d?\d).(\d?\d?\d)',ip_addr).groups() :
            ip_dec_str+=i+'.'
        ip_dec_str=ip_dec_str[:-1]
    else:
        print("No ip address")
        return
    command='tc'
    arg1='-s'
    arg2='-d'
    arg3='class'
    arg4='show'
    arg5='dev'
    if data_forward.lower() == 'down' :
        arg6='enp11s0f0'
    else:
        arg6='enp11s0f1'
    arg7='classid'
    result = subprocess.run([command,arg1,arg2,arg3,arg4,arg5,arg6,arg7,ip_id_dict[ip_dec_str]], stdout=subprocess.PIPE, universal_newlines=True)
    return result.stdout
    
#检查网段内每个IP的网络使用情况排名
def tc_check_network(data_forward, network_addr):
    ip_dec_str=''
#转换要查找到IP到字符串
    if network_addr != None :
        for i in re.search(r'(\d?\d?\d).(\d?\d?\d).(\d?\d?\d).(?=\d?\d?\d)',network_addr).groups() :
            ip_dec_str+=i+'.'
        ip_dec_str=ip_dec_str[:-1]
    else:
        return
    command='tc'
    arg1='-s'
    arg2='-d'
    arg3='class'
    arg4='show'
    arg5='dev'
    if data_forward.lower() == 'down' :
        arg6='enp11s0f0'
    else:
        arg6='enp11s0f1'
    result = subprocess.run([command,arg1,arg2,arg3,arg4,arg5,arg6], stdout=subprocess.PIPE, universal_newlines=True)
    result_list = result.stdout.split(sep='class')
    
    #查询的网段IP存于此
    network_ip=[]
    for i in result_list:
        for j in range(1,254):
            if re.search(r'(?<=htb )(\d:\d{4})',i) == None :
                break
            if  re.search(r'(?<=htb )(\d:\d{4})',i).group() == ip_id_dict[ip_dec_str+'.'+str(j)] :
                #将IP和tokens加入到列表中
                try:
                    network_ip.append((ip_dec_str+'.'+str(j), int(re.search(r'(?<=tokens: )(-?\d?\d?\d?\d?\d?\d?)',i).group())))
                    break
                except ValueError:
                    print(i)
                    break
            else:
                continue
    #对network_ip进行排序
    network_ip.sort(key=comp_func)
    return network_ip
    
    
def comp_func(sth):
    return sth[1]
'''
用法 command down|up IP_ADDRESS|NETWORK_ADDRESS
'''
if __name__=='__main__' :
    arg_forward = sys.argv[1]
    arg_ip = sys.argv[2]
    
    tc_ip2id(arg_forward)
    if re.search(r'(?<=\.)(\d?\d?\d)$',arg_ip) != None :
        if int(re.search(r'(?<=\.)(\d?\d?\d)$',arg_ip).group()) == 0 :
            print(tc_check_network(arg_forward,arg_ip))
        else :
            print(tc_check_ip(arg_forward, arg_ip))