#!/usr/bin/env lua

--%# family=auto
--%# capabilities=autoconf

-- Configuration with environment variable overrides
local DEFAULT_HOST = os.getenv("host") or os.getenv("GPON_HOST") or "192.168.1.254"
local USERNAME = os.getenv("username") or os.getenv("GPON_USER") or "userAdmin"
local PASSWORD = os.getenv("password") or os.getenv("GPON_PASS") or "xxxxxxxxxx"

local LABEL_KEYS = { "WiFi2", "WiFi6", "LAN1", "LAN2", "LAN3", "LAN4" }

local LABELS = {
    WiFi2 = "WiFi2(2.4G)",
    WiFi6 = "WiFi6(5G)",
    LAN1  = "LAN1(AirSt)",
    LAN2  = "LAN2(Proxy)",
    LAN3  = "LAN3(PC)",
    LAN4  = "TEL",
}

local COLOURS = {
    WiFi2 = "FFB000",
    WiFi6 = "FF4040",
    LAN1  = "00D0FF",
    LAN2  = "2060FF",
    LAN3  = "20FF00",
    LAN4  = "8040FF",
}

local function print_config()
    print("graph_title 光ファイバ GPON ネットワーク状態")
    print("graph_category network")
    print("graph_vlabel packet/秒 (-)受信 / (+)送信")
    print("graph_args --base 1024")
    print("graph_printf %4.1lf")

    -- WiFi fields
    for _, sr in ipairs({ "Received", "Sent" }) do
        for i = 1, 8 do
            local field = "WiFi" .. i
            if LABELS[field] then
                print(field .. sr .. ".type COUNTER")
                print(field .. sr .. ".min 0")
                if sr == "Received" then
                    print(field .. sr .. ".graph no")
                else
                    print(field .. sr .. ".draw AREASTACK")
                    print(field .. sr .. ".negative " .. field .. "Received")
                end
                print(field .. sr .. ".label " .. LABELS[field])
                print(field .. sr .. ".colour " .. COLOURS[field])
            end
        end
    end

    -- LAN fields
    for _, sr in ipairs({ "Received", "Sent" }) do
        for i = 1, 4 do
            local field = "LAN" .. i
            if LABELS[field] then
                print(field .. sr .. ".type COUNTER")
                print(field .. sr .. ".min  0")
                if sr == "Received" then
                    print(field .. sr .. ".graph no")
                else
                    print(field .. sr .. ".draw AREASTACK")
                    print(field .. sr .. ".negative " .. field .. "Received")
                end
                print(field .. sr .. ".label " .. LABELS[field])
                print(field .. sr .. ".colour " .. COLOURS[field])
            end
        end
    end

    print("LAN1Sent.warning 1000")
    print("LAN1Received.warning -1000")
end

