#!/usr/bin/env luajit


local args = {...}


verbosity = 0
keep = false
options =
{
    referenceSite =
	{
	    typ = "string",
	    help = "",
	    value = "pkgmaster.devuan.org",
	},
    tests =
	{
	    typ = "table",
	    help = "",
	    value =
	    {
	        "IPv4",
	        "IPv6",
--	        "ftp",
	        "http",
	        "https",
--	        "rsync",
	        "DNS-RR",
	        "Protocol",
--	        "URL-Sanity",
--	        "Integrity",
--	        "Updated",
	    },
	},
}

local defaultURL = {scheme = "http"}
local downloadLock = "flock -n results/wget-"
local download = "wget --timeout=300 -np -N -r -P results "	-- Note wget has a default read timeout of 900 seconds (15 minutes).
local releases = {"jessie", "ascii", "beowulf", "ceres"}
local releaseFiles =
{
    -- Release file.
    "/Release",
    "/InRelease",
    "/main/binary-all/Packages.gz",
    -- Contents files.
    "/main/Contents-all.gz",
    "/main/Contents-amd64.gz",
    "/main/Contents-arm64.gz",
    "-security/main/Contents-all.gz",
    "-security/main/Contents-amd64.gz",
    "-security/main/Contents-arm64.gz",
}
local notExist =
{
    "ceres-security"	-- This will never exist, it's our code name for the testing suite.
}
local referenceDebs =
{
    -- Devuan package.  NOTE this one should not get redirected, but that's more a warning than an error.
    "merged/pool/DEVUAN/main/d/desktop-base/desktop-base_2.0.3_all.deb",
    -- Debian package.
    "merged/pool/DEBIAN/main/d/dash/dash_0.5.8-2.4_amd64.deb",
    -- Debian security package.  NOTE this one should always be redirected?
    "merged/pool/DEBIAN-SECURITY/updates/main/a/apt/apt-transport-https_1.4.9_amd64.deb",
}
local arg = {}
local sendArgs = ""
local logFile


local socket = require 'socket'
local ftp = require 'socket.ftp'
local http = require 'socket.http'
local url = require 'socket.url'


-- Use this to print a table.
printTable = function (table, space, name)
    if "" == space then print(space .. name .. " =") else print(space .. "[" .. name .. "] =") end
    print(space .. "{")
    printTableSub(table, space .. "  ")
    if "" == space then print(space .. "}") else print(space .. "},") end
end
printTableSub = function (table, space)
    for k, v in pairs(table) do 
        if type(k) == "string" then k = '"' .. k .. '"' end
	if type(v) == "table" then
	    printTable(v, space, k)
	elseif type(v) == "string" then
	    print(space .. "[" .. k .. "] = '" .. v .. "';")
	elseif type(v) == "function" then
	    print(space .. "[" .. k .. "] = function ();")
	elseif type(v) == "userdata" then
	    print(space .. "userdata " .. "[" .. k .. "];")
	elseif type(v) == "boolean" then
	    if (v) then
		print(space .. "boolean " .. "[" .. k .. "] = true;")
	    else
		print(space .. "boolean " .. "[" .. k .. "] = false;")
	    end
	else
	    print(space .. "[" .. k .. "] = " .. v .. ";")
	end
    end
end

local log = function(v, t, s)
    if v <= verbosity then
	if 3 <= verbosity then t = os.date() .. " " .. t end
	print(t .. ": " .. s)
    end
    if nil ~= logFile then
	logFile:write(os.date() .. " " .. t .. ": " .. s .. "\n")
	logFile:flush()
    end
end
local D = function(s) log(3,  "DEBUG   ", s) end
local I = function(s) log(2,  "INFO    ", s) end
local W = function(s) log(1,  "WARNING ", s) end
local E = function(s) log(0,  "ERROR   ", s) end
local C = function(s) log(-1, "CRITICAL", s) end

local testing = function(t)
    for i, v in pairs(options.tests.value) do
	if t == v then return true end
    end
    return false
end

local checkExes = function (exe)
    local count = io.popen('ps x | grep "' .. exe .. '" | grep -v " grep " | wc -l'):read("*l")
    D(count .. " " .. exe .. " commands still running.")
    return tonumber(count)
end

local repoExists = function (r)
    r = r:match("([%a-]*)")
    if nil == r then return false end
    for k, v in pairs(notExist) do
	if v == r then return false end
    end
    return true
end

