#!/opt/lnxall_app/bin/lua

-- Created by jiaqiang.ye@lnxall.com
-- Simple Local Device Command Batch Executor
-- 2025/03/06

local cjson = require 'cjson'
local posix = require 'posix'
local invoker = require 'invoker'

local gfmt = string.format

local g_iplist         = nil -- local ip list
local g_linecmd        = nil -- shell command to be executed on remote device
local g_file0, g_file1 = nil, nil -- files to be copied to remote device
local g_onhost = nil -- whether the command line is invoked on host

local function badstr(str)
	if type(str) == "string" and #str > 0 then return false end
	return true
end

local function find_localip(ndev)
	local okay, output = invoker.invoke(invoker.OUTPUT, "ip", "neigh")
	if okay ~= 0 or type(output) ~= "string" then
		io.stderr:write("Error, command `ip neigh` has failed.\n")
		io.stderr:flush()
		return nil
	end

	local needle = "dev " .. ndev
	local nip, iptab = 0, {}
	for line in string.gmatch(output, "[^\r\n]+") do
		if string.find(line, needle, 1, true) then
			local ip = string.match(line, "^([%d%.]+)%s")
			if invoker.isipv4(ip) then
				nip = nip + 1
				iptab[nip] = ip
			end
		end
	end

	if nip == 0 then
		io.stderr:write(gfmt("Error, no device found via network device: %s\n", ndev))
		io.stderr:flush()
		return nil
	end
	if nip >= 2 then table.sort(iptab) end
	io.stdout:write("{\n\t\"iplist\":[\n")
	for i, ip in ipairs(iptab) do
		if i >= 2 then
			io.stdout:write(gfmt(",\n\t\t\"%s\"", ip))
		else
			io.stdout:write(gfmt("\t\t\"%s\"", ip))
		end
		io.stdout:flush()
	end
	io.stdout:write("\n\t]\n}\n")
	io.stdout:flush()
	return iptab
end

local function read_localip(iplist)
	local iph = io.open(iplist, "rb")
	if not iph then
		io.stderr:write(gfmt("Error, failed to open '%s'\n", iplist))
		io.stderr:flush()
		return nil
	end

	local ipd = iph:read("*a")
	iph:close(); iph = nil
	if badstr(ipd) then
		io.stderr:write(gfmt("Error, failed to read '%s'\n", iplist))
		io.stderr:flush()
		return nil
	end

	local okay, ipjson = pcall(cjson.decode, ipd)
	local ilist = type(ipjson) == "table" and ipjson["iplist"] or nil
	if not okay or type(ilist) ~= "table" then
		io.stderr:write(gfmt("Error, invalid iplist file: '%s'\n", iplist))
		io.stderr:flush()
		return nil
	end

	return ilist
end

local function check_args(iplist, linecmd, file0, file1)
	if badstr(iplist) or badstr(linecmd) then
		io.stderr:write("Error, invalid command-line arguments\n")
		io.stderr:flush()
		return false
	end

	if invoker.invoke(invoker.NOSTDIO, "ip", "link", "show", "dev", iplist) == 0 then
		g_iplist = find_localip(iplist)
	elseif posix.access(iplist) == 0 then
		g_iplist = read_localip(iplist)
	elseif invoker.isipv4(iplist) then
		g_iplist = { [1] = iplist }
	end

	if type(g_iplist) ~= "table" then
		io.stderr:write("Error, failed to construct IP list.\n")
		io.stderr:flush()
		return false
	end

	local lcmd = invoker.readfile(linecmd, invoker.TRIMEND)	
	if type(lcmd) ~= "string" or #lcmd == 0 or string.find(lcmd, "\r", 1, true) or string.find(lcmd, "\n", 1, true) then
		io.stderr:write(gfmt("Error, invalid one-line command file: '%s'\n", linecmd))
		io.stderr:flush()
		return false
	end
	g_linecmd = lcmd
	local needle = "# BATCH_HOST"
	local lencmd, lenn = string.len(lcmd), string.len(needle)
	if lencmd > lenn and lcmd:sub(lencmd - lenn + 1) == needle then
		g_onhost = true
	end

	if not badstr(file0) then
		if posix.access(file0) ~= 0 then
			io.stderr:write(gfmt("Error, cannot find file: '%s'\n", file0))
			io.stderr:flush()
			return false
		end
		g_file0 = file0
	end

	if not badstr(file1) then
		if posix.access(file1) ~= 0 then
			io.stderr:write(gfmt("Error, cannot find file: '%s'\n", file1))
			io.stderr:flush()
			return false
		end
		g_file1 = file1
	end

	return true
