require "libcppjwt"
crypto = require 'crypto'

local Synchronize = require "Synchronize"
local https = Synchronize("luasec_mutex", function()
	return require ("ssl.https")
end)
local m_model = require "authfilter.model"

--A reasonable frequency to check for updates to the public keys used by Azure AD is every 24 hours (https://learn.microsoft.com/en-us/azure/active-directory/develop/access-tokens)
JWKS_KEYS_RETENTION_SECONDS = 86400

local function splitByChunk(text, chunkSize)
    local s = {}
    for i=1, #text, chunkSize do
        s[#s+1] = text:sub(i,i+chunkSize - 1)
    end
    return s
end

local function GetPubKeyFromCert(certStr)
	local cert = crypto.x509_cert and crypto.x509_cert()
	if cert then
		cert:from_pem(certStr)
		local pubkey = cert:pubkey()
		return pubkey and pubkey:to_pem()
	end
end

local function GetRedis ()
	local _, redis = m_model.GetDatasource()
	return redis
end

local function LoadVccKeys(vccId, jwks_uri)
	local responseBody = {}

	local result, code, headers, statusline = https.request{
		url = jwks_uri,
		method = "GET",
		sink = ltn12.sink.table(responseBody)
	}
	if code == 200 then
		local data = table.concat(responseBody, "")
		local redis = GetRedis()
		local redisKey = "OpenIdKeys:" ..vccId
		redis:set(redisKey, data)
		redis:expire(redisKey, JWKS_KEYS_RETENTION_SECONDS)
		return data and json.decode(data) or {}
	end
	return {}
end

local function GetVccKeys(vccId, jwks_uri)
	local redis = GetRedis()
	local redisKey = "OpenIdKeys:" ..vccId
	local openIdVccKeysStr = redis:get(redisKey)
	return openIdVccKeysStr and json.decode(openIdVccKeysStr) or LoadVccKeys(vccId, jwks_uri)
end

local function GetVccKey(vccId, jwks_uri, kid)
	local vccKeys = GetVccKeys(vccId, jwks_uri)
	for _, key in ipairs(vccKeys.keys or {}) do
		if key and key.kid == kid then
			return key
		end
	end
end

local M = {}

function M.Verify(vccId, jwks_uri, token, claims)
	local kid = cppjwtlib.GetKid(token)	
	local key = GetVccKey(vccId, jwks_uri, kid)	
	local x5c = key and key.x5c and type(key.x5c) == "table" and #key.x5c > 0 and key.x5c[1]
	if x5c then
		local cert_table = splitByChunk(x5c, 64)
		local cert = "-----BEGIN CERTIFICATE-----\n" .. table.concat(cert_table, "\n") .. "\n-----END CERTIFICATE-----"
		local pubkey = GetPubKeyFromCert(cert)
		if not pubkey then
			return false, "Invalid token"
		end
		return cppjwtlib.Verify(token, pubkey, json.encode(claims))
	end
	return false, "can not get x5c cert on kid key " .. kid
end

return M