local hashes = require "util.hashes";

local cache = require "util.cache";

local http = require "util.http";

local jid = require "util.jid";

local json = require "util.json";

local usermanager = require "core.usermanager";

local errors = require "util.error";

local url = require "socket.url";

local id = require "util.id";

local encodings = require "util.encodings";

local base64 = encodings.base64;

local random = require "util.random";

local schema = require "util.jsonschema";

local set = require "util.set";

local jwt = require"util.jwt";

local it = require "util.iterators";

local array = require "util.array";

local st = require "util.stanza";

local function b64url(s)

	return (base64.encode(s):gsub("[+/=]", { ["+"] = "-", ["/"] = "_", ["="] = "" }))

end

local function read_file(base_path, fn, required)

	local f, err = io.open(base_path .. "/" .. fn);

	if not f then

		module:log(required and "error" or "debug", "Unable to load template file: %s", err);

		if required then

			return error("Failed to load templates");

		end

		return nil;

	end

	local data = assert(f:read("*a"));

	assert(f:close());

	return data;

end

local template_path = module:get_option_path("oauth2_template_path", "html");

local templates = {

	login = read_file(template_path, "login.html", true);

	consent = read_file(template_path, "consent.html", true);

	error = read_file(template_path, "error.html", true);

	css = read_file(template_path, "style.css");

	js = read_file(template_path, "script.js");

};

local site_name = module:get_option_string("site_name", module.host);

local _render_html = require"util.interpolation".new("%b{}", st.xml_escape);

local function render_page(template, data, sensitive)

	data = data or {};

	data.site_name = site_name;

	local resp = {

		status_code = 200;

		headers = {

			["Content-Type"] = "text/html; charset=utf-8";

			["Content-Security-Policy"] = "default-src 'self'";

			["X-Frame-Options"] = "DENY";

			["Cache-Control"] = (sensitive and "no-store" or "no-cache")..", private";

		};

		body = _render_html(template, data);

	};

	return resp;

end

local tokens = module:depends("tokenauth");

local default_access_ttl = module:get_option_number("oauth2_access_token_ttl", 86400);

local default_refresh_ttl = module:get_option_number("oauth2_refresh_token_ttl", nil);

-- Used to derive client_secret from client_id, set to enable stateless dynamic registration.

local registration_key = module:get_option_string("oauth2_registration_key");

local registration_algo = module:get_option_string("oauth2_registration_algorithm", "HS256");

local registration_options = module:get_option("oauth2_registration_options", { default_ttl = 60 * 60 * 24 * 90 });

local pkce_required = module:get_option_boolean("oauth2_require_code_challenge", false);

local verification_key;

local jwt_sign, jwt_verify;

if registration_key then

	-- Tie it to the host if global

	verification_key = hashes.hmac_sha256(registration_key, module.host);

	jwt_sign, jwt_verify = jwt.init(registration_algo, registration_key, registration_key, registration_options);

end

local function parse_scopes(scope_string)

	return array(scope_string:gmatch("%S+"));

end

local openid_claims = set.new({ "openid", "profile"; "email"; "address"; "phone" });

local function filter_scopes(username, requested_scope_string)

	local selected_role, granted_scopes = nil, array();

	if requested_scope_string then -- Specific role(s) requested

		local requested_scopes = parse_scopes(requested_scope_string);

		for _, scope in ipairs(requested_scopes) do

			if openid_claims:contains(scope) then

				granted_scopes:push(scope);

			end

			if selected_role == nil and usermanager.user_can_assume_role(username, module.host, scope) then

				selected_role = scope;

			end

		end

	end

	if not selected_role then

		-- By default use the users' default role

		selected_role = usermanager.get_user_role(username, module.host).name;

	end

	granted_scopes:push(selected_role);

	return granted_scopes:concat(" "), selected_role;

end

local function code_expires_in(code) --> number, seconds until code expires

	return os.difftime(code.expires, os.time());

end

local function code_expired(code) --> boolean, true: has expired, false: still valid

	return code_expires_in(code) < 0;

end

local codes = cache.new(10000, function (_, code)

	return code_expired(code)

end);

-- Periodically clear out unredeemed codes.  Does not need to be exact, expired

