-- Load the module
local ffi = require 'ffi'
local kres = require('kres')
local C = ffi.C

local trust_anchors -- the public pseudo-module, exported as global variable

-- Fetch over HTTPS with peert cert checked
local function https_fetch(url, ca)
	local ssl_ok, https = pcall(require, 'ssl.https')
	local ltn_ok, ltn12 = pcall(require, 'ltn12')
	if not ssl_ok or not ltn_ok then
		return nil, 'luasec and luasocket needed for root TA bootstrap'
	end
	local resp = {}
	local r, c = https.request{
	       url = url,
	       cafile = ca,
	       verify = {'peer', 'fail_if_no_peer_cert' },
	       protocol = 'tlsv1_2',
	       sink = ltn12.sink.table(resp),
	}
	if r == nil then return r, c end
	return resp[1]
end

-- remove UTC timezone specification if present or throw error
local function time2utc(orig_timespec)
	local patterns = {'[+-]00:00$', 'Z$'}
	for _, pattern in ipairs(patterns) do
		local timespec, removals = string.gsub(orig_timespec, pattern, '')
		if removals == 1 then
			return timespec
		end
	end
	error(string.format('unsupported time specification: %s', orig_timespec))
end

local function keydigest_is_valid(valid_from, valid_until)
	local format =            '%Y-%m-%dT%H:%M:%S'
	local time_now = os.date('!%Y-%m-%dT%H:%M:%S')  -- ! forces UTC
	local time_diff = ffi.new('double[1]')
	local err = ffi.C.kr_strptime_diff(
		format, time_now, time2utc(valid_from), time_diff)
	if (err ~= nil) then
	       error(string.format('failed to process "validFrom" constraint: %s',
				   ffi.string(err)))
	end
	local from_ok = time_diff[0] > 0

	-- optional attribute
	local until_ok = true
	if valid_until then
		err = ffi.C.kr_strptime_diff(
			format, time_now, time2utc(valid_until), time_diff)
		if (err ~= nil) then
			error(string.format('failed to process "validUntil" constraint: %s',
					    ffi.string(err)))
		end
		until_ok = time_diff[0] < 0
	end
	return from_ok and until_ok
end

local function parse_xml_keydigest(attrs, inside, output)
	local fields = {}
	local _, n = string.gsub(attrs, "([%w]+)=\"([^\"]*)\"", function (k, v) fields[k] = v end)
	assert(n >= 1,
		string.format('cannot parse XML attributes from "%s"', attrs))
	assert(fields['validFrom'],
		string.format('mandatory KeyDigest XML attribute validFrom ' ..
		'not found in "%s"', attrs))
	local valid_attrs = {id = true, validFrom = true, validUntil = true}
	for key, _ in pairs(fields) do
		assert(valid_attrs[key],
		       string.format('unsupported KeyDigest attribute "%s" found in "%s"',
				     key, attrs))
	end

	_, n = string.gsub(inside, "<([%w]+).->([^<]+)</[%w]+>", function (k, v) fields[k] = v end)
	assert(n >= 1,
		string.format('error parsing KeyDigest XML elements from "%s"',
			      inside))
	local mandatory_elements = {'KeyTag', 'Algorithm', 'DigestType', 'Digest'}
	for _, key in ipairs(mandatory_elements) do
		assert(fields[key],
			string.format('mandatory element %s is missing in "%s"',
				      key, inside))
	end
	assert(n == 4, string.format('found %d elements but expected 4 in %s', n, inside))
	table.insert(output, fields)  -- append to list of parsed keydigests
end

local function generate_ds(keydigests)
	local rrset = ''
	for _, fields in ipairs(keydigests) do
		local rr = string.format(
			'. 0 IN DS %s %s %s %s',
			fields.KeyTag, fields.Algorithm, fields.DigestType, fields.Digest)
		if keydigest_is_valid(fields['validFrom'], fields['validUntil']) then
			rrset = rrset .. '\n' .. rr
		else
			log('[ ta ] skipping trust anchor "%s" ' ..
			    'because it is outside of validity range', rr)
		end
	end
	return rrset
end

