local aes = require "resty.aes"
local md5 = require "resty.md5"
local str = require "resty.string"
local ipmatcher = require "resty.ipmatcher"
local kong_meta = require "kong.meta"
local redis = require "resty.redis"
local cjson_safe = require "cjson.safe"
cjson_safe.encode_empty_table_as_object(false)

local realm = "tidb.cloud"
local SECRET_AES_256_KEY = os.getenv("SECRET_AES_256_KEY")
local SECRET_AES_256_IV = os.getenv("SECRET_AES_256_IV")
local nonce_aes_128
local secret_aes_256
local REQUIRED_DIGEST_FIELDS = {"realm", "uri", "username", "nonce", "nc", "cnonce", "response", "qop"}

local DigestAuth = {PRIORITY = 1097, VERSION = kong_meta.version}

local function gen_error_response(code, message, error_format)
    if error_format == "openapi" then
        return cjson_safe.encode({code = code, message = message, details = {}})
    elseif error_format == "dataapi" then
        return cjson_safe.encode({
            type = "",
            data = {
                columns = {},
                rows = {},
                result = {
                    start_ms = 0,
                    end_ms = 0,
                    latency = "",
                    row_count = 0,
                    row_affect = 0,
                    limit = 0,
                    code = code,
                    message = message
                }
            }
        })
    else
        return ""
    end
end

local m = md5:new()
local function md5_digest(s)
    m:reset()
    m:update(s)
    local digest = m:final()
    return str.to_hex(digest)
end

local function hex_to_char(c)
    return string.char(tonumber(c, 16))
end

local function from_hex(s)
    if s ~= nil then
        s = string.gsub(s, "%x%x", hex_to_char)
    end
    return s
end

local function is_present(s)
    return s and s ~= "" and s ~= ngx.null
end

local sock_opts = {}
local function get_redis_connection(conf)
    local red = redis:new()
    red:set_timeout(conf.redis_timeout)

    sock_opts.ssl = conf.redis_ssl
    sock_opts.ssl_verify = conf.redis_ssl_verify

    -- use a special pool name only if redis_database is set to non-zero
    -- otherwise use the default pool name host:port
    if conf.redis_database ~= 0 then
        sock_opts.pool = string.format("%s:%d;%d", conf.redis_host,
                                       conf.redis_port, conf.redis_database)
    end

    local ok, err = red:connect(conf.redis_host, conf.redis_port, sock_opts)
    if not ok then
        return nil, string.format("failed to connect to Redis: %s", err)
    end

    local times, err = red:get_reused_times()
    if err then
        return nil, string.format("failed to get connect reused times: %s", err)
    end

    if times == 0 then
        if is_present(conf.redis_password) then
            local ok, err
            if is_present(conf.redis_username) then
                ok, err = red:auth(conf.redis_username, conf.redis_password)
            else
                ok, err = red:auth(conf.redis_password)
            end

            if not ok then
                return nil, string.format("failed to auth Redis: %s", err)
            end
        end

        if conf.redis_database ~= 0 then
            -- Only call select first time, since we know the connection is shared
            -- between instances that use the same redis database

            local ok, err = red:select(conf.redis_database)
            if not ok then
                return nil, string.format("failed to change Redis database: %s",
                                          err)
            end
        end
    end

    return red
end

local function generate_nonce()
    local now = math.ceil(ngx.now() * 1000)
    local ts_with_magic = string.format("M%sC", now)
    return str.to_hex(nonce_aes_128:encrypt(ts_with_magic))
end

local function extract_authorization(authorization)
    local auth_str = string.match(authorization, "%s*[Dd]igest%s*(.*)$")
    if not auth_str then
        return nil
    end

    local auth = {}
    -- Match any non-control character except double-quote to handle arbitrary URI characters
    for k, v in string.gmatch(auth_str, "(%w+)%s*=%s*\"([^%c\"]+)\"") do
        auth[k] = v
    end
    for k, v in string.gmatch(auth_str, "(%w+)%s*=%s*(%w+)") do
        auth[k] = v
    end
    return auth
end

local function challenge(response)
    local nonce = generate_nonce()
    local auth_header = 'Digest realm="' .. realm .. '", domain="", nonce="' ..
                            nonce .. '", algorithm=MD5, qop="auth", stale=false'
    kong.response
        .set_header("X-Debug-Trace-Id", kong.ctx.shared.trace_id or "-")
    kong.response.set_header("Access-Control-Allow-Headers",
                             "accept, content-type, authorization")
    kong.response.set_header("Access-Control-Allow-Methods",
                             "POST, GET, OPTIONS, PUT, DELETE, PATCH")
    kong.response.set_header("Access-Control-Allow-Origin", "*")
    kong.response.add_header("Access-Control-Expose-Headers",
                             "x-debug-trace-id,x-kong-proxy-latency,x-kong-upstream-latency,x-kong-response-latency,x-pingcap-requestid,x-ratelimit-remaining-minute,x-ratelimit-limit-minute,x-ratelimit-reset")
    kong.response.exit(401, response, {
        ["WWW-Authenticate"] = auth_header,
        ["Content-Type"] = "application/json"
    })