-- codes are rejected if tried. Mostly just to keep memory usage in check.

module:hourly("Clear expired authorization codes", function()

	local k, code = codes:tail();

	while code and code_expired(code) do

		codes:set(k, nil);

		k, code = codes:tail();

	end

end)

local function get_issuer()

	return (module:http_url(nil, "/"):gsub("/$", ""));

end

local loopbacks = set.new({ "localhost", "127.0.0.1", "::1" });

local function is_secure_redirect(uri)

	local u = url.parse(uri);

	return u.scheme ~= "http" or loopbacks:contains(u.host);

end

local function oauth_error(err_name, err_desc)

	return errors.new({

		type = "modify";

		condition = "bad-request";

		code = err_name == "invalid_client" and 401 or 400;

		text = err_desc and (err_name..": "..err_desc) or err_name;

		extra = { oauth2_response = { error = err_name, error_description = err_desc } };

	});

end

-- client_id / client_metadata are pretty large, filter out a subset of

-- properties that are deemed useful e.g. in case tokens issued to a certain

-- client needs to be revoked

local function client_subset(client)

	return { name = client.client_name; uri = client.client_uri; id = client.software_id; version = client.software_version };

end

local function new_access_token(token_jid, role, scope_string, client, id_token, refresh_token_info)

	local token_data = { oauth2_scopes = scope_string, oauth2_client = nil };

	if client then

		token_data.oauth2_client = client_subset(client);

	end

	if next(token_data) == nil then

		token_data = nil;

	end

	local refresh_token;

	local grant = refresh_token_info and refresh_token_info.grant;

	if not grant then

		-- No existing grant, create one

		grant = tokens.create_grant(token_jid, token_jid, default_refresh_ttl, token_data);

		-- Create refresh token for the grant if desired

		refresh_token = refresh_token_info ~= false and tokens.create_token(token_jid, grant, nil, nil, "oauth2-refresh");

	else

		-- Grant exists, reuse existing refresh token

		refresh_token = refresh_token_info.token;

		refresh_token_info.grant = nil; -- Prevent reference loop

	end

	local access_token, access_token_info = tokens.create_token(token_jid, grant, role, default_access_ttl, "oauth2");

	local expires_at = access_token_info.expires;

	return {

		token_type = "bearer";

		access_token = access_token;

		expires_in = expires_at and (expires_at - os.time()) or nil;

		scope = scope_string;

		id_token = id_token;

		refresh_token = refresh_token or nil;

	};

end

local function get_redirect_uri(client, query_redirect_uri) -- record client, string : string

	if not query_redirect_uri then

		if #client.redirect_uris ~= 1 then

			-- Client registered multiple URIs, it needs specify which one to use

			return;

		end

		-- When only a single URI is registered, that's the default

		return client.redirect_uris[1];

	end

	-- Verify the client-provided URI matches one previously registered

	for _, redirect_uri in ipairs(client.redirect_uris) do

		if query_redirect_uri == redirect_uri then

			return redirect_uri

		end

	end

end

local grant_type_handlers = {};

local response_type_handlers = {};

local verifier_transforms = {};

function grant_type_handlers.password(params)

	local request_jid = assert(params.username, oauth_error("invalid_request", "missing 'username' (JID)"));

	local request_password = assert(params.password, oauth_error("invalid_request", "missing 'password'"));

	local request_username, request_host, request_resource = jid.prepped_split(request_jid);

	if not (request_username and request_host) or request_host ~= module.host then

		return oauth_error("invalid_request", "invalid JID");

	end

	if not usermanager.test_password(request_username, request_host, request_password) then

		return oauth_error("invalid_grant", "incorrect credentials");

	end

	local granted_jid = jid.join(request_username, request_host, request_resource);

	local granted_scopes, granted_role = filter_scopes(request_username, params.scope);

	return json.encode(new_access_token(granted_jid, granted_role, granted_scopes, nil));

end