local function assert_str_match(str, pattern, expected)
	local count = 0
	for _ in string.gmatch(str, pattern) do
		count = count + 1
	end
	assert(count == expected,
	       string.format('expected %d occurences of "%s" but got %d in "%s"',
			     expected, pattern, count, str))
end

-- Fetch root anchors in XML over HTTPS, returning a zone-file-style string
-- or false in case of error, and a message.
local function bootstrap(url, ca)
	-- RFC 7958, sec. 2, but we don't do precise XML parsing.
	-- @todo ICANN certificate is verified against current CA
	--       this is not ideal, as it should rather verify .xml signature which
	--       is signed by ICANN long-lived cert, but luasec has no PKCS7
	local xml, err = https_fetch(url, ca)
	if not xml then
		return false, string.format('[ ta ] fetch of "%s" failed: %s', url, err)
	end

	-- we support only minimal subset of https://tools.ietf.org/html/rfc7958
	assert_str_match(xml, '<?xml version="1%.0" encoding="UTF%-8"%?>', 1)
	assert_str_match(xml, '<TrustAnchor ', 1)
	assert_str_match(xml, '<Zone>.</Zone>', 1)
	assert_str_match(xml, '</TrustAnchor>', 1)

	-- Parse root trust anchor, one digest at a time, converting to a zone-file-style string.
	local keydigests = {}
	string.gsub(xml, "<KeyDigest([^>]*)>(.-)</KeyDigest>", function(attrs, inside)
		parse_xml_keydigest(attrs, inside, keydigests)
	end)
	local rrset = generate_ds(keydigests)
	if rrset == '' then
		return false, string.format('[ ta ] no valid trust anchors found at "%s"', url)
	end
	local msg = '[ ta ] Root trust anchors bootstrapped over https with pinned certificate.\n'
			 .. '       You SHOULD verify them manually against original source:\n'
			 .. '       https://www.iana.org/dnssec/files\n'
			 .. '[ ta ] Current root trust anchors are:'
			 .. rrset
	return rrset, msg
end

-- RFC5011 state table
local key_state = {
	Start = 'Start', AddPend = 'AddPend', Valid = 'Valid',
	Missing = 'Missing', Revoked = 'Revoked', Removed = 'Removed'
}