end

local function check_if_nonce_valid(nonce)
    local nonce_decoded = from_hex(nonce)
    if not nonce_decoded then
        return false, "can not decode nonce"
    end
    local ts_with_magic = nonce_aes_128:decrypt(nonce_decoded)
    if not ts_with_magic then
        return false, "can not decrypt nonce"
    end
    local ts = tonumber(string.match(ts_with_magic, "M(.*)C"))
    if not ts then
        return false, "can not transfer nonce to timestamp"
    end
    -- The max expiration nonce can tolerate is 5 seconds.
    if ngx.now() * 1000 - ts > 5000 then
        return false, "nonce is expired"
    end
    return true, ""
end

local function check_if_response_valid(auth, passwd)
    local A1 = string.format("%s:%s:%s", auth.username, auth.realm, passwd)
    local HA1 = md5_digest(A1)
    local A2 = string.format("%s:%s", kong.request.get_method(), auth.uri)
    local HA2 = md5_digest(A2)
    local expected_response = md5_digest(
                                  string.format("%s:%s:%s:%s:%s:%s", HA1,
                                                auth.nonce, auth.nc,
                                                auth.cnonce, auth.qop, HA2))
    return auth.response == expected_response
end

local function check_api_key_cidrs(cidrs, error_format)
    if not cidrs then
        return
    end
    if table.getn(cidrs) == 0 then
        return
    end
    local ip, err = ipmatcher.new(cidrs)
    if err then
        kong.log.err("new ip matcher failed: ", err, ", trace_id: ",
                     kong.ctx.shared.trace_id or "-")
        return
    end
    local real_ip = kong.client.get_forwarded_ip()
    local ok, err = ip:match(real_ip)
    if not ok then
        kong.response.set_header("X-Debug-Trace-Id",
                                 kong.ctx.shared.trace_id or "-")
        local resp = gen_error_response(49900006, string.format(
                                            "IP address %s is not allowed to access this resource.",
                                            real_ip), error_format)
        kong.response.exit(403, resp, {["Content-Type"] = "application/json"})
    end
end

local function reportUnauth(conf)
    if string.find(conf.redis_host,"^replica")~=nil then
        return
    end
    local red, err = get_redis_connection(conf)
    if not red then
        kong.log.err("reportUnauth failed to get redis connection",err)
        return
    end
    local path=kong.request.get_path()
    local host=kong.request.get_host()
    local res, err = red:hincrby("gateway_unauth_report",host.."/"..path,1)
    if err then
        kong.log.err("reportUnauth failed to exec hincrby redis",err)
        return
    end
    red:expire("gateway_unauth_report",60)
    red:set_keepalive(10000, 100)
end


local function validate(auth, conf)
    if not auth then
        return nil, "empty auth"
    end
    local required = REQUIRED_DIGEST_FIELDS
    for _, field in ipairs(required) do
        if auth[field] == nil or auth[field] == "" then
            return nil, string.format("missing required field: %s", field)
        end
    end
    if auth.realm ~= realm then
        return nil, string.format("invalid realm: %s", auth.realm)
    end
    if auth.uri ~= kong.request.get_path_with_query() then
        return nil, string.format("invalid uri: %s", auth.uri)
    end
    if auth.algorithm and string.upper(auth.algorithm) ~= "MD5" then
        return nil, string.format("invalid algorithm: %s", auth.algorithm)
    end
    local ok, msg = check_if_nonce_valid(auth.nonce)
    if not ok then
        return nil, string.format("invalid nonce: %s", msg)
    end

    local red, err = get_redis_connection(conf)
    if not red then
        if conf.fault_tolerant then
            kong.log.err(string.format(
                             "get redis connection failed, downgrade to fault tolerant mode, username:%s, nonce:%s, nc:%s, cnonce:%s, response:%s, err:%s, trace_id:%s",
                             auth.username, auth.nonce, auth.nc, auth.cnonce,
                             auth.response, err, kong.ctx.shared.trace_id or "-"))
            return {}, nil
        end
        return nil, string.format("failed to get redis connection: %s", err)
    end

    local user_secret, err = red:get(string.format("dbaas:account:api_keys:%s",
                                                   auth.username))
    if err then
        if conf.fault_tolerant then
            kong.log.err(string.format(
                             "exec redis get command failed, downgrade to fault tolerant mode, username:%s, nonce:%s, nc:%s, cnonce:%s, response:%s, err:%s, trace_id:%s",
                             auth.username, auth.nonce, auth.nc, auth.cnonce,
                             auth.response, err, kong.ctx.shared.trace_id or "-"))
            return {}, nil
        end
        return nil, string.format("failed to exec get redis: %s", err)
    end
    local ok, err = red:set_keepalive(10000, 100)
    if not ok then
        return nil, string.format("failed to set Redis keepalive: %s", err)
    end
    if user_secret == ngx.null then
        return nil,
               string.format(
                   "get empty user secret, result is null: username %s",
                   auth.username)
    end

    local secret_obj = cjson_safe.decode(user_secret)
    if not secret_obj then
        return nil, "can't decode redis user secret"
    end

    local expired_at = secret_obj["expired_ts"]
    if expired_at~=nil and expired_at>0 and ngx.now()*1000>expired_at then
        return nil, "KEY_EXPIRED"
    end
    if not secret_obj["private_key"] then
        return nil, "can't get user private_key"
    end
    local passwd = secret_aes_256:decrypt(from_hex(secret_obj["private_key"]))
    if not passwd then
        return nil, "can't decrypt passwd"
    end
    if not check_if_response_valid(auth, passwd) then
        return nil, string.format(
                   "response check failed: username:%s, nonce:%s, nc:%s, cnonce:%s, response:%s",
                   auth.username, auth.nonce, auth.nc, auth.cnonce,
                   auth.response)
    end
    if secret_obj["cidrs"] then
        check_api_key_cidrs(secret_obj["cidrs"], conf.error_format)
    end

    return secret_obj
