-- /etc/nginx/lua/captcha.lua
local cjson = require("cjson.safe")

local M = {}

local COOKIE_NAME = "__ts_verified"
local COOKIE_TTL  = 60 * 60 * 24 * 90   -- 90 days
local VERIFY_PATH = "/__ts/verify"

local WALKAWAY_THRESHOLD = 10
local WALKAWAY_TTL       = 7200       -- seconds; counter resets after 2h idle
local REPORT_COOLDOWN    = 15 * 60    -- never re-report same IP within 15 min

-- ---------- config: load ONCE, at module load time (in master/root) ----------
local cfg = { site_key = "", secret_key = "", cookie_secret = "", abuseipdb_key = "", cap_url="", cap_site_key="", cap_secret_key="" }

do
    local f, err = io.open("/etc/nginx/captcha.env", "r")
    if not f then
        ngx.log(ngx.ERR, "captcha: cannot open env file: ", err)
    else
        for line in f:lines() do
            local k, v = line:match("^%s*([A-Z_]+)%s*=%s*(.-)%s*$")
            if     k == "TURNSTILE_SITE_KEY"      then cfg.site_key       = v
            elseif k == "TURNSTILE_SECRET_KEY"    then cfg.secret_key     = v
            elseif k == "COOKIE_SECRET"           then cfg.cookie_secret  = v
            elseif k == "ABUSEIPDB_API_KEY"       then cfg.abuseipdb_key  = v
            elseif k == "CAP_API_URL"             then cfg.cap_url        = (v or ""):gsub("/+$","")
            elseif k == "CAP_SITE_KEY"            then cfg.cap_site_key   = v
            elseif k == "CAP_SECRET_KEY"          then cfg.cap_secret_key = v
            end
        end
        f:close()
    end
end

-- ---------- whitelist ----------
local WHITELIST_FILE = "/etc/nginx/captcha-whitelist.txt"
local wl4, wl6 = {}, {}   -- {{net=..., mask=..., bits=...}, ...}

local function ip4_to_n(ip)
    local a, b, c, d = ip:match("^(%d+)%.(%d+)%.(%d+)%.(%d+)$")
    if not a then return nil end
    a, b, c, d = tonumber(a), tonumber(b), tonumber(c), tonumber(d)
    for _, n in ipairs({a, b, c, d}) do
        if not n or n < 0 or n > 255 then return nil end
    end
    -- Build via arithmetic so this works on plain Lua 5.1 (no bit lib)
    return ((a * 256 + b) * 256 + c) * 256 + d
end