-- Find key in current keyset
local function ta_find(keyset, rr)
	local rr_tag = C.kr_dnssec_key_tag(rr.type, rr.rdata, #rr.rdata)
	assert(rr_tag >= 0 and rr_tag <= 65535, string.format('invalid RR: %s: %s',
	       kres.rr2str(rr), ffi.string(C.knot_strerror(rr_tag))))
	for i, ta in ipairs(keyset) do
		-- Match key owner and content
		local ta_tag = C.kr_dnssec_key_tag(ta.type, ta.rdata, #ta.rdata)
		assert(ta_tag >= 0 and ta_tag <= 65535, string.format('invalid RR: %s: %s',
		       kres.rr2str(ta), ffi.string(C.knot_strerror(ta_tag))))
		if ta.owner == rr.owner then
			if ta.type == rr.type then
				if rr.type == kres.type.DNSKEY then
					if C.kr_dnssec_key_match(ta.rdata, #ta.rdata, rr.rdata, #rr.rdata) == 0 then
						return ta
					end
				elseif rr.type == kres.type.DS and ta.rdata == rr.rdata then
					return ta
				end
			-- DNSKEY superseding DS, inexact match
			elseif rr.type == kres.type.DNSKEY and ta.type == kres.type.DS then
				if ta.key_tag == rr_tag then
					keyset[i] = rr -- Replace current DS
					rr.state = ta.state
					rr.key_tag = ta.key_tag
					return rr
				end
			-- DS key matching DNSKEY, inexact match
			elseif rr.type == kres.type.DS and ta.type == kres.type.DNSKEY then
				if rr_tag == ta_tag then
					return ta
				end
			end
		end
	end
	return nil
end

-- Evaluate TA status of a RR according to RFC5011.  The time is in seconds.
local function ta_present(keyset, rr, hold_down_time, force_valid)
	if rr.type == kres.type.DNSKEY and not C.kr_dnssec_key_ksk(rr.rdata) then
		return false -- Ignore
	end
	-- Find the key in current key set and check its status
	local now = os.time()
	local key_revoked = (rr.type == kres.type.DNSKEY) and C.kr_dnssec_key_revoked(rr.rdata)
	local key_tag = C.kr_dnssec_key_tag(rr.type, rr.rdata, #rr.rdata)
	assert(key_tag >= 0 and key_tag <= 65535, string.format('invalid RR: %s: %s',
	       kres.rr2str(rr), ffi.string(C.knot_strerror(key_tag))))
	local ta = ta_find(keyset, rr)
	if ta then
		-- Key reappears (KeyPres)
		if ta.state == key_state.Missing then
			ta.state = key_state.Valid
			ta.timer = nil
		end
		-- Key is revoked (RevBit)
		if ta.state == key_state.Valid or ta.state == key_state.Missing then
			if key_revoked then
				ta.state = key_state.Revoked
				ta.timer = now + hold_down_time
			end
		end
		-- Remove hold-down timer expires (RemTime)
		if ta.state == key_state.Revoked and os.difftime(ta.timer, now) <= 0 then
			ta.state = key_state.Removed
			ta.timer = nil
		end
		-- Add hold-down timer expires (AddTime)
		if ta.state == key_state.AddPend and os.difftime(ta.timer, now) <= 0 then
			ta.state = key_state.Valid
			ta.timer = nil
		end
		if rr.state ~= key_state.Valid or verbose() then
			log('[ ta ] key: ' .. key_tag .. ' state: '..ta.state)
		end
		return true
	elseif not key_revoked then -- First time seen (NewKey)
		rr.key_tag = key_tag
		if force_valid then
			rr.state = key_state.Valid
		else
			rr.state = key_state.AddPend
			rr.timer = now + hold_down_time
		end
		if rr.state ~= key_state.Valid or verbose() then
			log('[ ta ] key: ' .. key_tag .. ' state: '..rr.state)
		end
		table.insert(keyset, rr)
		return true
	end
	return false
end

-- TA is missing in the new key set.  The time is in seconds.
local function ta_missing(ta, hold_down_time)
	-- Key is removed (KeyRem)
	local keep_ta = true
	local key_tag = C.kr_dnssec_key_tag(ta.type, ta.rdata, #ta.rdata)
	assert(key_tag >= 0 and key_tag <= 65535, string.format('invalid RR: %s: %s',
	       kres.rr2str(ta), ffi.string(C.knot_strerror(key_tag))))
	if ta.state == key_state.Valid then
		ta.state = key_state.Missing
		ta.timer = os.time() + hold_down_time

	-- Remove key that is missing for too long
	elseif ta.state == key_state.Missing and os.difftime(ta.timer, os.time()) <= 0 then
		ta.state = key_state.Removed
		log('[ ta ] key: '..key_tag..' removed because missing for too long')
		keep_ta = false

	-- Purge pending key
	elseif ta.state == key_state.AddPend then
		log('[ ta ] key: '..key_tag..' purging')
		keep_ta = false
	end
	log('[ ta ] key: '..key_tag..' state: '..ta.state)
	return keep_ta
end

local active_refresh, update -- forwards

-- Plan an event for refreshing the root DNSKEYs and re-scheduling itself
local function refresh_plan(keyset, delay, is_initial)
	local owner_str = kres.dname2str(keyset.owner) -- maybe fix converting back and forth?
	keyset.refresh_ev = event.after(delay, function ()
		resolve(owner_str, kres.type.DNSKEY, kres.class.IN, 'NO_CACHE',
		function (pkt)
			-- Schedule itself with updated timeout
			local delay_new = active_refresh(keyset, kres.pkt_t(pkt), is_initial)
			delay_new = keyset.refresh_time or trust_anchors.refresh_time or delay_new
			log('[ ta ] next refresh for ' .. owner_str .. ' in '
					.. delay_new/hour .. ' hours')
			refresh_plan(keyset, delay_new)
		end)
	end)
end

-- Refresh the DNSKEYs from the packet, and return time to the next check.
active_refresh = function (keyset, pkt, is_initial)
	local retry = true
	if pkt:rcode() == kres.rcode.NOERROR then
		local records = pkt:section(kres.section.ANSWER)
		local new_keys = {}
		for _, rr in ipairs(records) do
			if rr.type == kres.type.DNSKEY then
				table.insert(new_keys, rr)
			end
		end
		update(keyset, new_keys, is_initial)
		retry = false
	else
		warn('[ ta ] active refresh failed for ' .. kres.dname2str(keyset.owner)
				.. ' with rcode: ' .. pkt:rcode())
	end
	-- Calculate refresh/retry timer (RFC 5011, 2.3)
	local min_ttl = retry and day or 15 * day
	for _, rr in ipairs(keyset) do -- 10 or 50% of the original TTL
		min_ttl = math.min(min_ttl, (retry and 100 or 500) * rr.ttl)
	end
	return math.max(hour, min_ttl)
end

-- Write keyset to a file.  States and timers are stored in comments.
local function keyset_write(keyset)
	if not keyset.filename then return false end -- not to be persisted
	local fname_tmp = keyset.filename .. '.lock.' .. tostring(worker.pid);
	local file = assert(io.open(fname_tmp, 'w'))
	for i = 1, #keyset do
		local ta = keyset[i]
		ta.comment = ' ' .. ta.state .. ':' .. (ta.timer or '')
					.. ' ; KeyTag:' .. ta.key_tag -- the tag is just for humans
		local rr_str = kres.rr2str(ta) .. '\n'
		if ta.state ~= key_state.Valid and ta.state ~= key_state.Missing then
			rr_str = '; '..rr_str -- Invalidate key string (for older kresd versions)
		end
		file:write(rr_str)
	end
	file:close()
	assert(os.rename(fname_tmp, keyset.filename))
end

-- Search the values of a table and return the corrseponding key (or nil).
local function table_search(t, val)
	for k, v in pairs(t) do
		if v == val then
			return k
		end
	end
	return nil
end

-- For each RR, parse .state and .timer from .comment.
local function keyset_parse_comments(tas, default_state)
	for _, ta in pairs(tas) do
		ta.state = default_state
		if ta.comment then
			string.gsub(ta.comment, '^%s*(%a+):(%d*)', function (state, time)
				if table_search(key_state, state) then
					ta.state = state
				end
				ta.timer = tonumber(time) -- nil on failure
			end)
			ta.comment = nil
		end
	end
	return tas
end

-- Read keyset from a file.  (This includes the key states and timers.)
local function keyset_read(path)
	-- First load the regular entries, trusting them.
	local zonefile = require('zonefile')
	local tas, err = zonefile.file(path)
	if not tas then
		return tas, err
	end
	keyset_parse_comments(tas, key_state.Valid)

	-- The untrusted keys are commented out but important to load.
	for line in io.lines(path) do
		if line:sub(1, 2) == '; ' then
			-- Ignore the line if it fails to parse including recognized .state.
			local l_set = zonefile.string(line:sub(3))
			if l_set and l_set[1] then
				keyset_parse_comments(l_set)
				if l_set[1].state then
					table.insert(tas, l_set[1])
				end
			end
		end
	end

	for _, ta in pairs(tas) do
		local ta_keytag = C.kr_dnssec_key_tag(ta.type, ta.rdata, #ta.rdata)
		if not (ta_keytag >= 0 and ta_keytag <= 65535) then
			return nil, string.format('invalid key: "%s": %s',
				kres.rr2str(ta), ffi.string(C.knot_strerror(ta_keytag)))
		end
		ta.key_tag = ta_keytag
	end
	return tas
end

-- Replace current TAs for given owner by the "trusted" ones from passed keyset.
-- Return the number of trusted keys for the owner.
local function keyset_publish(keyset)
	local store = kres.context().trust_anchors
	local count = 0
	C.kr_ta_del(store, keyset.owner)
	for _, ta in ipairs(keyset) do
		-- Key MAY be used as a TA only in these two states (RFC5011, 4.2)
		if ta.state == key_state.Valid or ta.state == key_state.Missing then
			if C.kr_ta_add(store, ta.owner, ta.type, ta.ttl, ta.rdata, #ta.rdata) == 0 then
				count = count + 1
			end
		end
	end
	if count == 0 then
		warn('[ ta ] ERROR: no anchors are trusted for ' ..
		     kres.dname2str(keyset.owner) .. ' !')
	end
	return count
end


-- Update existing keyset; return true if successful.
-- Param `is_initial` (bool): force .NewKey states to .Valid, i.e. init empty keyset.
update = function (keyset, new_keys, is_initial)
	if not new_keys then return false end

	-- Filter TAs to be purged from the keyset (KeyRem), in three steps
	-- 1: copy TAs to be kept to `keepset`
	local hold_down = (keyset.hold_down_time or trust_anchors.hold_down_time) / 1000
	local keepset = {}
	local keep_removed = keyset.keep_removed or trust_anchors.keep_removed
	for _, ta in ipairs(keyset) do
		local keep = true
		if not ta_find(new_keys, ta) then
			-- Ad-hoc: RFC 5011 doesn't mention removing a Missing key.
			-- Let's do it after a very long period has elapsed.
			keep = ta_missing(ta, hold_down * 4)
		end
		-- Purge removed keys
		if ta.state == key_state.Removed then
			if keep_removed > 0 then
				keep_removed = keep_removed - 1
			else
				keep = false
			end
		end
		if keep then
			table.insert(keepset, ta)
		end
	end
	-- 2: remove all TAs - other settings etc. will remain in the keyset
	for i, _ in ipairs(keyset) do
		keyset[i] = nil
	end
	-- 3: move TAs to be kept into the keyset (same indices)
	for k, ta in pairs(keepset) do
		keyset[k] = ta
	end

	-- Evaluate new TAs
	for _, rr in ipairs(new_keys) do
		if (rr.type == kres.type.DNSKEY or rr.type == kres.type.DS) and rr.rdata ~= nil then
			ta_present(keyset, rr, hold_down, is_initial)
		end
	end

	-- Store the keyset
	keyset_write(keyset)

	-- Start using the new TAs.
	if keyset_publish(keyset) == 0 then
		-- TODO: try to rebootstrap if for root?
		return false
	elseif verbose() then
		log('[ ta ] refreshed trust anchors for domain ' .. kres.dname2str(keyset.owner) .. ' are:\n'
		    .. trust_anchors.summary(keyset.owner))
	end

	return true
end

local add_file = function (path, unmanaged)
	if not unmanaged then
		if not io.open(path .. '.lock', 'w') then
			error("[ ta ] ERROR: write access needed to keyfile dir '"..path.."'")
		end
		os.remove(path .. ".lock")
	end

	-- Bootstrap if requested and keyfile doesn't exist
	if not unmanaged and not io.open(path, 'r') then
		log("[ ta ] keyfile '%s': doesn't exist, bootstrapping", path);
		local tas, msg = bootstrap(trust_anchors.bootstrap_url, trust_anchors.bootstrap_ca)
		if not tas then
			msg = msg .. '\n'
				.. '[ ta ] Failed to bootstrap root trust anchors; see:\n'
				.. '       https://knot-resolver.readthedocs.io/en/latest/daemon.html#enabling-dnssec'
			error(msg)
		end
		print(msg)
		trustanchor(tas)
		-- Fetch DNSKEY immediately
		if not trust_anchors.keysets['\0'] then
			trust_anchors.keysets['\0'] = { owner = '\0' }
		end
		local keyset = trust_anchors.keysets['\0']
		keyset.filename = path
		if keyset.refresh_ev then event.cancel(keyset.refresh_ev) end
		refresh_plan(keyset, 0, true)
		return
	end
	if not unmanaged and path == (trust_anchors.keysets['\0'] or {}).filename then
		return
	end

	-- Parse the file and check its sanity
	local keyset, err = keyset_read(path)
	if not keyset then
		panic("[ ta ] ERROR: failed to read anchors from '%s' (%s)", path, err)
	end
	if not unmanaged then keyset.filename = path end
	if not keyset[1] then
		panic("[ ta ] ERROR: failed to read anchors from '%s'", path)
	end
	if not unmanaged then keyset.filename = path end
	local owner = keyset[1].owner
	for _, ta in ipairs(keyset) do
		if ta.owner ~= owner then
			panic("[ ta ] ERROR: mixed owner names found in file '%s'! " ..
			      "Do not mix %s and %s TAs in single file",
			      path, kres.dname2str(ta.owner), kres.dname2str(owner))
		end
	end
	keyset.owner = owner

	local owner_str = kres.dname2str(owner)
	if trust_anchors.keysets[owner] then
		warn('[ ta ] warning: overriding previously set trust anchors for ' .. owner_str)
		local refresh_ev = trust_anchors.keysets[owner].refresh_ev
		if refresh_ev then event.cancel(refresh_ev) end
	end
	trust_anchors.keysets[owner] = keyset

	-- Parse new keys, refresh eventually
	if keyset_publish(keyset) ~= 0 and verbose() then
		log('[ ta ] installed trust anchors for domain ' .. owner_str .. ' are:\n'
		    .. trust_anchors.summary(owner))
	end
	-- TODO: if failed and for root, try to rebootstrap?

	refresh_plan(keyset, 10 * sec, false)
end

local function ta_str(owner)
	local owner_str = kres.dname2str(owner) .. ' '
	local msg = ''
	for _, nta in pairs(trust_anchors.insecure) do
		if owner == kres.str2dname(nta) then
			msg = owner_str .. 'is negative trust anchor\n'
		end
	end
	if not trust_anchors.keysets[owner] then
		if #msg > 0 then  -- it is normal that NTA does not have explicit TA
			return msg
		else
			return owner_str .. 'has no explicit trust anchors\n'
		end
	end
	if #msg > 0 then
		msg = msg .. 'WARNING! negative trust anchor also has an explicit TA\n'
	end
	for _, ta in ipairs(trust_anchors.keysets[owner]) do
		msg = msg .. kres.rr2str(ta) .. '\n'
	end
	return msg
end

-- TA store management, for user docs see ../README.rst
trust_anchors = {
	-- [internal] table indexed by dname;
	--  each item is a list of RRs and additionally contains:
	--   - owner - that dname (for simplicity)
	--   - [optional] filename in which to persist the state
	--   - [optional] overrides for global defaults of
	--     hold_down_time, refresh_time, keep_removed
	-- The RR tables also contain some additional TA-specific fields.
	keysets = {},

	-- Documented properties:
	insecure = {},
	hold_down_time = 30 * day,
	refresh_time = nil,
	keep_removed = 0,

	bootstrap_url = 'https://data.iana.org/root-anchors/root-anchors.xml',
	bootstrap_ca = '@ETCDIR@/icann-ca.pem',
	-- change empty string to nil
	keyfile_default = ('@KEYFILE_DEFAULT@' ~= '' and '@KEYFILE_DEFAULT@') or nil,

	-- Load keys from a file, 5011-managed by default.
	-- If managed and the file doesn't exist, try bootstrapping the root into it.
	add_file = add_file,
	config = add_file,

	-- Add DS/DNSKEY record(s) (unmanaged)
	add = function (keystr)
		local ret = trustanchor(keystr)
		if verbose() then log(trust_anchors.summary()) end
		return ret
	end,
	-- Negative TA management
	set_insecure = function (list)
		assert(type(list) == 'table', 'parameter must be list of domain names (e.g. {"a.test", "b.example"})')
		local store = kres.context().negative_anchors
		C.kr_ta_clear(store)
		for i = 1, #list do
			local dname = kres.str2dname(list[i])
			C.kr_ta_add(store, dname, kres.type.DS, 0, nil, 0)
		end
		trust_anchors.insecure = list
	end,
	summary = function (single_owner)
		if single_owner then  -- single domain
			return ta_str(single_owner)
		end

		-- all domains
		local msg = ''
		local ta_count = 0
		local seen = {}
		for _, nta_str in pairs(trust_anchors.insecure) do
			local owner = kres.str2dname(nta_str)
			seen[owner] = true
			msg = msg .. ta_str(owner)
		end
		for owner, _ in pairs(trust_anchors.keysets) do
			if not seen[owner] then
				ta_count = ta_count + 1
				msg = msg .. ta_str(owner)
			end
		end
		if ta_count == 0 then
			msg = msg .. 'No valid trust anchors, DNSSEC validation is disabled\n'
		end
		return msg
	end,
}

-- Syntactic sugar for TA store
setmetatable(trust_anchors, {
	__newindex = function (t,k,v)
	if     k == 'file' then t.config(v)
	elseif k == 'negative' then t.set_insecure(v)
	else   rawset(t, k, v) end
	end,
})

return trust_anchors