local helpers = require "spec.helpers"
local md5 = require "resty.md5"
local str = require "resty.string"

local REDIS_HOST = helpers.redis_host
local REDIS_PORT = helpers.redis_port
local REDIS_PASSWORD = ""
local REDIS_DATABASE = 1
local REDIS_DATABASE_RATE_LIMIT = 2

-- This performs the test up to two times (and no more than two).
-- We are **not** retrying to "give it another shot" in case of a flaky test.
-- The reason why we allow for a single retry in this test suite is because
-- tests are dependent on the value of the current minute. If the minute
-- flips during the test (i.e. going from 03:43:59 to 03:44:00), the result
-- will fail. Since each test takes less than a minute to run, running it
-- a second time right after that failure ensures that another flip will
-- not occur. If the second execution failed as well, this means that there
-- was an actual problem detected by the test.
local function it_with_retry(desc, test)
    return it(desc, function(...)
        if not pcall(test, ...) then
            ngx.sleep(61 - (ngx.now() % 60)) -- Wait for minute to expire
            test(...)
        end
    end)
end

local function init_dependency()
    local redis = require "resty.redis"
    local red = redis:new()
    red:set_timeout(2000)
    local ok, err = red:connect(REDIS_HOST, REDIS_PORT)
    if not ok then
        error("failed to connect to Redis: " .. err)
    end

    if REDIS_PASSWORD and REDIS_PASSWORD ~= "" then
        local ok, err = red:auth(REDIS_PASSWORD)
        if not ok then
            error("failed to connect to Redis: " .. err)
        end
    end

    local ok, err = red:select(REDIS_DATABASE)
    if not ok then
        error("failed to change Redis database: " .. err)
    end

    local ok, err = red:set("dbaas:account:api_keys:abc",
                            "{\"private_key\":\"895e6c10bad41ad3adc353baa50138cd\", \"cidrs\":[\"0.0.0.0/0\"], \"rate_limit_rpm\":10}")
    if not ok then
        error("failed to set `abc` api keys: " .. err)
    end

    local ok, err = red:set("dbaas:account:api_keys:digest",
                            "{\"private_key\":\"895e6c10bad41ad3adc353baa50138cd\", \"cidrs\":[\"0.0.0.0/0\"], \"rate_limit_rpm\":10}")
    if not ok then
        error("failed to set `digest` api keys: " .. err)
    end

    local ok, err = red:set("dbaas:account:api_keys:basic",
                            "{\"private_key\":\"895e6c10bad41ad3adc353baa50138cd\", \"cidrs\":[\"0.0.0.0/0\"], \"rate_limit_rpm\":10}")
    if not ok then
        error("failed to set `basic` api keys: " .. err)
    end

    helpers.setenv("SECRET_AES_256_KEY", "12345678901234561234567890123456")
    helpers.setenv("SECRET_AES_256_IV", "1234567890123456")

    red:close()
end

local function extract_www_authenticate(res)
    local auth_str = string.match(res.headers["WWW-Authenticate"],
                                  "%s*[Dd]igest%s*(.*)$")
    if not auth_str then
        error("get bad WWW-Authenticate header")
    end

    local auth = {}
    for k, v in string.gmatch(auth_str, "(%w+)%s*=%s*\"?([%w./=:_?&-]+)\"?") do
        auth[k] = v
    end
    return auth
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 gen_response(username, passwd, realm, method, uri, nonce, nc,
                            cnonce, qop)
    local A1 = string.format("%s:%s:%s", username, realm, passwd)
    local HA1 = md5_digest(A1)
    local A2 = string.format("%s:%s", method, uri)
    local HA2 = md5_digest(A2)
    local response = md5_digest(string.format("%s:%s:%s:%s:%s:%s", HA1, nonce,
                                              nc, cnonce, qop, HA2))
    return response
end

-- Digest username="apjcejac", realm="tidb.cloud", nonce="4tfOBX/XnfNQ7XAuJ8WTuCEhR2t1SpwW", uri="/status/200", cnonce="YzI2ZTYzNTFiMDk0MzFlMDQxOGNjMDEwMTYyNjRkNmE=", nc=00000001, qop=auth, response="53c2f533bc405ee75dff2c67431b518e", algorithm=MD5