local function to_base64(data)
    local b64chars = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"
    local res = {}
    local len = #data
    for i = 1, len, 3 do
        local b1 = data:byte(i)
        local b2 = data:byte(i + 1)
        local b3 = data:byte(i + 2)

        local n = b1 * 65536 + (b2 or 0) * 256 + (b3 or 0)
        local c1 = math.floor(n / 262144) % 64 + 1
        local c2 = math.floor(n / 4096) % 64 + 1
        local c3 = math.floor(n / 64) % 64 + 1
        local c4 = n % 64 + 1

        res[#res + 1] = b64chars:sub(c1, c1)
        res[#res + 1] = b64chars:sub(c2, c2)
        res[#res + 1] = b2 and b64chars:sub(c3, c3) or "="
        res[#res + 1] = b3 and b64chars:sub(c4, c4) or "="
    end
    return table.concat(res)
end

local function base64url_escape(b64)
    return b64:gsub("%+", "-"):gsub("/", "_"):gsub("=", ".")
end

local function to_base64url(data)
    return to_base64(data):gsub("%+", "-"):gsub("/", "_"):gsub("=+$", "")
end

local function url_encode(str)
    return (str:gsub("[^%w%-_%.!~%*'%(%)]", function(c)
        return string.format("%%%02X", string.byte(c))
    end))
end

local function to_hex(data)
    return (data:gsub(".", function(c)
        return string.format("%02x", string.byte(c))
    end))
end

local function random_bytes(n)
    local f = io.open("/dev/urandom", "rb")
    if not f then
        error("Cannot open /dev/urandom for random bytes")
    end
    local bytes = f:read(n)
    f:close()
    if #bytes ~= n then
        error("Failed to read required number of random bytes")
    end
    return bytes
end

local function write_temp_file(content)
    local path = os.tmpname()
    local f = io.open(path, "wb")
    if not f then
        error("Cannot open temporary file: " .. tostring(path))
    end
    f:write(content)
    f:close()
    return path
end

local function aes_128_cbc_encrypt(plaintext, key, iv)
    local in_file = write_temp_file(plaintext)
    local out_file = os.tmpname()
    local cmd = string.format("openssl enc -aes-128-cbc -K %s -iv %s -in %q -out %q 2>/dev/null",
        to_hex(key), to_hex(iv), in_file, out_file)
    local ok = os.execute(cmd)
    os.remove(in_file)
    if ok ~= 0 and ok ~= true then
        os.remove(out_file)
        error("AES-128-CBC encryption failed")
    end
    local f = io.open(out_file, "rb")
    if not f then
        os.remove(out_file)
        error("Failed to read AES encrypted output")
    end
    local ct = f:read("*a")
    f:close()
    os.remove(out_file)
    return ct
end

local function rsa_pkcs1_encrypt(plaintext, pem)
    local key_file = write_temp_file(pem)
    local in_file = write_temp_file(plaintext)
    local out_file = os.tmpname()
    local cmd = string.format("openssl pkeyutl -encrypt -pubin -inkey %q -pkeyopt rsa_padding_mode:pkcs1 -in %q -out %q 2>/dev/null",
        key_file, in_file, out_file)
    local ok = os.execute(cmd)
    os.remove(key_file)
    os.remove(in_file)
    if ok ~= 0 and ok ~= true then
        os.remove(out_file)
        error("RSA PKCS#1 v1.5 encryption failed")
    end
    local f = io.open(out_file, "rb")
    if not f then
        os.remove(out_file)
        error("Failed to read RSA encrypted output")
    end
    local ct = f:read("*a")
    f:close()
    os.remove(out_file)
    return ct
end

-- HTTP Client Implementation (LuaSocket with fallback to curl command)
local has_socket, http = pcall(require, "socket.http")
local has_ltn12, ltn12 = pcall(require, "ltn12")

local function http_request(req)
    if has_socket and has_ltn12 then
        local resp = {}
        local headers = {}
        for k, v in pairs(req.headers or {}) do
            headers[k] = v
        end
        if req.body then
            headers["Content-Length"] = tostring(#req.body)
        end
        local prev_timeout = http.TIMEOUT
        http.TIMEOUT = req.timeout or 10
        local _, code, resp_headers = http.request{
            url = req.url,
            method = req.method or "GET",
            headers = headers,
            source = req.body and ltn12.source.string(req.body) or nil,
            sink = ltn12.sink.table(resp)
        }
        http.TIMEOUT = prev_timeout
        return table.concat(resp), tonumber(code) or 0, resp_headers or {}
    else
        -- Fallback to curl CLI
        local header_args = {}
        for k, v in pairs(req.headers or {}) do
            header_args[#header_args + 1] = string.format("-H %q", k .. ": " .. v)
        end
        local body_arg = ""
        local temp_body_file = nil
        if req.body then
            temp_body_file = write_temp_file(req.body)
            body_arg = string.format("--data-binary @%q", temp_body_file)
        end
        local method_arg = ""
        if req.method == "POST" and not req.body then
            method_arg = "-X POST"
        end
        local timeout = req.timeout or 10
        local resp_file = os.tmpname()
        local cmd = string.format("curl -s -i --max-time %d %s %s %s %q > %q 2>/dev/null",
            timeout,
            method_arg,
            table.concat(header_args, " "),
            body_arg,
            req.url,
            resp_file)
        os.execute(cmd)
        if temp_body_file then os.remove(temp_body_file) end

        local f = io.open(resp_file, "rb")
        if not f then
            os.remove(resp_file)
            return "", 0, {}
        end
        local content = f:read("*a")
        f:close()
        os.remove(resp_file)

        local h_end = content:find("\r?\n\r?\n")
        if not h_end then return "", 0, {} end
        local raw_headers = content:sub(1, h_end - 1)
        local _, body_start = content:find("\r?\n\r?\n")
        local body = content:sub(body_start + 1)

        local status_line = raw_headers:match("^[^\r\n]+")
        local status_code = tonumber(status_line and status_line:match("%s(%d%d%d)%s")) or 0

        local headers = {}
        for line in raw_headers:gmatch("[^\r\n]+") do
            local k, v = line:match("^([^:]+):%s*(.*)$")
            if k and v then
                local lk = k:lower()
                if headers[lk] then
                    headers[lk] = headers[lk] .. "; " .. v
                else
                    headers[lk] = v
                end
            end
        end
        return body, status_code, headers
    end
end

local function get_header(headers, name)
    name = name:lower()
    for k, v in pairs(headers or {}) do
        if k:lower() == name then
            return v
        end
    end
    return nil
end

local function parse_pubkey(html)
    local raw_key = html:match("var%s+pubkey%s*=%s*'([^']+)';")
    if not raw_key then
        error("Could not find pubkey in login page")
    end
    local body = raw_key:gsub("%-%-%-%-%-BEGIN PUBLIC KEY%-%-%-%-%-", "")
                        :gsub("%-%-%-%-%-END PUBLIC KEY%-%-%-%-%-", "")
                        :gsub("\\", "")
                        :gsub("%s+", "")
    local lines = {}
    for i = 1, #body, 64 do
        lines[#lines + 1] = body:sub(i, i + 63)
    end
    return "-----BEGIN PUBLIC KEY-----\n" .. table.concat(lines, "\n") .. "\n-----END PUBLIC KEY-----\n"
end

local function parse_lan_status(html)
    local stats = { wifi = {}, lan = {} }

    -- wlan_status parsing
    local wlan_block = html:match("var%s+wlan_status%s*=%s*{(.-)%s*};")
    if wlan_block then
        local cur_id = nil
        local cur_enable = 0
        local sent = nil
        local recv = nil

        for line in wlan_block:gmatch("[^\r\n]+") do
            local id_m = line:match("^[,%s]*(%d+):{")
            if id_m then
                cur_id = id_m
                cur_enable = 0
                sent = nil
                recv = nil
            end
            local en_m = line:match('"Enable":(%d+),')
            if en_m then cur_enable = tonumber(en_m) end

            local sent_m = line:match('"TotalPacketsSent":(%d+),')
            if sent_m then sent = tonumber(sent_m) end

            local recv_m = line:match('"TotalPacketsReceived":(%d+),')
            if recv_m then recv = tonumber(recv_m) end

            if cur_id and cur_enable ~= nil and sent ~= nil and recv ~= nil then
                stats.wifi[cur_id] = { sent = sent, recv = recv, enable = cur_enable }
                cur_id = nil
            end
        end
    end

    -- lan_ether parsing
    local lan_block = html:match("var%s+lan_ether%s*=%s*{(.-)%s*};")
    if lan_block then
        local cur_id = nil
        local sent = nil
        local recv = nil

        for line in lan_block:gmatch("[^\r\n]+") do
            local id_m = line:match("^[,%s]*(%d+):{")
            if id_m then
                cur_id = id_m
                sent = nil
                recv = nil
            end
            local sent_m = line:match("PacketsSent:(%d+),")
            if sent_m then sent = tonumber(sent_m) end

            local recv_m = line:match("PacketsReceived:(%d+),")
            if recv_m then recv = tonumber(recv_m) end

            if cur_id and sent ~= nil and recv ~= nil then
                stats.lan[cur_id] = { sent = sent, recv = recv }
                cur_id = nil
            end
        end
    end

    return stats
end

local function fetch_router_data(host, username, password)
    local base_url = "http://" .. host

    -- 1. GET /
    local html1, _, headers1 = http_request{
        url = base_url .. "/",
        timeout = 10
    }
    local cookies1 = get_header(headers1, "set-cookie") or ""

    if not html1:find("GPON") then
        error(string.format("Target %s does not appear to be GPON gateway", host))
    end

    local pem = parse_pubkey(html1)
    local nonce = html1:match('var%s+nonce%s*=%s*"(.-)";')
    local token = html1:match('var%s+token%s*=%s*"(.-)";')
    if not nonce or not token then
        error("Could not find nonce/token in login page")
    end

    -- Generate random dec_key and dec_iv (16 bytes each)
    local dec_key_bytes = random_bytes(16)
    local dec_iv_bytes = random_bytes(16)
    local dec_key_b64 = to_base64(dec_key_bytes)
    local dec_iv_b64 = to_base64(dec_iv_bytes)

    local postdata = "&username=" .. username ..
                     "&password=" .. url_encode(password) ..
                     "&csrf_token=" .. token ..
                     "&nonce=" .. nonce ..
                     "&enckey=" .. base64url_escape(dec_key_b64) ..
                     "&enciv=" .. base64url_escape(dec_iv_b64)

    -- AES-128-CBC encryption of postdata
    local aes_key = random_bytes(16)
    local aes_iv = random_bytes(16)
    local ct_buffer = aes_128_cbc_encrypt(postdata, aes_key, aes_iv)
    local ct = to_base64url(ct_buffer)

    -- RSA PKCS#1 v1.5 encryption of aesinfo
    local aesinfo = to_base64(aes_key) .. " " .. to_base64(aes_iv)
    local ck_buffer = rsa_pkcs1_encrypt(aesinfo, pem)
    local ck = base64url_escape(to_base64(ck_buffer))

    local encrypt_data = "encrypted=1&ct=" .. ct .. "&ck=" .. ck

    -- 2. POST /login.cgi
    local _, r2_code, headers2 = http_request{
        url = base_url .. "/login.cgi",
        method = "POST",
        headers = {
            ["Content-Type"] = "application/x-www-form-urlencoded; charset=UTF-8",
            ["X-Requested-With"] = "XMLHttpRequest",
            ["Referer"] = base_url .. "/",
            ["Cookie"] = cookies1
        },
        body = encrypt_data,
        timeout = 10
    }

    local sid = get_header(headers2, "x-sid")
    local cookies2 = get_header(headers2, "set-cookie") or ""
    if r2_code ~= 299 and not sid then
        error(string.format("Login failed with status %s", tostring(r2_code)))
    end

    local session_cookies = {}
    if cookies1 ~= "" then session_cookies[#session_cookies + 1] = cookies1 end
    if cookies2 ~= "" then session_cookies[#session_cookies + 1] = cookies2 end
    if sid and sid ~= "" then session_cookies[#session_cookies + 1] = "sid=" .. sid end
    local cookie_header = table.concat(session_cookies, "; ")

local function perform_logout(base_url, cookie_header, sid)
    local logout_url = base_url .. "/login.cgi?out"
    local sid_str = sid or ""

    -- Asynchronous logout in background via curl (&), so the caller receives metrics immediately
    -- without waiting ~1.5s for the logout CGI response.
    local cmd = string.format("curl -s -m 5 -H %q -H %q %q >/dev/null 2>&1 &",
        "Cookie: " .. cookie_header,
        "X-SID: " .. sid_str,
        logout_url)
    local ok_exec, _, code_exec = os.execute(cmd)
    if ok_exec == true or ok_exec == 0 or code_exec == 0 then
        return
    end

    -- Fallback: synchronous logout if curl execution failed
    pcall(function()
        http_request{
            url = logout_url,
            headers = {
                ["Cookie"] = cookie_header,
                ["X-SID"] = sid_str
            },
            timeout = 5
        }
    end)
end

    -- 3. GET /lan_status.cgi?lan & 4. Logout
    local ok_req, res = pcall(function()
        local status_html, _, _ = http_request{
            url = base_url .. "/lan_status.cgi?lan",
            headers = {
                ["Cookie"] = cookie_header,
                ["X-SID"] = sid or ""
            },
            timeout = 10
        }
        return parse_lan_status(status_html)
    end)

    -- Always ensure logout is executed so the session is cleanly closed on the router
    perform_logout(base_url, cookie_header, sid)

    if not ok_req then
        error(res)
    end
    return res
end

local function main()
    local arg1 = arg and arg[1]

    if arg1 == "autoconf" then
        print("yes")
        os.exit(0)
    end

    if arg1 == "config" then
        print_config()
        os.exit(0)
    end

    -- Try target host. If connection fails or target is not GPON and host was not 192.168.1.254,
    -- fallback to 192.168.1.254.
    local host = DEFAULT_HOST
    local stats = nil

    local ok, res = pcall(fetch_router_data, host, USERNAME, PASSWORD)
    if ok then
        stats = res
    else
        local err = res
        if host ~= "192.168.1.254" then
            local ok2, res2 = pcall(fetch_router_data, "192.168.1.254", USERNAME, PASSWORD)
            if ok2 then
                stats = res2
                host = "192.168.1.254"
            else
                error(err)
            end
        else
            error(err)
        end
    end

    -- Output stats for munin
    for _, key in ipairs(LABEL_KEYS) do
        if key:find("^WiFi") then
            local id = key:gsub("^WiFi", "")
            local w = stats.wifi[id]
            if w and w.enable and w.enable ~= 0 then
                print(string.format("%sSent.value %d", key, w.sent))
                print(string.format("%sReceived.value %d", key, w.recv))
            end
        elseif key:find("^LAN") then
            local id = key:gsub("^LAN", "")
            local l = stats.lan[id]
            if l then
                print(string.format("%sSent.value %d", key, l.sent))
                print(string.format("%sReceived.value %d", key, l.recv))
            end
        end
    end
    io.stdout:flush()
end

local ok, err = pcall(main)
if not ok then
    local msg = tostring(err):gsub("^.-:%d+:%s*", "")
    io.stderr:write("Error fetching router stats: " .. msg .. "\n")
    os.exit(1)
end