malware-filter.lua

-- File: malware-filter.lua
-- Author: Dylan Blanqué

local re -- Regex Module
if f.isModuleAvailable("rex_pcre") then
	re = require"rex_pcre"
elseif f.isModuleAvailable("rex_pcre2") then
	re = require"rex_pcre2"
else
	mainlog(
		"pdns-recursor-scripts malware-filter.lua requires rex_pcre or rex_pcre2 to be installed",
		pdns.loglevels.Error
	)
	return false
end

-- Pre-compiled Regex Patterns
-- Matches Adblock Format
local re_pat_adblock=re.new(
	"^(\\|){2}(.*)\\^$"
)
-- Matches PCRE Format
local re_pat_pcre=re.new(
	"(.*)[$^|](.*)"
)
-- Matches Wildcard Format
local re_pat_wild=re.new(
	"^\\*\\..*"
)
-- Matches HOSTS Format
local re_pat_hosts=re.new(
	"^(0\\.0\\.0\\.0|127\\.0\\.0\\.1|::|2001:1::1)\\s+.*"
)
-- Matches RFC for TLDs
local re_pat_domain=re.new(
	"^(?:(?!-)[A-Za-z0-9-]{1,63}(?<!-)\\.)+[A-Za-z]{2,63}$"
)

-- Blocklist Objects (DS, CAS, NMG, Tables, etc.)
local mf_dnsbl = newDS()
local mf_rebl = {}
local mf_ipbl = newCAS()
local mf_cidrbl = newNMG()
local mf_whitelist_ds = newDS()
local mf_nxdomain_ds = newDS()

-- DNSBL Metrics
local dnsbl_hits = getMetric("dnsbl-hits")
local dnsbl_parsed_loaded_metric = getMetric("dnsbl-parsed-success")
local dnsbl_parsed_ignored_metric = getMetric("dnsbl-parsed-ignored")
local dnsbl_parsed_errors_metric = getMetric("dnsbl-parsed-errors")
local dnsbl_loaded_adblock = getMetric("dnsbl-loaded-adblock")
local dnsbl_loaded_pcre = getMetric("dnsbl-loaded-pcre")
local dnsbl_loaded_wildcard = getMetric("dnsbl-loaded-wildcard")
local dnsbl_loaded_hosts = getMetric("dnsbl-loaded-hosts")
local dnsbl_loaded_normal = getMetric("dnsbl-loaded-normal")

-- IPBL / CIDRBL Metrics
local ipbl_loaded_metric = getMetric("ipbl-loaded")
local cidrbl_loaded_metric = getMetric("cidrbl-loaded")
local ipbl_hits = getMetric("ipbl-hits")

-- For debugging
local whitelist_files_loaded = getMetric("whitelist-files-loaded")
local dnsbl_files_loaded = getMetric("dnsbl-files-loaded")
local ipbl_files_loaded = getMetric("ipbl-files-loaded")

-- Set NXDOMAIN Hosts
local nxdomain_hosts={
	"localhost",
	"localhost.localdomain",
	"local",
	"broadcasthost",
	"ip6-localhost",
	"ip6-loopback",
	"ip6-localnet",
	"ip6-mcastprefix",
	"ip6-allnodes",
	"ip6-allrouters",
	"ip6-allhosts",
}
for _, host in ipairs(nxdomain_hosts) do
	mf_nxdomain_ds:add(host)
end

local function get_list_files_in_dir(search_dir)
	local files = {}
	for dir in io.popen("ls -pa " .. search_dir .. " | grep -v / | grep -E \".*\\.(list|txt|hosts)$\""):lines() 
	do
		table.insert(files, search_dir .. "/" .. dir)
	end
	return files
end

-- Get relevant files
local dnsbl_file_table = get_list_files_in_dir(g.pdns_scripts_path.."/dnsbl.d")
local ipbl_file_table = get_list_files_in_dir(g.pdns_scripts_path.."/ipbl.d")

local function is_dq_internal(dq)
	pdnslog(
		"is_dq_internal(): checking if internal record for " .. tostring(dq.qname),
		pdns.loglevels.Debug
	)
	if mf_nxdomain_ds:check(dq.qname) then
		pdnslog(
			"is_dq_internal(): " .. tostring(dq.qname) .. " is internal",
			pdns.loglevels.Debug
		)
		return true
	end

	return false
end

local function is_dq_whitelisted(dq)
	local qname_no_dot = dq.qname:toStringNoDot()
	pdnslog(
		"is_dq_whitelisted(): checking whitelist for " .. qname_no_dot,
		pdns.loglevels.Debug
	)

	-- Check in DNSSuffixMatchGroup
	if mf_whitelist_ds:check(dq.qname) then
		pdnslog(
			"is_dq_whitelisted(): " .. qname_no_dot .. " is whitelisted",
			pdns.loglevels.Debug
		)
		return true
	end

	return false