local IP = {}
local gatherIPs = function (host)
    if nil == IP[host] then
        local IPs
	local dig = io.popen('dig +keepopen +noall +nottlid +answer ' .. host .. ' A ' .. host .. ' AAAA ' .. host .. ' CNAME ' .. host .. ' SRV | sort -r | uniq')
	repeat
	    IPs = dig:read("*l")
	    if nil ~= IPs then
		for k, t, v in IPs:gmatch("([%w_%-%.]*)%.%s*IN%s*(%a*)%s*(.*)") do
		    if "." == v:sub(-1, -1) then v = v:sub(1, -2) end
		    if nil == IP[k] then IP[k] = {} end
		    IP[k][v] = t
		end
	    end
	until nil == IPs
    end
end

-- Returns FTP directory listing
local nlst = function (u)
    local t = {}
    local p = url.parse(u)
    p.command = "nlst"
    p.sink = ltn12.sink.table(t)
    local r, e = ftp.get(p)
    return r and table.concat(t), e
end

checkURL = function (host, URL, r, retry)
    if nil == r then r = 0 end
    if nil == retry then retry = 0 end
    local check = "Checking file"
    if 0 < r then
	check = "Redirecting to"
	--[[ The hard part here is that we end up throwing ALL of the test files at the redirected location.
	    Not good for deb.debian.org, which we should only be throwing .debs at.
	    What we should do is loop through the DNS entries, and only test the specific protocol & file being tested here.

	    This is what I came up with for checking if we are already testing a specific URL.
	    Still duplicates a tiny bit, but much less than the previous find based method.
	]]
	local pl = url.parse(URL, defaultURL)
	local file = host .. "://" .. pl.path
        local f = io.popen(string.format('if [ ! -f results/%s.check ] ; then touch results/%s.check; echo -n "check"; fi', file:gsub("/", "_"), file:gsub("/", "_") )):read("*a")
	if (nil == f) or ("check" == f) then
	    I("  Now checking redirected host " .. file)
	    checkRedirects(host, host, nil, "redir", pl.path)
	else
	    D("  Already checking " .. file)
	end

    end
    if 0 < retry then
-- TODO - should have a random sleep here before retrying.
	check = "Retry " .. retry .. " " .. check
    end
    if 10 < r then
	E("too many redirects!  " .. check .. " " .. host .. " -> " .. URL)
	return
    end
    if 10 < retry then
	E("too many retries!  " .. check .. " " .. host .. " -> " .. URL)
	return
    end
    local PU = url.parse(URL, defaultURL)
    D(PU.scheme .. " ://  " .. check .. " " .. host .. " -> " .. URL)
    if not testing(PU.scheme) then D("not testing " .. PU.scheme .. " " .. host .. " -> " .. URL);  return end
    local hd = {Host = host}
    local p, c, h, s = http.request{method = "HEAD", redirect = false, url = URL, headers = hd}
    if nil == s then s = "" end
    if nil == p then
	E("  " .. c .. " " .. s .. "!   " .. check .. " " .. host .. " -> " .. URL)
	-- So far the only errors are "timeout", "Network is unreachable", and "closed", and I suspect "closed" is due to saturating my bandwidth.
	-- Might be worthwhile retrying those, some number of times.
	if ("closed" == c) or ("Network is unreachable" == c) or ("timeout" == c) then checkURL(host, URL, r, retry + 1) end
    else
	if ("4" == tostring(c):sub(1, 1)) or ("5" == tostring(c):sub(1, 1)) then
	    E("  " .. c .. " " .. s .. ".   " .. check .. " " .. host .. " -> " .. URL)
	else
	    I("  " .. c .. " " .. s .. ".   " .. check .. " " .. host .. " -> " .. URL)
	end
	l = h.location
	if nil ~= l then
	    local pu = url.parse(l, defaultURL)
	    if l == URL then
		E("    redirect loop!  " .. check .. " " .. host .. " -> " .. URL)
	    else
		if nil == pu.host then
		    W("    no location host!  " .. check .. " " .. host .. " -> " .. URL .. " -> " .. l)
		    checkURL(host, PU.scheme .. "://" .. PU.host .. l, r + 1)
		else
		    if testing("Protocol") and pu.scheme ~= PU.scheme then
			W("    protocol changed during redirect!  " .. check .. " " .. host .. " -> " .. URL .. " -> " .. l)
		    end
		    checkURL(pu.host, l, r + 1, retry)
		end
	    end
	end
    end
end

