#!/usr/bin/env lua
local posix = require("Posix")
local ubus = require ("ubus")
local fs     = require "nixio.fs"

local cfg = {
	['postfile'] = "/tmp/postfile.dat",
	['postfilesize'] = 1024,  -- kbyte
	['posturl'] = "http://netsp.master.qq.com/cgi-bin/netspeed",
	['geturl'] = "http://dlied6.qq.com/invc/qqdoctor/other/test32mb.dat",
	['nr'] = 320,        --Number of requests to perform
	['nc'] = 8,     --Number of multiple requests to make at a time
	['timelimit'] = 9,
	['timestep'] = 1,
	['interval'] = 1,
	['weight'] = 0.95, -- smooth net tarffic burst at first second
	['qos_weight'] = 0.98,
	['burstrate'] = 2,
	['ab'] = "/usr/bin/ab",
	['dd'] = "/bin/dd",
	['debug'] = 0,
	['xmlfile'] = "/usr/share/speedtest.xml",
        ['tmp_speedtest_xml'] = "/tmp/speedtest_urls.xml",
}

VERSION="__UNDEFINED__"

if VERSION == "LESSMEM" then
    cfg.postfilesize = 640
    cfg.nc = 6
end

local filename = ""
filexml = io.open(cfg.tmp_speedtest_xml)
if filexml then
    filexml:close()
    filename = cfg.tmp_speedtest_xml
else
    filename = cfg.xmlfile
end

local pp = io.open(filename)
local line = pp:read("*line")
local size = 0
local resources = {}
local u = ""
local pids = {}

function die(err)
	posix.openlog(arg[0], "cp", posix.LOG_LOCAL7)
	posix.syslog(posix.LOG_ERR, err)
	posix.closelog()
	os.exit(1)
end

function logger(loglevel,msg)
    posix.openlog("speedtest","np",LOG_USER)
    posix.syslog(loglevel,msg)
    posix.closelog()
end

print(string.format("upload using %s...", filename))
print(string.format("upload using file size: %d  nc: %d", cfg.postfilesize, cfg.nc))
logger(3, string.format("upload using %s...", filename))
logger(3, string.format("upload using file size: %d  nc: %d", cfg.postfilesize, cfg.nc))

function mrandom(min,max,num)
	local reverse = {}
	local t = {}
	local ret = {}
	local i = min
	local index
	while i <= max do
		table.insert(t, i)
		i = i + 1
	end
	i = num
	math.randomseed(os.time())

	while i > 0 do
		index = math.random(table.getn(t))
		table.insert(ret,t[index])
		if index == table.getn(t) then
			table.remove(t)
		else
			local top = table.remove(t)
			t[index] = top
		end
		i = i - 1
	end
	return ret
end


function execa(cmd)
	local p = io.popen(cmd)
	local line = p:read("*l")
	while(line) do
		print(line)
		line = p:read("*l")
	end
	p:close()
end

function wget_work(url)
	local _url = url
	pid = posix.fork()
	if pid < 0 then
		print("fork error")
		return -1
	elseif pid > 0 then
		--print(string.format("child pid %d\n", pid))
	else
		os.execute('for i in $(seq '.. math.floor(cfg.nr/cfg.nc) ..'); do wget '.. url  ..
		" -q -O /dev/null; done")
	end
	return pid
end

function wan_device()
	local conn = ubus.connect()
	if not conn then
		elog("Failed to connect to ubusd")
	end
	local status = conn:call("network.interface.wan", "status",{})
	conn:close()
	if not status then
		return nil
	else
		return (status.l3_device and status.l3_device) or status.device
	end
end

function get_pstree(root, pids)
	local pp = io.popen("pgrep -P "..root)
	pid = pp:read("*line")
	while pid do
		data = table.insert(pids, pid)
		get_pstree(pid, pids)
		pid = pp:read("*line")
	end
	pp:close()
end