end

local function is_qname_whitelisted(qname_no_dot)
	if not qname_no_dot or f.empty_str(qname_no_dot) then
		return false
	end

	-- Check in DNSSuffixMatchGroup
	if mf_whitelist_ds:check(newDN(qname_no_dot)) then
		mainlog(
			"is_qname_whitelisted(): " .. qname_no_dot .. " is whitelisted",
			pdns.loglevels.Debug
		)
		return true
	end

	return false
end

-- Precompile regex patterns for better performance
local comment_or_empty = function(line) 
	return f.is_comment(line) or f.empty_str(line) 
end

-- loads contents of a file line by line with the given processing function
local function load_list_file(fn_name, fn_process, filename)
	if not f.fileExists(filename) then
		mainlog(
			fn_name .. "(): could not open file " .. filename,
			pdns.loglevels.Warning
		)
		return
	end

	-- Process each line
	local line_errors = 0
	local line_idx = 1
	for line in io.lines(filename) do
		local ok, result = pcall(function()
			return fn_process(line)
		end)
		if not ok then
			line_errors = line_errors + 1
			local _filename = f.string_split(filename, "/")
			mainlog(
				string.format(
					"%s(): could not parse line %d for file %s",
					fn_name,
					line_idx,
					_filename[#_filename]
				),
				pdns.loglevels.Error
			)
		end
		line_idx = line_idx + 1
	end

	mainlog(
		string.format(
			"%s(): %s successfully loaded",
			fn_name,
			filename
		),
		pdns.loglevels.Notice
	)
	return line_errors
end

-- loads contents of a file line by line into the whitelist table
local function load_whitelist_file(filename, whitelist_ds)

	-- Validate input parameters
	if type(filename) ~= "string" or not whitelist_ds then
		mainlog("load_whitelist_file(): invalid parameters", pdns.loglevels.Error)
		return false
	end

	local function process_line(line)
		-- Trim whitespace from line
		line = line:match("^%s*(.-)%s*$")

		-- Ignore commented lines
		if comment_or_empty(line) then
			return
		end

		if re_pat_wild:match(line) then
			local trimmed_dn = line:gsub('*%.', "")
			whitelist_ds:add(newDN(trimmed_dn))
		else
			whitelist_ds:add(newDN(line))
		end
		mainlog(
			line .. " is whitelisted",
			pdns.loglevels.Notice
		)
	end

	if getRecursorThreadId() == 1 then
		whitelist_files_loaded:inc()
	end
	return load_list_file("load_whitelist_file", process_line, filename)
end

-- Loads contents of a file line by line into the dnsbl table
local function load_dnsbl_file(filename, dnsbl_ds, rebl_list)

	-- Validate input parameters
	if type(filename) ~= "string" or not dnsbl_ds or not rebl_list then
		mainlog("load_dnsbl_file(): invalid parameters", pdns.loglevels.Error)
		return false
	end

	local ignored_lines = 0
	local loaded_lines = 0
	local adblock_lines = 0
	local pcre_lines = 0
	local wildcard_lines = 0
	local hosts_lines = 0
	local normal_lines = 0
	local start_time = os.time()

	local function process_line(line)
		-- Trim whitespace from line
		line = line:match("^%s*(.-)%s*$")

		if comment_or_empty(line) then
			ignored_lines = ignored_lines + 1
			return
		end

		-- ADBLOCK FORMAT
		if re_pat_adblock:match(line) then
			local stripped = line:gsub("||", ""):gsub("%^", "")
			if is_qname_whitelisted(stripped) then
				return
			end
			dnsbl_ds:add(stripped)
			loaded_lines = loaded_lines + 1
			adblock_lines = adblock_lines + 1

		-- PCRE FORMAT
		elseif re_pat_pcre:match(line) then
			-- Insert pre-compiled regex into table
			table.insert(rebl_list, re.new(line))
			loaded_lines = loaded_lines + 1
			pcre_lines = pcre_lines + 1

		-- WILDCARD FORMAT
		elseif re_pat_wild:match(line) then
			local wilded = line:gsub('*%.', "")
			if is_qname_whitelisted(wilded) then
				return
			end
			dnsbl_ds:add(wilded)
			loaded_lines = loaded_lines + 1
			wildcard_lines = wildcard_lines + 1

		-- HOSTS FORMAT
		elseif re_pat_hosts:match(line) then
			local host = f.extract_hosts_domain(line)
			host = f.trim_hosts_comment(host)
			if is_qname_whitelisted(host) then
				return
			end
			dnsbl_ds:add(host)
			loaded_lines = loaded_lines + 1
			hosts_lines = hosts_lines + 1

		-- STANDARD FORMAT
		elseif re_pat_domain:match(line) then
			dnsbl_ds:add(line)
			loaded_lines = loaded_lines + 1
			normal_lines = normal_lines + 1
		else
			error("Could not parse line.")
		end
	end

	-- Load file
	local error_lines = load_list_file("load_dnsbl_file", process_line, filename)
	local duration = os.time() - start_time

	-- Set metrics if main thread only
	if getRecursorThreadId() == 1 then
		-- Set general metrics
		dnsbl_parsed_loaded_metric:incBy(loaded_lines)
		-- Set detailed metrics
		dnsbl_loaded_adblock:incBy(adblock_lines)
		dnsbl_loaded_pcre:incBy(pcre_lines)
		dnsbl_loaded_wildcard:incBy(wildcard_lines)
		dnsbl_loaded_hosts:incBy(hosts_lines)
		dnsbl_loaded_normal:incBy(normal_lines)
		-- Set ignored metric
		dnsbl_parsed_ignored_metric:incBy(ignored_lines)
		dnsbl_parsed_errors_metric:incBy(error_lines)
	end

	dnsbl_files_loaded:inc()
	return true
end

-- loads contents of a file line by line into the ipbl table
local function load_ipbl_file(filename, ipbl_cas, cidrbl_nmg)

	-- Validate input parameters
	if type(filename) ~= "string" or not ipbl_cas or not cidrbl_nmg then
		mainlog("load_ipbl_file(): invalid parameters", pdns.loglevels.Error)
		return false
	end

	local function process_line(line)
		-- Ignore commented lines
		if comment_or_empty(line) then
			return
		end

		-- Check if it's a CIDR
		line = f.trim_hosts_comment(line)
		if string.find(line, "/") then
			mainlog(
				"load_ipbl_file(): loaded CIDR " .. line,
				pdns.loglevels.Debug
			)
			cidrbl_nmg:addMask(line)
			if getRecursorThreadId() == 1 then
				cidrbl_loaded_metric:inc()
			end
		else -- Assume IP Address
			mainlog(
				"load_ipbl_file(): loaded IP " .. line,
				pdns.loglevels.Debug
			)
			ipbl_cas:add(line)
			if getRecursorThreadId() == 1 then
				ipbl_loaded_metric:inc()
			end
		end
	end

	-- Load file
	load_list_file("load_ipbl_file", process_line, filename)
	if getRecursorThreadId() == 1 then
		ipbl_files_loaded:inc()
	end
	return true
end

local function get_sinkhole(dq)
	if dq.qtype == pdns.A or dq.qtype == pdns.ANY then
		return {
			["0.0.0.0"]=pdns.A
		}
	elseif dq.qtype == pdns.AAAA or dq.qtype == pdns.ANY then
		return {
			["2001:1::1"]=pdns.AAAA
		}
	-- elseif dq.qtype == pdns.CNAME or dq.qtype == pdns.ANY then
	-- 	return {
	-- 		["0.0.0.0"]=pdns.A,
	-- 		["::"]=pdns.AAAA
	-- 	}
	end
	if dq.qtype == pdns.HTTPS then
		return {
			["1 . alpn=h3,h3-29,h2 ipv4hint=0.0.0.0 ipv6hint=2001:1::1"]=pdns.HTTPS
		}
	end
	return nil
end

-- this function is hooked before resolving starts
local function preresolve_mf(dq)
	if not g.initialized_mf then
		dq.appliedPolicy.policyKind = pdns.policykinds.Drop
		return false
	end

	if is_dq_internal(dq) then
		dq.appliedPolicy.policyKind = pdns.policykinds.NXDOMAIN
		return false
	end

	if is_dq_whitelisted(dq) then
		return false
	end

	-- If the last char is a dot, remove it (for non-canonical format)
	local qname_no_trailing_dot = f.qname_remove_trailing_dot(dq)
	local resolve_sinkhole = mf_dnsbl:check(dq.qname) -- NORMAL BL Check

	-- PCRE BL Check
	if not resolve_sinkhole then
		for k, p in pairs(mf_rebl) do
			if resolve_sinkhole then
				break
			end
			-- Log PCRE Pattern
			pdnslog(
				"PCRE Pattern: ".. qname_no_trailing_dot .." | " .. tostring(p),
				pdns.loglevels.Debug
			)
			resolve_sinkhole = p:match(qname_no_trailing_dot)
		end
	end

	pdnslog("Pre-Resolving " .. tostring(dq.qname), pdns.loglevels.Debug)
	pdnslog("PCRE Check: " .. tostring(re_check), pdns.loglevels.Debug)
	pdnslog("DNSBL Check: " .. tostring(ds_check), pdns.loglevels.Debug)
	-- check blocklist
	if resolve_sinkhole then
		local sinkhole = get_sinkhole(dq)
		if not sinkhole then
			dq.appliedPolicy.policyKind = pdns.policykinds.Drop
			return false
		end
		for dq_val, dq_qtype in pairs(sinkhole) do
			dq:addAnswer(dq_qtype, dq_val, g.options.default_ttl)
		end
		dnsbl_hits:inc()
		return true
	end

	-- default, do not rewrite this response
	pdnslog("Did not re-write " .. tostring(dq.qname), pdns.loglevels.Debug)
	return false
end

local function postresolve_mf(dq)
	if not g.initialized_mf then
		dq.appliedPolicy.policyKind = pdns.policykinds.Drop
		return false
	end

	if is_dq_whitelisted(dq) then
		return false
	end

	if dq.qtype ~= pdns.A and dq.qtype ~= pdns.AAAA then
		pdnslog(
			"Skipping postresolve_mf for ".. tostring(dq.qname),
			pdns.loglevels.Debug
		)
		return false
	end

	local dq_records = dq:getRecords()
	local blocked_records = {}
	local result_dq = {}
	local update_dq = false

	for dr_index, dr in ipairs(dq_records) do
		local dr_content = dr:getContent()
		local ip_check = false
		local cidr_check = false
		if not dr_content then
			pdnslog(
				"No DNSR Content for ".. tostring(dq.qname),
				pdns.loglevels.Debug
			)
			goto continue
		end
		-- Call function without raising exception to parent process
		local ok, dr_ca = pcall(newCA, dr_content)
		if not ok then
			goto continue
		end
		pdnslog("DNSR Content: " .. dr_ca:toString(), pdns.loglevels.Debug)

		if not ip_check and not cidr_check then
			cidr_check = mf_cidrbl:match(dr_ca)
			ip_check = mf_ipbl:check(dr_ca)
		end
		pdnslog("IP Check: " .. tostring(ip_check), pdns.loglevels.Debug)
		pdnslog("CIDR Check: " .. tostring(cidr_check), pdns.loglevels.Debug)

		if ip_check or cidr_check then
			update_dq = true
			table.insert(blocked_records, dr_index)
			for dq_val, dq_qtype in pairs(get_sinkhole(dq)) do
				dr:changeContent(dq_val)
			end
			pdnslog(
				string.format("Blocked Post-Resolving %s", tostring(dq.qname)),
				pdns.loglevels.Debug
			)
		else
			pdnslog(
				string.format("Allowed Post-Resolving %s", tostring(dq.qname)),
				pdns.loglevels.Debug
			)
		end
		::continue::
	end

	if update_dq then
		for i,v in ipairs(dq_records) do
			if not f.table_contains(blocked_records, i) then
				table.insert(result_dq, v)
			end
		end

		local use_result
		if f.table_len(result_dq) > 1 then
			use_result = result_dq
		else
			use_result = dq_records
		end

		dq:setRecords(use_result)
		pdnslog(
			string.format("Query Result %s", tostring(use_result)),
			pdns.loglevels.Debug
		)
		ipbl_hits:inc()
		return true
	end
	return false
end

local function main_mf()
	if g.initialized_mf then
		return
	end

	mainlog(
		"Executing main_mf().",
		pdns.loglevels.Debug
	)

	-- Load Whitelist
	if f.fileExists(g.pdns_scripts_path.."/conf.d/whitelist.txt") then
		load_whitelist_file(
			g.pdns_scripts_path.."/conf.d/whitelist.txt",
			mf_whitelist_ds
		)
	end

	-- Load DNSBL
	if g.options.use_dnsbl then
		for key, filename in pairs(dnsbl_file_table) do
			load_dnsbl_file(filename, mf_dnsbl, mf_rebl)
		end

		-- Add preresolve function to table
		mainlog(
			"Loading preresolve_mf into pre-resolve functions.",
			pdns.loglevels.Debug
		)
		f.addHookFunction("pre", "preresolve_mf", preresolve_mf)
	else
		mainlog("DNSBL Function not enabled. Set overrides in file overrides.lua", pdns.loglevels.Notice)
	end

	-- Load IPBL
	if g.options.use_ipbl then
		for key, filename in pairs(ipbl_file_table) do
			load_ipbl_file(filename, mf_ipbl, mf_cidrbl)
		end

		-- Add postresolve function to table
		mainlog(
			"Loading postresolve_mf into post-resolve functions.",
			pdns.loglevels.Debug
		)
		f.addHookFunction("post", "postresolve_mf", postresolve_mf)
	else
		mainlog("IPBL Function not enabled. Set overrides in file overrides.lua", pdns.loglevels.Notice)
	end

	g.initialized_mf = true
end

-- Execute malware-filter main logic
main_mf()