end

function DigestAuth:access(conf)

    if kong.ctx.shared.has_authed then
        kong.log.info("skip digest auth since has authed, trace_id: ",
                      kong.ctx.shared.trace_id or "-")
        return
    end

    local authorization = kong.request.get_header("authorization")
    if not authorization and conf.allow_anonymous then
        -- Only allow anonymous when X-TIDBCLOUD-ZERO-INSTANCE-ID is present
        -- (zero-instance auth uses instance secret, no user credentials needed).
        -- For all other requests (e.g. X-TIDBCLOUD-CLUSTER-ID), issue a Digest
        -- challenge so that digest.Transport clients can complete the handshake.
        local zero_instance = kong.request.get_header("X-TIDBCLOUD-ZERO-INSTANCE-ID")
        if not is_present(zero_instance) then
            kong.log.info("anonymous not allowed without zero-instance header, issuing challenge, trace_id: ",
                          kong.ctx.shared.trace_id or "-")
            nonce_aes_128 = aes:new(conf.nonce_key, nil, aes.cipher(128, "cbc"),
                                    {iv = conf.nonce_iv})
            challenge("")
            return
        end

        kong.service.request.set_header("x-auth-type", "openapi")
        kong.service.request.set_header("x-auth-method", "anonymous")
        kong.service.request.clear_header("x-auth-username")
        kong.service.request.clear_header("x-auth-content")
        kong.service.request.clear_header("x-auth-raw")
        kong.ctx.shared.auth_content = nil
        kong.log.info("skip digest auth since anonymous access is allowed for zero-instance, trace_id: ",
                      kong.ctx.shared.trace_id or "-")
        return
    end

    nonce_aes_128 = aes:new(conf.nonce_key, nil, aes.cipher(128, "cbc"),
                            {iv = conf.nonce_iv})
    secret_aes_256 = aes:new(SECRET_AES_256_KEY, nil, aes.cipher(256, "cbc"),
                             {iv = SECRET_AES_256_IV})
    if not authorization then
        challenge("")
    end
    local auth = extract_authorization(authorization)
    local secret_obj, err = validate(auth, conf)
    if not secret_obj then
        if auth then
            kong.log.err("authenticate failed: username: ", auth.username,
                         ", err: ", err, ", trace_id: ",
                         kong.ctx.shared.trace_id or "-")
        else
            kong.log.err("authenticate failed: err: ", err, ", trace_id",
                         kong.ctx.shared.trace_id or "-")
        end
        local resp = ""

        if conf.error_format == "dataapi" then
            local msg="unauthorized"
            if err == "KEY_EXPIRED" then
                msg="API Key is no longer valid"
            end
            resp = gen_error_response(49900002, msg,
                                      conf.error_format)
        end
        reportUnauth(conf)
        challenge(resp)
    end
    if auth and secret_obj then
        local auth_content = {}
        auth_content["public_key"] = auth["username"]
        auth_content["cidrs"] = secret_obj["cidrs"]
        auth_content["rate_limit_rpm"] = secret_obj["rate_limit_rpm"]
        kong.service.request.set_header("x-auth-type", "openapi")
        kong.service.request.set_header("x-auth-method", "digest")
        kong.service.request.set_header("x-auth-username", auth["username"])
        kong.service.request.set_header("x-auth-content",
                                        cjson_safe.encode(auth_content))
        kong.service.request.clear_header("authorization")
        kong.ctx.shared.auth_content = auth_content
    end
end

return DigestAuth