function execl2(command)
    local pp   = io.popen(command)
    local line = ""
    local data = {}

    while true do
        line = pp:read()
        if line == nil then
            break
        end
        data[#data+1] = line
    end
    pp:close()
    return data
end

function get_pstree2(root, pids)
    local cmd = "pstree -p "..root
    local res = execl2(cmd)

    print("type root is: "..type(root))
    if res and next(res)~=nil then
        for k,v in ipairs(res) do
            local j = 0
            while true do
                _,j,pid  = string.find(v, '%(([0-9]+)%)', j+1)
                if pid then
                    if tonumber(pid) ~= root then
                        table.insert(pids, pid)
                    end
                else
                    break
                end
            end
        end
    end
end

function done(signo)
	io.output("/dev/null")
	local fd = io.open("/dev/null", "rw")
	if not posix.dup(fd, io.stdout) then
		die("error dup2-ing")
	end
	if not posix.dup(fd, io.stderr) then
		die("error dup2-ing")
	end
	get_pstree2(posix.getpid("pid"), pids)
	for k, pid in ipairs(pids) do
		print("kill pid:" .. pid)
		posix.kill(pid, posix.SIGINT)
	end
	os.exit(0)
end

function read_line(filename)
	local fd = io.open(filename)
	local line = fd:read("*line")
	fd:close()
	return line
end


function get_uptime()
	local _, _, uptime, idle = string.find(read_line("/proc/uptime"),'^([0-9.]+)%s+([0-9.]+)$')
	return tonumber(uptime)
end


function get_rt(ifname)
	local line
	local face, r_bytes, r_packets, r_errs, r_drop, r_fifo, r_frame, r_compressed, r_multicast
	local t_bytes, t_packets, t_errs, t_drop, t_fifo, t_colls, t_carrier, t_compressed
	local _nic = {}
	if fs.access("/proc/net/dev") then
		for line in io.lines("/proc/net/dev") do
			_, _, face, r_bytes, r_packets, r_errs, r_drop, r_fifo, r_frame, r_compressed, r_multicast,
			t_bytes, t_packets, t_errs, t_drop, t_fifo, t_colls, t_carrier, t_compressed = string.find(line,
			'%s*(%S+):%s+(%d+)%s+(%d+)%s+(%d+)%s+(%d+)%s+(%d+)%s+(%d+)%s+(%d+)%s+(%d+)%s+(%d+)%s+(%d+)%s+(%d+)%s+(%d+)%s+(%d+)%s+(%d+)%s+(%d+)%s+(%d+)%s*')
			if (face == ifname) then
				return {
					r_bytes   = tonumber(r_bytes),
					r_packets = tonumber(r_packets),
					t_bytes   = tonumber(t_bytes),
					t_packets = tonumber(t_packets),
					uptime    = get_uptime()
				}
			end
		end
	end
end

function print_rt(rt, rt_last)
	local delta_time = rt.uptime - rt_last.uptime
	print(string.format("Time(ms):%.0f\tr_x:%.2f t_x:%.2f",
	delta_time*1000,
	(rt.r_bytes - rt_last.r_bytes) / delta_time / 1024,
	(rt.t_bytes - rt_last.t_bytes) / delta_time / 1024))
end

function ab_work(url)
	pid = posix.fork()
	if pid < 0 then
		print("fork error")
		return -1
	elseif pid > 0 then
		--print(string.format("child pid %d\n", pid))
	else
		local cmd = (string.format("%s -N  -p %s '%s' > /dev/null", cfg.ab, cfg.postfile, url))	
		--print("ab work cmd: "..cmd)
		os.execute('for i in $(seq '.. math.floor(cfg.nr/cfg.nc) ..'); do ' .. cmd .. "; done")	

	end
	return pid
end

function ab_work_lowmem()
	local cmd = string.format("%s -N '%s' '%s' '%s' '%s' '%s' '%s' -L 6 -R 100 -t 15 -M 15 > /dev/null", cfg.ab, 'http://www.taobao.com/', 
				'http://www.so.com/', 'http://www.qq.com/', 'http://www.sohu.com/', 'http://www.tudou.com/', 'http://www.kankan.com/')
	pid = posix.fork()
	if pid < 0 then
		print("fork error")
		return -1
	elseif pid > 0 then
		--print(string.format("child pid %d\n", pid))
	else
		os.execute(cmd)

	end
	return pid
end

function dump_dot_data(dotdata)
	for i=1,#dotdata do
	print(string.format("[%.2d] %.2f %.2f %.2f", i, dotdata[i].r_bytes,dotdata[i].t_bytes, dotdata[i].uptime))
	end
end
-----------------------------------------------------------------------



posix.signal(posix.SIGTERM, done);

local wan_ifname = wan_device()
if not wan_ifname then
        print("got invalid wan device")
        print(string.format("avg tx:%.2f", 0))
	logger(3, "stat_points_privacy network_speedtest=10|upload|no_wan_dev")
	os.remove(cfg.postfile)
	done()
end

while line do
	local _, _, url = string.find(line,'<item uploadurl="(.*)"/>')
	if url then
		table.insert(resources, url)
	end
	line = pp:read("*line")
end
pp:close()

--创建post文件
if VERSION ~= "LESSMEM" then
	os.execute(string.format("%s if=/dev/zero of=%s bs=1k count=%d >/dev/null 2>&1",
    		cfg.dd, cfg.postfile, cfg.postfilesize))

	if (posix.stat(cfg.postfile) == nil) then
    		print("create postfile error")
    		logger(3, "stat_points_privacy network_speedtest=10|upload|postfile_err")
    		os.exit(1)
	end
end

local urls = mrandom(1, table.getn(resources), cfg.nc)

if VERSION == "LESSMEM" then
	local pid = ab_work_lowmem()
	if(pid == 0) then
            	os.exit(0)
        elseif(pid == -1) then
               	done()
        end
else
	for k, v in ipairs(urls) do
        	local pid = ab_work(resources[v])
        	if(pid == 0) then
                	os.exit(0)
        	elseif(pid == -1) then
                	os.remove(cfg.postfile)
                	done()
        	end
	end
end

local rt = get_rt(wan_ifname)
if not rt then
        print("got invalid data")
        print(string.format("avg tx:%.2f", 0))
	logger(3, "stat_points_privacy network_speedtest=10|upload|no_wan_info")
	os.remove(cfg.postfile)
	done()
end


local rt_last
local dot_datas = {}
local dot = 1
dot_datas[dot] = rt

while dot <= cfg.timelimit  do
	posix.sleep(1)

	dot = dot + 1
	rt_last = rt
	rt = get_rt(wan_ifname)
	if not rt then
        	print("got invalid data")
        	print(string.format("avg tx:%.2f", 0))
		logger(3, "stat_points_privacy network_speedtest=10|upload|no_wan_info")
        	os.remove(cfg.postfile)
        	done()
	end

	dot_datas[dot] = rt
	print_rt(rt, rt_last)
end


--计算平均值
function tx_avg_speed(dotdata)
	local sum,j = 0,1
	local dlt_time

	if #dotdata < 2 then 
		return 0
	end

	while j + 1 <= #dotdata do
		dlt_time = dotdata[j+1].uptime - dotdata[j].uptime
		if (dotdata[j+1].t_bytes <= dotdata[j].t_bytes) then 
			return 0
		end
		sum = sum + (dotdata[j+1].t_bytes - dotdata[j].t_bytes) / dlt_time / 1024
		j = j + 1
	end

	return sum/(#dotdata-1)
end

--每间隔cfg.interval秒的tx数据计算速度，取其中的最大值
function tx_max_avg_speed(dotdata, net_burst)
	local tx_speed = 0
	local i = 2  --i从2开始，避开第一秒的burst数据

	while (i + cfg.interval) <= (cfg.timelimit + 1) do
		local delta_time = dotdata[i+cfg.interval].uptime - dotdata[i].uptime
		if (dotdata[i+cfg.interval].t_bytes <= dotdata[i].t_bytes) then
			return 0
		end

		local tmp = (dotdata[i+cfg.interval].t_bytes - dotdata[i].t_bytes)/ delta_time / 1024
		--print(string.format("sec[%d]%.2f  sec[%d]%.2f delta_time %.2f, speed is %.2f",
		--				i+cfg.interval, dotdata[i+cfg.interval].t_bytes, i, 
		--				dotdata[i].t_bytes, delta_time, tmp))

		if (i == 1) and net_burst then
			--print(string.format("sec[%d] weight %.2f burstrate %.2f : smooth from %.2f to %.2f", 
			--		i, cfg.weight, cfg.burstrate, tmp, tmp*cfg.weight))
			tmp = tmp * cfg.weight
		end

		if tmp > tx_speed then
			tx_speed = tmp
		end
		i = i + 1
	end

	return tx_speed
end

-- 去掉一个最大值，去掉一个最小值, 求平均
function  avg_remove_max_min(tdata)
	local sum = 0
	table.sort(tdata)

	--print(string.format("table remove max %.2f  min %.2f",
	--					tdata[#tdata], tdata[1]))

	table.remove(tdata,#tdata)
	table.remove(tdata,1)

	for key, value in pairs(tdata) do
		sum = sum + value
		--print(string.format("%.2f", value))
	end

	return sum/#tdata
end

--每间隔cfg.interval秒的tx数据计算速度，去掉一个最大值，去掉一个最小值, 求平均
function tx_avg_interval_speed(dotdata)
	local sum = {}
	local tx_speed = 0
	local i = 1

	while (i + cfg.interval) <= (cfg.timelimit + 1) do
		if  (dotdata[i+cfg.interval].t_bytes <= dotdata[i].t_bytes) then
			print("t_x data "..(i+cfg.interval).." <= "..i)
			logger(3, "stat_points_privacy network_speedtest=10|upload|data_invalid|"..i+cfg.interval.."|"..dotdata[i+cfg.interval].r_bytes.."|"..i.."|"..dotdata[i].r_bytes)
			return 0
		end

		local delta_time = dotdata[i+cfg.interval].uptime - dotdata[i].uptime
		local tmp = (dotdata[i+cfg.interval].t_bytes - dotdata[i].t_bytes)/ delta_time / 1024
		--print(string.format("sec[%d]%.2f  sec[%d]%.2f delta_time %.2f, speed is %.2f",
		--			i+cfg.interval, dotdata[i+cfg.interval].t_bytes, i, 
		--			dotdata[i].t_bytes, delta_time, tmp))

		table.insert(sum, tmp)
		i = i + 1
	end

	tx_speed = avg_remove_max_min(sum)
	if tx_speed*8 < 56 then
    		logger(3, "stat_points_privacy network_speedtest=10|upload|data_small_value|"..tx_speed*8)
	end
	return tx_speed
end

--判断第一秒是否有net burst
function if_net_burst(dotdata)
	if ((dotdata[2].t_bytes < dotdata[1].t_bytes) or (dotdata[3].t_bytes < dotdata[2].t_bytes)) then 
		return false
	end

	local delta_time = dotdata[2].uptime - dotdata[1].uptime
	local first_speed = (dotdata[2].t_bytes - dotdata[1].t_bytes) / delta_time / 1024

	delta_time = dotdata[3].uptime - dotdata[2].uptime
	local second_speed = (dotdata[3].t_bytes - dotdata[2].t_bytes) / delta_time / 1024

	if first_speed > (cfg.burstrate * second_speed) then
		return true
	end
		return false
end

--dump_dot_data(dot_datas)
local res_speed =  tx_avg_speed(dot_datas)
print(string.format("avg t_x:%.2f  %.2fMB  %.2fMbit", res_speed*8, res_speed/1024, res_speed*8/1024))

res_speed = tx_avg_interval_speed(dot_datas)

print(string.format("qos weight %.3f:  from %.2f to %.2f",
					cfg.qos_weight, res_speed, res_speed*cfg.qos_weight))

res_speed = res_speed * cfg.qos_weight
print(string.format("avg tx:%.2f  %.2fMB  %.2fMbit", res_speed*8, res_speed/1024, res_speed*8/1024))


os.remove(cfg.postfile)
done()