function response_type_handlers.code(client, params, granted_jid, id_token)

	local request_username, request_host = jid.split(granted_jid);

	if not request_host or request_host ~= module.host then

		return oauth_error("invalid_request", "invalid JID");

	end

	local granted_scopes, granted_role = filter_scopes(request_username, params.scope);

	if pkce_required and not params.code_challenge then

		return oauth_error("invalid_request", "PKCE required");

	end

	local code = id.medium();

	local ok = codes:set(params.client_id .. "#" .. code, {

		expires = os.time() + 600;

		granted_jid = granted_jid;

		granted_scopes = granted_scopes;

		granted_role = granted_role;

		challenge = params.code_challenge;

		challenge_method = params.code_challenge_method;

		id_token = id_token;

	});

	if not ok then

		return {status_code = 429};

	end

	local redirect_uri = get_redirect_uri(client, params.redirect_uri);

	if redirect_uri == "urn:ietf:wg:oauth:2.0:oob" then

		-- TODO some nicer template page

		-- mod_http_errors will set content-type to text/html if it catches this

		-- event, if not text/plain is kept for the fallback text.

		local response = { status_code = 200; headers = { content_type = "text/plain" } }

		response.body = module:context("*"):fire_event("http-message", {

			response = response;

			title = "Your authorization code";

			message = "Here's your authorization code, copy and paste it into " .. (client.client_name or "your client");

			extra = code;

		}) or ("Here's your authorization code:\n%s\n"):format(code);

		return response;

	elseif not redirect_uri then

		return 400;

	end

	local redirect = url.parse(redirect_uri);

	local query = http.formdecode(redirect.query or "");

	if type(query) ~= "table" then query = {}; end

	table.insert(query, { name = "code", value = code });

	table.insert(query, { name = "iss", value = get_issuer() });

	if params.state then

		table.insert(query, { name = "state", value = params.state });

	end

	redirect.query = http.formencode(query);

	return {

		status_code = 303;

		headers = {

			location = url.build(redirect);

		};

	}

end

-- Implicit flow

function response_type_handlers.token(client, params, granted_jid)

	local request_username, request_host = jid.split(granted_jid);

	if not request_host or request_host ~= module.host then

		return oauth_error("invalid_request", "invalid JID");

	end

	local granted_scopes, granted_role = filter_scopes(request_username, params.scope);

	local token_info = new_access_token(granted_jid, granted_role, granted_scopes, client, nil);

	local redirect = url.parse(get_redirect_uri(client, params.redirect_uri));

	token_info.state = params.state;

	redirect.fragment = http.formencode(token_info);

	return {

		status_code = 303;

		headers = {

			location = url.build(redirect);

		};

	}

end

local function make_client_secret(client_id) --> client_secret

	return hashes.hmac_sha256(verification_key, client_id, true);

end

local function verify_client_secret(client_id, client_secret)

	return hashes.equals(make_client_secret(client_id), client_secret);

end

function grant_type_handlers.authorization_code(params)

	if not params.client_id then return oauth_error("invalid_request", "missing 'client_id'"); end

	if not params.client_secret then return oauth_error("invalid_request", "missing 'client_secret'"); end

	if not params.code then return oauth_error("invalid_request", "missing 'code'"); end

	if params.scope and params.scope ~= "" then

		return oauth_error("invalid_scope", "unknown scope requested");

	end

	local client_ok, client = jwt_verify(params.client_id);

	if not client_ok then

		return oauth_error("invalid_client", "incorrect credentials");

	end

	if not verify_client_secret(params.client_id, params.client_secret) then

		module:log("debug", "client_secret mismatch");

		return oauth_error("invalid_client", "incorrect credentials");

	end

	local code, err = codes:get(params.client_id .. "#" .. params.code);

	if err then error(err); end

	-- MUST NOT use the authorization code more than once, so remove it to

	-- prevent a second attempted use

	codes:set(params.client_id .. "#" .. params.code, nil);

	if not code or type(code) ~= "table" or code_expired(code) then

		module:log("debug", "authorization_code invalid or expired: %q", code);

		return oauth_error("invalid_client", "incorrect credentials");

	end

	-- TODO Decide if the code should be removed or not when PKCE fails

	local transform = verifier_transforms[code.challenge_method or "plain"];

	if not transform then

		return oauth_error("invalid_request", "unknown challenge transform method");

	elseif transform(params.code_verifier) ~= code.challenge then