local checkPaths = function (host, ip, path, file)
    if nil ~= file then
	if "redirect" == ip then ip = host end
	I("  Checking IP for file " .. host .. " -> " .. ip .. " " .. path .. " " .. file)
	if testing("http")  then checkURL(host, "http://"  .. ip .. path .. "/" .. file) end
	if testing("https") then checkURL(host, "https://" .. ip .. path .. "/" .. file) end
    else
	I("  Checking IP " .. host .. " -> " .. ip .. " " .. path)
	for i, s in pairs(referenceDebs) do
	    if testing("http")  then checkURL(host, "http://"  .. ip .. path .. "/" .. s) end
	    if testing("https") then checkURL(host, "https://" .. ip .. path .. "/" .. s) end
	end

	for i, s in pairs(releases) do
	    for j, k in pairs(releaseFiles) do
		if repoExists(s .. k) then
		    if testing("http")  then checkURL(host, "http://"  .. ip .. path .. "/merged/dists/" .. s .. k) end
		    if testing("https") then checkURL(host, "https://" .. ip .. path .. "/merged/dists/" .. s .. k) end
		end
	    end
	end
    end
end

local execute = function (s)
    D("  executing " .. s)
    os.execute(s)
end

checkRedirects = function (orig, host, path, ip, file)
    if nil == host then host = orig end
    if nil == path then path = "" end
    if nil == file then file = "" end
    local po = url.parse("http://" .. orig)
    local ph = url.parse("http://" .. host)
    if (nil ~= ip) and ("redir" ~= ip) then
	if "" ~= file then
	    D("checking redirected file " .. po.host .. " " .. file)
	    checkPaths(po.host, ip, path, file)
	else
	    checkPaths(po.host, ip, path)
	end
    else
	if orig == host then D("checkRedirects " .. orig .. "" .. file) else D("checkRedirects " .. orig .. " -> " .. host) end
-- TODO - use checkPaths() here ^^^ on the original domain name, coz that's not getting caught yet.  I think.
	gatherIPs(ph.host)
	for k, v in pairs(IP[ph.host]) do
	    D("  DNS record " .. v .. " " .. ph.host .. " -> " .. k)
	    if v == "A" then
    		if testing("IPv4") then execute("ionice -c3 ./mirror-checker.lua " .. sendArgs .. " " .. orig .. path .. " " .. k .. " " .. file .." &") end
	    elseif v == "AAAA" then
		if testing("IPv6") then execute("ionice -c3 ./mirror-checker.lua " .. sendArgs .. " " .. orig .. path .. " [" .. k .. "] " .. file .. " &") end
	    elseif v == "CNAME" then
		checkRedirects(orig, k, path, ip, file)  -- Check the original, with the DNS records from the CNAME.  Do not check the CNAME, it's just asource of IPs.
	    elseif v == "SRV" then
		print("SVR record found, now what do we do?")
	    end
	end
    end
end

local downloads = function (cut, host, URL)
    if 0 ~= cut then cd = " --cut-dirs=" .. cut .. " " else cd = "" end
    if nil == URL then URL = "/" end
    local lock = "%s-" .. host .. ".log "
    local log = " --rejected-log=results/wget-%s_REJECTS-" .. host .. ".log -a results/wget-%s-" .. host .. ".log "
    I("starting file download commands for " .. host .. " " .. URL)
    local cm = "ionice -c3 " .. downloadLock .. lock:format("debs") .. download .. log:format("debs", "debs") .. cd
    for i, s in pairs(referenceDebs) do
	cm = cm .. " https://" .. host .. URL .. "/" .. s
    end
    for i, s in pairs(releases) do
	execute(cm .. " &")
	cm = "ionice -c3 " .. downloadLock .. lock:format(s) .. download .. log:format(s, s) .. cd
	if repoExists(s .. k) then
	    for j, k in pairs(releaseFiles) do
		cm = cm .. " https://" .. host .. URL .. "/merged/dists/" .. s .. k
	    end
	end
    end
    execute(cm .. " &")
end

local getMirrors = function ()
    local mirrors = {}
    local host = ""
    local m = {}
    local URL = "https://" .. options.referenceSite.value .. "/mirror_list.txt"
    I("getting mirrors.")
    local p, c, h = http.request(URL)
    if nil == p then E(c .. " fetching " .. URL) else
	for l in p:gmatch("\n*([^\n]+)\n*") do
	    local t, d = l:match("(%a*):%s*(.*)")
	    if "FQDN" == t then
		if "" ~= host then
		    mirrors[host] = m
		    m = {}
		end
		host = d
	    end
	    m[t] = d
	end
	if "" ~= host then
	    mirrors[host] = m
	end
    end
    return mirrors
end