for _, strategy in helpers.each_strategy() do
    describe("Plugin: digest-auth [#" .. strategy .. "]", function()
        local proxy_client
        lazy_setup(function()
            init_dependency()

            local bp = helpers.get_db_utils(strategy)

            local service1 = bp.services:insert{
                protocol = "http",
                host = helpers.mock_upstream_host,
                port = helpers.mock_upstream_port
            }

            local route1 = bp.routes:insert{
                protocols = {"http", "https"},
                hosts = {"test1.com"},
                service = service1
            }

            local route2 = bp.routes:insert{
                protocols = {"http", "https"},
                hosts = {"test2.com"},
                service = service1
            }

            bp.plugins:insert{
                route = {id = route1.id},
                name = "digest-auth",
                config = {
                    nonce_key = "1234567890123456",
                    nonce_iv = "1234567890123456",
                    redis_host = REDIS_HOST,
                    redis_database = REDIS_DATABASE
                }
            }
            bp.plugins:insert{
                route = {id = route1.id},
                name = "rate-limit",
                config = {
                    header_name = "x-auth-username",
                    minute = 6,
                    fault_tolerant = false,
                    redis_host = REDIS_HOST,
                    redis_port = REDIS_PORT,
                    redis_database = REDIS_DATABASE_RATE_LIMIT,
                    error_format = "openapi"
                }
            }

            bp.plugins:insert{
                route = {id = route2.id},
                name = "custom-basic-auth",
                config = {
                    nonce_key = "1234567890123456",
                    nonce_iv = "1234567890123456",
                    redis_host = REDIS_HOST,
                    redis_database = REDIS_DATABASE
                }
            }
            bp.plugins:insert{
                route = {id = route2.id},
                name = "digest-auth",
                config = {
                    nonce_key = "1234567890123456",
                    nonce_iv = "1234567890123456",
                    redis_host = REDIS_HOST,
                    redis_database = REDIS_DATABASE
                }
            }
            bp.plugins:insert{
                route = {id = route2.id},
                name = "rate-limit",
                config = {
                    header_name = "x-auth-username",
                    minute = 6,
                    fault_tolerant = false,
                    redis_host = REDIS_HOST,
                    redis_port = REDIS_PORT,
                    redis_database = REDIS_DATABASE_RATE_LIMIT,
                    error_format = "dataapi"
                }
            }

            assert(helpers.start_kong({
                database = strategy,
                nginx_conf = "spec/fixtures/custom_nginx.template",
                plugins = "rate-limit,digest-auth,custom-basic-auth"
            }))
        end)

        lazy_teardown(function()
            helpers.stop_kong()
        end)

        before_each(function()
            proxy_client = helpers.proxy_client()
        end)

        after_each(function()
            if proxy_client then
                proxy_client:close()
            end
        end)

        describe("Digest Auth With Rate Limit", function()
            it_with_retry("exceed requests return 429", function()
                for i = 1, 10 do
                    local res = assert(proxy_client:send({
                        method = "GET",
                        path = "/status/200",
                        headers = {["Host"] = "test1.com"}
                    }))
                    assert.res_status(401, res)
                    local auth = extract_www_authenticate(res)
                    local response = gen_response("abc", "123456", auth.realm,
                                                  "GET", "/status/200",
                                                  auth.nonce, "00000001",
                                                  "YzI2ZTYzNTFiMDk0MzFlMDQxOGNjMDEwMTYyNjRkNmE=",
                                                  "auth")
                    local header =
                        "Digest username=\"abc\", realm=\"tidb.cloud\", uri=\"/status/200\", cnonce=\"YzI2ZTYzNTFiMDk0MzFlMDQxOGNjMDEwMTYyNjRkNmE=\", nc=00000001, qop=auth, algorithm=MD5, " ..
                            "nonce=" .. auth.nonce .. ", response=" .. response

                    local res200 = assert(proxy_client:send({
                        method = "GET",
                        path = "/status/200",
                        headers = {
                            ["Host"] = "test1.com",
                            ["Authorization"] = header
                        }
                    }))
                    assert.res_status(200, res200)

                    helpers.wait_timer("rate-limiting", true, "any-finish")
                end

                for i = 1, 10 do
                    local res = assert(proxy_client:send({
                        method = "GET",
                        path = "/status/200",
                        headers = {["Host"] = "test1.com"}
                    }))
                    assert.res_status(401, res)
                    local auth = extract_www_authenticate(res)
                    local response = gen_response("abc", "123456", auth.realm,
                                                  "GET", "/status/200",
                                                  auth.nonce, "00000001",
                                                  "YzI2ZTYzNTFiMDk0MzFlMDQxOGNjMDEwMTYyNjRkNmE=",
                                                  "auth")
                    local header =
                        "Digest username=\"abc\", realm=\"tidb.cloud\", uri=\"/status/200\", cnonce=\"YzI2ZTYzNTFiMDk0MzFlMDQxOGNjMDEwMTYyNjRkNmE=\", nc=00000001, qop=auth, algorithm=MD5, " ..
                            "nonce=" .. auth.nonce .. ", response=" .. response

                    local res429 = assert(proxy_client:send({
                        method = "GET",
                        path = "/status/200",
                        headers = {
                            ["Host"] = "test1.com",
                            ["Authorization"] = header
                        }
                    }))
                    assert.res_status(429, res429)

                    helpers.wait_timer("rate-limiting", true, "any-finish")
                end
            end)
        end)
        describe("Digest Auth With Custom Rate Limit", function()
            it_with_retry("exceed requests return 429", function()
                for i = 1, 10 do
                    local res = assert(proxy_client:send({
                        method = "GET",
                        path = "/status/200",
                        headers = {["Host"] = "test2.com"}
                    }))
                    assert.res_status(401, res)
                    local auth = extract_www_authenticate(res)
                    local response = gen_response("digest", "123456",
                                                  auth.realm, "GET",
                                                  "/status/200", auth.nonce,
                                                  "00000001",
                                                  "YzI2ZTYzNTFiMDk0MzFlMDQxOGNjMDEwMTYyNjRkNmE=",
                                                  "auth")
                    local header =
                        "Digest username=\"digest\", realm=\"tidb.cloud\", uri=\"/status/200\", cnonce=\"YzI2ZTYzNTFiMDk0MzFlMDQxOGNjMDEwMTYyNjRkNmE=\", nc=00000001, qop=auth, algorithm=MD5, " ..
                            "nonce=" .. auth.nonce .. ", response=" .. response

                    local res200 = assert(proxy_client:send({
                        method = "GET",
                        path = "/status/200",
                        headers = {
                            ["Host"] = "test2.com",
                            ["Authorization"] = header
                        }
                    }))
                    assert.res_status(200, res200)

                    helpers.wait_timer("rate-limiting", true, "any-finish")
                end

                for i = 1, 6 do
                    local res = assert(proxy_client:send({
                        method = "GET",
                        path = "/status/200",
                        headers = {["Host"] = "test2.com"}
                    }))
                    assert.res_status(401, res)
                    local auth = extract_www_authenticate(res)
                    local response = gen_response("digest", "123456",
                                                  auth.realm, "GET",
                                                  "/status/200", auth.nonce,
                                                  "00000001",
                                                  "YzI2ZTYzNTFiMDk0MzFlMDQxOGNjMDEwMTYyNjRkNmE=",
                                                  "auth")
                    local header =
                        "Digest username=\"digest\", realm=\"tidb.cloud\", uri=\"/status/200\", cnonce=\"YzI2ZTYzNTFiMDk0MzFlMDQxOGNjMDEwMTYyNjRkNmE=\", nc=00000001, qop=auth, algorithm=MD5, " ..
                            "nonce=" .. auth.nonce .. ", response=" .. response

                    local res429 = assert(proxy_client:send({
                        method = "GET",
                        path = "/status/200",
                        headers = {
                            ["Host"] = "test2.com",
                            ["Authorization"] = header
                        }
                    }))
                    assert.res_status(429, res429)

                    helpers.wait_timer("rate-limiting", true, "any-finish")
                end
            end)
        end)
        describe("Basic Auth With Custom Rate Limit", function()
            it_with_retry("exceed requests return 429", function()
                for i = 1, 10 do
                    local res200 = assert(proxy_client:send({
                        method = "GET",
                        path = "/status/200",
                        headers = {
                            ["Host"] = "test2.com",
                            ["Authorization"] = "Basic YmFzaWM6MTIzNDU2Cg=="
                        }
                    }))
                    assert.res_status(200, res200)

                    helpers.wait_timer("rate-limiting", true, "any-finish")
                end

                for i = 1, 6 do
                    local res429 = assert(proxy_client:send({
                        method = "GET",
                        path = "/status/200",
                        headers = {
                            ["Host"] = "test2.com",
                            ["Authorization"] = "Basic YmFzaWM6MTIzNDU2Cg=="
                        }
                    }))
                    assert.res_status(429, res429)

                    helpers.wait_timer("rate-limiting", true, "any-finish")
                end
            end)
        end)
    end)
end