end

local function ssh_copy(ip, file)
	io.stdout:write(gfmt("[%s] Copying '%s' to '%s' ...\n", os.date(), file, ip))
	io.stdout:flush()
	local okay = invoker.invoke(0, "sshpass", "-p", "lnxall123", "scp", "-r",
		"-oConnectTimeout=5", "-oUserKnownHostsFile=/dev/null", "-oStrictHostKeyChecking=no",
		file, gfmt("root@%s:/tmp/", ip))
	return okay == 0
end

local function take_operation(ip)
	if g_file0 and not ssh_copy(ip, g_file0) then
		io.stderr:write(gfmt("Error, failed to copy '%s' to '%s'\n", g_file0, ip))
		io.stderr:flush()
		return false
	end

	if g_file1 and not ssh_copy(ip, g_file1) then
		io.stderr:write(gfmt("Error, failed to copy '%s' to '%s'\n", g_file1, ip))
		io.stderr:flush()
		return false
	end

	local linecmd, okay, needle = nil, nil, "DEVIPADDR"
	if string.find(g_linecmd, needle, 1, true) then
		linecmd = string.gsub(g_linecmd, needle, ip)
	else
		linecmd = g_linecmd
	end

	if g_onhost then
		io.stdout:write(gfmt("[%s] Invoking command on host '%s': %s\n", os.date(), ip, linecmd))
		io.stdout:flush()
		okay = invoker.invoke(0, "sh", "-c", linecmd)
	else
		io.stdout:write(gfmt("[%s] Invoking command on device '%s': %s\n", os.date(), ip, linecmd))
		io.stdout:flush()
		okay = invoker.invoke(0, "sshpass", "-p", "lnxall123", "ssh",
			"-oConnectTimeout=5", "-oUserKnownHostsFile=/dev/null", "-oStrictHostKeyChecking=no",
			gfmt("root@%s", ip), linecmd)
	end
	return okay == 0
end

local function batch_operation()
	local ni = 0
	local sep0 = string.rep("*", 120) .. "\n"
	local sep1 = string.rep("=", 120) .. "\n"
	while true do
		ni = ni + 1
		local ip = g_iplist[ni]
		if ip == nil then break end
		if invoker.isipv4(ip) then
			io.stdout:write(sep0) ; io.stdout:flush()
			take_operation(ip)
			io.stdout:write(sep1) ; io.stdout:flush()
		else
			ip = tostring(ip)
			if type(ip) ~= "string" then ip = "unknown" end
			io.stderr:write(gfmt("Error, invalid IPv4 address: %s\n", ip))
			io.stderr:flush()
			break
		end
	end
end

local function batchlocal_help()
	io.stdout:write("Usage:\n")
	io.stdout:write("\tbatchlocal NETDEV|IPLIST.JSON ONELINE-CMD.SH [COPY_FILE0] [COPY_FILE1]\n")
	io.stdout:write("\tExample: batchlocal eth0 oneline-cmd.sh XXXX0.elf\n")
	io.stdout:write("\tExample: batchlocal iplist.json oneline-cmd.sh XXXX0.elf XXXX1.elf\n")
	io.stdout:write("Note:\n\tContents of `oneline-cmd.sh` could be:\n")
	io.stdout:write("cd /tmp && chmod 755 XXXX0.elf && ./XXXX0.elf -n ; rm -f -v *.elf\n")
	io.stdout:flush()
end

if not check_args(arg[1], arg[2], arg[3], arg[4]) then
	batchlocal_help()
	os.exit(1)
end
batch_operation()
os.exit(0)