if 0 ~= #args then
    local option = ""
    for i, a in pairs(args) do
	if ("--help" == a) or ("-h" == a) then
	    print("I should write some docs, huh?")
	elseif "--version" == a then
	    print("mirror-checker-lua version 0.1 WIP development version")
	elseif "-v" == a then
	    verbosity = verbosity + 1
	    sendArgs = sendArgs .. a .. " "
	elseif "-q" == a then
	    verbosity = -1
	    sendArgs = sendArgs .. a .. " "
	elseif "-k" == a then
	    keep = true
	elseif "--" == a:sub(1, 2) then
	    local s, e = a:find("=")
	    if nil == s then e = -1 end
	    option = a:sub(3, e - 1)
	    local o = options[option]
	    if nil == o then
		print("Unknown option --" .. option)
		option = ""
	    else
		option = a
		sendArgs = sendArgs .. a .. " "
		local s, e = a:find("=")
		if nil == s then e = 0 end
		option = a:sub(3, e - 1)
		if "table" == options[option].typ then
		    local result = {}
		    for t in (a:sub(e + 1) .. ","):gmatch("([+%-]?%w*),") do
			local f = t:sub(1, 1)
			local n = t:sub(2, -1)
			if ("+" ~= f) and ("-" ~= f) then
			    table.insert(result, t)
			end
		    end
		    if 0 ~= #result then
			options[option].value = result
		    else
			for t in (a:sub(e + 1) .. ","):gmatch("([+%-]?%w*),") do
			    local f = t:sub(1, 1)
			    local n = t:sub(2, -1)
			    if "+" == f then
				table.insert(options[option].value, n)
			    elseif "-" == f then
				local r = {}
				for i, k in pairs(options[option].value) do
				    if k ~= n then table.insert(r, k) end
				end
				options[option].value = r
			    end
			end
		    end
		else
		    options[option].value = a
		end
		option = ""
	    end
	elseif "-" == a:sub(1, 1) then
	    print("Unknown option " .. a)
	else
	    table.insert(arg, a)
	end
    end
end

--printTable(options.tests.value, "", "tests")

execute("mkdir -p results")

if 0 < #arg then
    if "/" == arg[1]:sub(-1, -1) then
	W("slash at end of path!  " .. arg[1])
	arg[1] = arg[1]:sub(1, -2)
    end
    local pu = url.parse("http://" .. arg[1])
    if nil ~= arg[2] then
	logFile, e = io.open("results/mirror-checker-lua_" .. pu.host .. "_" .. arg[2] .. ".log", "a+")
    else
	logFile, e = io.open("results/mirror-checker-lua_" .. pu.host .. ".log", "a+")
    end
    if nil == logFile then C("opening log file - " .. e); return end
    I("Starting tests for " ..arg[1] .. " with these tests - " .. table.concat(options.tests.value, ", "))
    if nil ~= arg[2] then I("  Using IP " .. arg[2]) end
    if nil ~= arg[3] then I("  Using file " .. arg[3]); end
    if testing("Integrity") or testing("Updated") then
	if nil == arg[3] then 
	    if not keep then execute("rm -fr results/" .. pu.host) end
	    cut = 0
	    for t in arg[1]:gmatch("(/)") do
		cut = cut + 1
	    end
	    downloads(cut, pu.host, pu.path)
	    checkExes("mirror-checker.lua " .. sendArgs)
	    checkExes(downloadLock)
	end
    end
    checkRedirects(pu.host, pu.host, pu.path, arg[2], arg[3])
    logFile:close()
else
    if not keep then os.execute("rm -f results/*.log") end
    os.execute("rm -f results/*.check")
    os.execute("mkdir -p results; touch results/stamp")
    logFile, e = io.open("results/mirror-checker-lua.log", "a+")
    if nil == logFile then C("opening log file - " .. e); return end
    I("Starting tests " .. table.concat(options.tests.value, ", "))
    execute("mkdir -p results")
    local mirrors = getMirrors()
    mirrors[options.referenceSite.value] = nil
    checkRedirects(options.referenceSite.value)
--    forkIP(options.referenceSite.value)
    checkRedirects("deb.devuan.org")
--    forkIP("deb.devuan.org")
    for k, m in pairs(mirrors) do
	if "/" == m.BaseURL:sub(-1, -1) then
	    W("slash at end of BaseURL in mirror_list.txt!  " .. m.BaseURL)
	    m.BaseURL = m.BaseURL:sub(1, -2)
	end
	local pu = url.parse("http://" .. m.BaseURL)
	checkRedirects(m.BaseURL)
--	forkIP(m.BaseURL)
        checkExes("mirror-checker.lua " .. sendArgs)
        if testing("Integrity") or testing("Updated") then checkExes(downloadLock) end
    end
    while 1 <= checkExes("mirror-checker.lua " .. sendArgs) do os.execute("sleep 30") end
    if testing("Integrity") or testing("Updated") then
	while 0 < checkExes(downloadLock) do os.execute("sleep 30") end
    end
    os.execute("rm -f results/*.check")
    logFile:close()
end