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 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");

-- 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 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 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 scope == "openid" 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:add_timer(900, function()

	local k, code = codes:tail();

	while code and code_expired(code) do

		codes:set(k, nil);

		k, code = codes:tail();

	end

	return code and code_expires_in(code) + 1 or 900;

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 };

end

local function new_access_token(token_jid, role, scope, ttl, client, id_token)

	local token_data = {};

	if client then

		token_data.oauth2_client = client_subset(client);

	end

	if next(token_data) == nil then

		token_data = nil;

	end

	local token = tokens.create_jid_token(token_jid, token_jid, role, ttl, token_data, "oauth2");

	return {

		token_type = "bearer";

		access_token = token;

		expires_in = ttl;

		scope = scope;

		id_token = id_token;

		-- TODO: include refresh_token when implemented

	};

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 = {};

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, 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);

	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;

		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, nil, 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_secret(client_id) --> client_secret

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

end

local function verify_secret(client_id, client_secret)

	return hashes.equals(make_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_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

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

end

-- Used to issue/verify short-lived tokens for the authorization process below

local new_user_token, verify_user_token = jwt.init("HS256", random.bytes(32), nil, { default_ttl = 600 });

-- From the given request, figure out if the user is authenticated and has granted consent yet

-- As this requires multiple steps (seek credentials, seek consent), we have a lot of state to

-- carry around across requests. We also need to protect against CSRF and session mix-up attacks

-- (e.g. the user may have multiple concurrent flows in progress, session cookies aren't unique

--  to one of them).

-- Our strategy here is to preserve the original query string (containing the authz request), and

-- encode the rest of the flow in form POSTs.

local function get_auth_state(request)

	local form = request.method == "POST"

	         and request.body

	         and #request.body > 0

	         and request.headers.content_type == "application/x-www-form-urlencoded"

	         and http.formdecode(request.body);

	if not form then return {}; end

	if not form.user_token then

		-- First step: login

		local username = encodings.stringprep.nodeprep(form.username);

		local password = encodings.stringprep.saslprep(form.password);

		if not (username and password) or not usermanager.test_password(username, module.host, password) then

			return {

				error = "Invalid username/password";

			};

		end

		return {

			user = {

				username = username;

				host = module.host;

				token = new_user_token({ username = username, host = module.host });

			};

		};

	elseif form.user_token and form.consent then

		-- Second step: consent

		local ok, user = verify_user_token(form.user_token);