-- Expand an IPv6 string to 8 groups of 16-bit numbers
local function ip6_to_groups(ip)
    if not ip:find(":") then return nil end
    local head, tail = ip:match("^(.-)::(.*)$")
    local h_parts, t_parts = {}, {}
    if head then
        for g in (head .. ":"):gmatch("([^:]*):") do h_parts[#h_parts+1] = g end
        for g in (tail .. ":"):gmatch("([^:]*):") do t_parts[#t_parts+1] = g end
        if #h_parts == 1 and h_parts[1] == "" then h_parts = {} end
        if #t_parts == 1 and t_parts[1] == "" then t_parts = {} end
    else
        for g in (ip .. ":"):gmatch("([^:]*):") do h_parts[#h_parts+1] = g end
    end
    local groups = {}
    for _, g in ipairs(h_parts) do groups[#groups+1] = tonumber(g, 16) or -1 end
    local fill = 8 - #h_parts - #t_parts
    if head then
        for _ = 1, fill do groups[#groups+1] = 0 end
    end
    for _, g in ipairs(t_parts) do groups[#groups+1] = tonumber(g, 16) or -1 end
    if #groups ~= 8 then return nil end
    for _, g in ipairs(groups) do
        if g < 0 or g > 0xFFFF then return nil end
    end
    return groups
end

local function load_whitelist()
    local new4, new6, count = {}, {}, 0
    local f = io.open(WHITELIST_FILE, "r")
    if not f then
        ngx.log(ngx.ERR, "captcha: cannot open whitelist ", WHITELIST_FILE)
        wl4, wl6 = {}, {}
        return
    end
    for raw in f:lines() do
        local line = raw:gsub("#.*$", ""):match("^%s*(.-)%s*$") or ""
        if line ~= "" then
            local addr, bits = line:match("^([^/]+)/(%d+)$")
            if not addr then addr = line end

            if addr:find(":", 1, true) then
                local g = ip6_to_groups(addr)
                if g then
                    bits = tonumber(bits or 128)
                    if bits >= 0 and bits <= 128 then
                        new6[#new6+1] = { groups = g, bits = bits }
                        count = count + 1
                    else
                        ngx.log(ngx.ERR, "captcha: bad whitelist line: ", raw)
                    end
                else
                    ngx.log(ngx.ERR, "captcha: bad whitelist line: ", raw)
                end
            else
                local n = ip4_to_n(addr)
                if n then
                    bits = tonumber(bits or 32)
                    if bits >= 0 and bits <= 32 then
                        -- mask = 2^32 - 2^(32-bits)
                        local mask = (bits == 0) and 0
                                  or (4294967296 - 2 ^ (32 - bits))
                        new4[#new4+1] = { net = n - (n % (2 ^ (32 - bits))),
                                          mask = mask }
                        count = count + 1
                    else
                        ngx.log(ngx.ERR, "captcha: bad whitelist line: ", raw)
                    end
                else
                    ngx.log(ngx.ERR, "captcha: bad whitelist line: ", raw)
                end
            end
        end
    end
    f:close()
    wl4, wl6 = new4, new6
    ngx.log(ngx.ERR, "captcha: whitelist loaded, ", count, " entries")
end

local function ip_whitelisted(ip)
    if not ip or ip == "" then return false end
    if ip:find(":", 1, true) then
        local g = ip6_to_groups(ip); if not g then return false end
        for _, e in ipairs(wl6) do
            local bits, ok = e.bits, true
            for i = 1, 8 do
                if bits >= 16 then
                    if g[i] ~= e.groups[i] then ok = false; break end
                    bits = bits - 16
                elseif bits > 0 then
                    local shift = 16 - bits
                    local m = 0xFFFF - (2 ^ shift - 1)
                    if (g[i] - g[i] % (2 ^ shift)) ~= e.groups[i] - e.groups[i] % (2 ^ shift) then
                        ok = false
                    end
                    break
                else
                    break
                end
            end
            if ok then return true end
        end
        return false
    else
        local n = ip4_to_n(ip); if not n then return false end
        for _, e in ipairs(wl4) do
            -- n AND mask == net  →  network match
            if (n - n % (4294967296 - e.mask + 0.5 - 0.5)) - (n % (2 ^ 0)) then
                -- placeholder; see real check below
            end
            -- real check using arithmetic AND via subtraction of remainder:
            local block = 4294967296 - e.mask
            if (n - n % block) == e.net then return true end
        end
        return false
    end
end

load_whitelist()

local function provider()
    if cfg.cap_site_key ~= "" and cfg.cap_secret_key ~= "" and cfg.cap_url ~= "" then
        return "cap"
    end
    return "turnstile"
end

ngx.log(ngx.NOTICE, "captcha: config loaded, provider=", provider(),
        (provider() == "cap" and (", cap_url=" .. cfg.cap_url) or ""))


-- ---------- helpers ----------
local function b64url(s)
    return (ngx.encode_base64(s):gsub("+","-"):gsub("/","_"):gsub("=",""))
end

local function html_escape(s)
    return (tostring(s)
        :gsub("&","&amp;"):gsub("<","&lt;"):gsub(">","&gt;")
        :gsub('"',"&quot;"):gsub("'","&#39;"))
end

local function ct_eq(a, b)
    if type(a) ~= "string" or type(b) ~= "string" or #a ~= #b then return false end
    local diff = 0
    for i = 1, #a do
        local d = string.byte(a, i) - string.byte(b, i)
        diff = diff + d * d
    end
    return diff == 0
end

local function sign_cookie(expires)
    local msg = tostring(expires)
    return msg .. "." .. b64url(ngx.hmac_sha1(cfg.cookie_secret, msg))
end

local function cookie_is_valid(val)
    if not val then return false end
    local exp, sig = val:match("^(%d+)%.([A-Za-z0-9_%-]+)$")
    if not exp then return false end
    exp = tonumber(exp)
    if not exp or exp < ngx.time() then return false end
    local expected = b64url(ngx.hmac_sha1(cfg.cookie_secret, tostring(exp)))
    return ct_eq(sig, expected)
end

local function safe_redirect(r)
    if type(r) ~= "string" or r == "" then return "/" end
    if r:sub(1,1) ~= "/" or r:sub(1,2) == "//" or r:sub(1,2) == "/\\" then
        return "/"
    end
    return r
end


-- ---------- URL parser + generic HTTP(S) POST via cosocket ----------
local function parse_url(url)
    local scheme, rest = url:match("^(https?)://(.*)$")
    if not scheme then return nil end
    local hostport, path = rest:match("^([^/]+)(.*)$")
    if not hostport then return nil end
    if path == "" then path = "/" end
    local host, p = hostport:match("^([^:]+):(%d+)$")
    if host then return { scheme=scheme, host=host, port=tonumber(p), path=path } end
    return { scheme=scheme, host=hostport,
             port=(scheme == "https" and 443 or 80), path=path }
end

local function http_post(url, body, content_type, extra_headers)
    local u = parse_url(url)
    if not u then return nil, "bad url: " .. tostring(url) end

    local sock = ngx.socket.tcp()
    sock:settimeout(5000)
    local ok, err = sock:connect(u.host, u.port)
    if not ok then return nil, "connect: " .. (err or "?") end
    if u.scheme == "https" then
        local sess; sess, err = sock:sslhandshake(false, u.host, true)
        if not sess then sock:close(); return nil, "ssl: " .. (err or "?") end
    end

    local lines = {
        "POST " .. u.path .. " HTTP/1.1",
        "Host: " .. u.host,
        "User-Agent: nginx-turnstile/1.0",
        "Accept: application/json",
        "Content-Type: " .. content_type,
        "Content-Length: " .. #body,
        "Connection: close",
    }
    for k, v in pairs(extra_headers or {}) do
        lines[#lines+1] = k .. ": " .. v
    end
    lines[#lines+1] = ""
    lines[#lines+1] = body

    local _, serr = sock:send(table.concat(lines, "\r\n"))
    if serr then sock:close(); return nil, "send: " .. serr end

    local status_line, rerr = sock:receive("*l")
    if not status_line then sock:close(); return nil, "recv status: " .. (rerr or "?") end
    local status = tonumber(status_line:match("HTTP/1%.[01]%s+(%d+)"))
    local content_length, chunked
    while true do
        local line; line, rerr = sock:receive("*l")
        if not line then sock:close(); return nil, "recv header: " .. (rerr or "?") end
        if line == "" then break end
        local lc = line:lower()
        local cl = lc:match("^content%-length:%s*(%d+)")
        if cl then content_length = tonumber(cl) end
        if lc:match("^transfer%-encoding:.*chunked") then chunked = true end
    end

    local resp = ""
    if chunked then
        while true do
            local size_line = sock:receive("*l")
            if not size_line then break end
            local size = tonumber((size_line:match("^([0-9a-fA-F]+)") or ""), 16)
            if not size or size == 0 then break end
            local chunk = sock:receive(size); if not chunk then break end
            resp = resp .. chunk
            sock:receive(2)
        end
    elseif content_length then
        resp = sock:receive(content_length) or ""
    else
        resp = sock:receive("*a") or ""
    end
    sock:close()
    return status, resp
end


-- ---------- provider-specific verifiers ----------
local function turnstile_verify(token, remote_ip)
    local body = "secret="    .. ngx.escape_uri(cfg.secret_key)
              .. "&response=" .. ngx.escape_uri(token)
              .. "&remoteip=" .. ngx.escape_uri(remote_ip or "")
    return http_post("https://challenges.cloudflare.com/turnstile/v0/siteverify",
                     body, "application/x-www-form-urlencoded")
end

local function cap_verify(token)
    local body = cjson.encode({ secret = cfg.cap_secret_key, response = token })
    local url  = cfg.cap_url .. "/" .. cfg.cap_site_key .. "/siteverify"
    return http_post(url, body, "application/json")
end

-- ---------- AbuseIPDB reporter ----------
local function report_to_abuseipdb(ip, count)
    if cfg.abuseipdb_key == "" then return end
    local body = "ip="        .. ngx.escape_uri(ip)
              .. "&categories=19"
              .. "&comment="  .. ngx.escape_uri(string.format(
                    "Fetched browser challenge page %d times in <2h without solving. Likely bad bot.", count))
              .. "&timestamp=".. ngx.escape_uri(os.date("!%Y-%m-%dT%H:%M:%S+00:00"))
    local status, resp = http_post("https://api.abuseipdb.com/api/v2/report",
                                   body, "application/x-www-form-urlencoded",
                                   { ["Key"] = cfg.abuseipdb_key })
    if status == 200 or status == 201 then
        ngx.log(ngx.ERR, "captcha: reported ", ip, " to AbuseIPDB after ", count, " walk-aways")
    else
        ngx.log(ngx.WARN, "captcha: abuseipdb returned ", tostring(status), ": ", tostring(resp))
    end
end

local function note_walkaway(ip)
    if not ip or ip == "" or cfg.abuseipdb_key == "" then return end
    local counts = ngx.shared.ts_walkaway; if not counts then return end
    local n = counts:incr(ip, 1, 0, WALKAWAY_TTL)
    if not n then return end
    if n >= WALKAWAY_THRESHOLD then
        local reported = ngx.shared.ts_reported; if not reported then return end
        if reported:add(ip, 1, REPORT_COOLDOWN) then
            ngx.header["X-Turnstile"] = "reported"
            local ok, terr = ngx.timer.at(0, function(prem, ip_arg, n_arg)
                if prem then return end
                report_to_abuseipdb(ip_arg, n_arg)
            end, ip, n)
            if not ok then ngx.log(ngx.ERR, "ts timer: ", terr); reported:delete(ip) end
        end
    end
end

local function clear_walkaway(ip)
    if not ip or ip == "" then return end
    local counts = ngx.shared.ts_walkaway
    if counts then counts:delete(ip) end
end



-- ---------- challenge page (provider-dispatched widget block) ----------
local function widget_turnstile(redirect)
    return table.concat({
        '<script src="https://challenges.cloudflare.com/turnstile/v0/api.js" async defer></script>',
        '<form id="ts-form" action="', VERIFY_PATH, '" method="POST">',
        '<input type="hidden" name="redirect" value="', redirect, '">',
        '<div class="cf-turnstile" data-sitekey="', html_escape(cfg.site_key),
            '" data-callback="onTsSuccess" data-theme="auto"></div>',
        '</form>',
        '<script>function onTsSuccess(){document.getElementById("ts-form").submit();}</script>',
    })
end

local function widget_cap(redirect)
    local endpoint = html_escape(cfg.cap_url .. "/" .. cfg.cap_site_key .. "/")
    return table.concat({
        '<script src="https://cdn.jsdelivr.net/npm/cap-widget" async></script>',
        '<form id="ts-form" action="', VERIFY_PATH, '" method="POST">',
        '<input type="hidden" name="redirect" value="', redirect, '">',
        '<cap-widget id="cap" data-cap-api-endpoint="', endpoint, '"></cap-widget>',
        '</form>',
        '<script>',
        'document.addEventListener("DOMContentLoaded",function(){',
        '  var w=document.getElementById("cap");if(!w)return;',
        '  w.addEventListener("solve",function(){document.getElementById("ts-form").submit();});',
        '});',
        '</script>',
    })
end

local function serve_challenge()
    local redirect = html_escape(ngx.var.request_uri or "/")
    local widget   = (provider() == "cap") and widget_cap(redirect) or widget_turnstile(redirect)

    local html = table.concat({
        '<!doctype html><html lang="en"><head>',
        '<meta charset="utf-8">',
        '<meta name="viewport" content="width=device-width,initial-scale=1">',
        '<title>Checking your browser…</title>',
        '<style>html,body{height:100%;margin:0;font-family:system-ui,-apple-system,Segoe UI,Roboto,sans-serif;background:#f5f6f8;color:#222}',
        '.wrap{min-height:100%;display:flex;align-items:center;justify-content:center;padding:1rem}',
        '.card{background:#fff;border-radius:10px;box-shadow:0 4px 20px rgba(0,0,0,.08);padding:2rem 2.25rem;max-width:420px;text-align:center}',
        'h1{font-size:1.25rem;margin:0 0 .5rem}p{margin:.25rem 0 1.25rem;color:#555;font-size:.95rem}',
        '.cf-turnstile,cap-widget{display:flex;justify-content:center;margin:0 auto}</style></head><body>',
        '<div class="wrap"><div class="card">',
        '<h1>One quick check…</h1>',
        '<p>Please complete the security check to continue.</p>',
        widget,
        '</div></div></body></html>',
    })

    ngx.header["X-Turnstile"]   = "challenge"
    ngx.header["Content-Type"]  = "text/html; charset=utf-8"
    ngx.header["Cache-Control"] = "no-store"
    ngx.header["X-Robots-Tag"]  = "noindex"
    ngx.status = 200
    ngx.print(html)
    return ngx.exit(200)
end


-- ---------- public: gate ----------
function M.guard()
    if ngx.var.uri == VERIFY_PATH then return end
    if ip_whitelisted(ngx.var.remote_addr) then return end
    if cookie_is_valid(ngx.var["cookie_" .. COOKIE_NAME]) then return end

    -- (re-add any monitoring bypasses you set up earlier here)

    local accept = ngx.var.http_accept or ""
    if ngx.req.get_method() ~= "GET" or not accept:find("text/html", 1, true) then
        ngx.status = 403
        ngx.header["Content-Type"] = "text/plain"
        ngx.say("Forbidden: browser verification required")
        return ngx.exit(403)
    end

    note_walkaway(ngx.var.remote_addr)
    return serve_challenge()
end


-- ---------- public: verify endpoint ----------
function M.verify()
    if ngx.req.get_method() ~= "POST" then
        ngx.status = 405; ngx.header["Allow"] = "POST"; return ngx.exit(405)
    end

    ngx.req.read_body()
    local args = ngx.req.get_post_args() or {}
    local back = safe_redirect(args["redirect"])
    local prov = provider()

    local token = (prov == "cap") and args["cap-token"] or args["cf-turnstile-response"]
    if not token or token == "" then
        ngx.status = 400; ngx.say("Missing CAPTCHA token"); return ngx.exit(400)
    end

    local status, body
    if prov == "cap" then
        status, body = cap_verify(token)
    else
        status, body = turnstile_verify(token, ngx.var.remote_addr)
    end

    if not status or status ~= 200 then
        ngx.log(ngx.ERR, "captcha: ", prov, " siteverify error: ", tostring(body))
        ngx.status = 502; ngx.say("Verification service unavailable"); return ngx.exit(502)
    end

    local data = cjson.decode(body or "")
    if not data or data.success ~= true then
        ngx.log(ngx.WARN, "captcha: ", prov, " rejected: ", body or "")
        ngx.status = 403; ngx.say("Verification failed. Please try again."); return ngx.exit(403)
    end

    clear_walkaway(ngx.var.remote_addr)

    local expires = ngx.time() + COOKIE_TTL
    ngx.header["Set-Cookie"] = string.format(
        "%s=%s; Path=/; Expires=%s; Max-Age=%d; HttpOnly; Secure; SameSite=Lax",
        COOKIE_NAME, sign_cookie(expires), ngx.cookie_time(expires), COOKIE_TTL
    )
    ngx.header["Cache-Control"] = "no-store"
    ngx.header["X-Turnstile"]   = "verified"
    return ngx.redirect(back, 302)
end

return M