-- lx-bundle:true

-- This file is part of Lx, created by Liqube <www.liqube.com>
-- Copyright Liqube Audio. All rights reserved.
-- Lx Interim License. See LICENSE for details.

--- Minimal Compound File Binary reader for legacy Microsoft Office files
-- Supports regular and mini streams used by DOC and XLS readers
-- @module plugins.msoutils.cfb

local cfb = {}

local MAGIC = "\208\207\017\224\161\177\026\225"
local FREE_SECTOR = 0xffffffff
local END_OF_CHAIN = 0xfffffffe

local function u16le(data, offset)
	local a, b = data:byte(offset, offset + 1)
	if not b then
		return nil
	end
	return a | (b << 8)
end

local function u32le(data, offset)
	local a, b, c, d = data:byte(offset, offset + 3)
	if not d then
		return nil
	end
	return a | (b << 8) | (c << 16) | (d << 24)
end

local function read_all(filename)
	local file, err = io.open(filename, "rb")
	if not file then
		return nil, err or "open failed"
	end

	local data = file:read("*a") or ""
	file:close()
	return data
end

local function sector_offset(sector_id, sector_size)
	return 512 + (sector_id * sector_size) + 1
end

local function utf16le_name(data, offset, byte_len)
	local out = {}
	local limit = math.max(0, byte_len - 2)

	for i = 0, limit - 1, 2 do
		local code = u16le(data, offset + i)
		if not code or code == 0 then
			break
		end
		out[#out + 1] = utf8.char(code)
	end

	return table.concat(out)
end

local Archive = {}
Archive.__index = Archive

--- Read one regular-sector chain and enforce its declared byte count
-- @param first_sector integer first regular sector id
-- @param max_bytes integer|nil declared stream bytes
-- @return string|nil data
-- @return string|nil err
function Archive:read_chain(first_sector, max_bytes)
	if first_sector == nil or first_sector < 0 or first_sector == END_OF_CHAIN or first_sector == FREE_SECTOR then
		if max_bytes and max_bytes > 0 then
			return nil, "CFB sector chain ended before its declared size"
		end

		return ""
	end

	local chunks = {}
	local sector = first_sector
	local visited = {}
	local total = 0

	while sector ~= END_OF_CHAIN and sector ~= FREE_SECTOR do
		local pos = sector_offset(sector, self.sector_size)
		-- Final data sector may contain only declared stream bytes
		if sector < 0 or pos > #self.data then
			return nil, "CFB sector is outside the file"
		end
		if visited[sector] then
			return nil, "CFB sector chain loop"
		end
		visited[sector] = true

		local chunk = self.data:sub(pos, pos + self.sector_size - 1)
		chunks[#chunks + 1] = chunk
		total = total + #chunk
		if max_bytes and total >= max_bytes then
			break
		end

		sector = self.fat[sector]
		if sector == nil then
			break
		end
	end

	local data = table.concat(chunks)
	if max_bytes and #data < max_bytes then
		return nil, "CFB sector chain ended before its declared size"
	end
	if max_bytes and #data > max_bytes then
		data = data:sub(1, max_bytes)
	end
	return data
end

--- Read one mini-sector chain and enforce its declared byte count
-- @param first_sector integer first mini sector id
-- @param max_bytes integer|nil declared stream bytes
-- @return string|nil data
-- @return string|nil err
function Archive:read_mini_chain(first_sector, max_bytes)
	if not self.mini_stream or first_sector == nil or first_sector < 0
		or first_sector == END_OF_CHAIN or first_sector == FREE_SECTOR then
		if max_bytes and max_bytes > 0 then
			return nil, "CFB mini sector chain ended before its declared size"
		end

		return ""
	end

	local chunks = {}
	local sector = first_sector
	local visited = {}
	local total = 0

	while sector ~= END_OF_CHAIN and sector ~= FREE_SECTOR do
		local pos = (sector * self.mini_sector_size) + 1
		if sector < 0 or pos > #self.mini_stream then
			return nil, "CFB mini sector is outside the root stream"
		end
		if visited[sector] then
			return nil, "CFB mini sector chain loop"
		end
		visited[sector] = true

		local chunk = self.mini_stream:sub(pos, pos + self.mini_sector_size - 1)
		chunks[#chunks + 1] = chunk
		total = total + #chunk
		if max_bytes and total >= max_bytes then
			break
		end

		sector = self.minifat[sector]
		if sector == nil then
			break
		end
	end

	local data = table.concat(chunks)
	if max_bytes and #data < max_bytes then
		return nil, "CFB mini sector chain ended before its declared size"
	end
	if max_bytes and #data > max_bytes then
		data = data:sub(1, max_bytes)
	end
	return data
end

function Archive:read_stream(name)
	local entry = self.streams[name] or self.streams[name:upper()]
	if not entry then
		return nil, "CFB stream not found: " .. tostring(name)
	end

	if entry.size < self.mini_cutoff and entry.name ~= "Root Entry" then
		return self:read_mini_chain(entry.start_sector, entry.size)
	end

	return self:read_chain(entry.start_sector, entry.size)
end

--- Read and validate the chained DIFAT sectors
-- @param data string complete archive bytes
-- @param sector_size integer regular sector size
-- @param first_difat_sector integer first DIFAT sector id
-- @param difat_count integer declared DIFAT sector count
-- @param sector_count integer sectors physically present in the archive
-- @param fat_sector_count integer declared FAT sector count
-- @param difat table header DIFAT entries
-- @return boolean|nil ok
-- @return string|nil err
local function read_difat(
	data,
	sector_size,
	first_difat_sector,
	difat_count,
	sector_count,
	fat_sector_count,
	difat
)
	local sector = first_difat_sector
	local ids_per_sector = (sector_size // 4) - 1
	local visited = {}

	for _ = 1, difat_count do
		if sector == END_OF_CHAIN or sector == FREE_SECTOR then
			return nil, "CFB DIFAT chain ended before its declared count"
		end
		if sector >= sector_count then
			return nil, "CFB DIFAT sector is outside the file"
		end
		if visited[sector] then
			return nil, "CFB DIFAT sector chain loop"
		end
		visited[sector] = true

		local pos = sector_offset(sector, sector_size)
		for i = 0, ids_per_sector - 1 do
			local id = u32le(data, pos + (i * 4))
			if id and id ~= FREE_SECTOR and #difat < fat_sector_count then
				difat[#difat + 1] = id
			end
		end
		sector = u32le(data, pos + (ids_per_sector * 4))
	end

	return true
end

--- Read FAT sectors after validating every referenced sector id
-- @param data string complete archive bytes
-- @param sector_size integer regular sector size
-- @param sector_count integer sectors physically present in the archive
-- @param difat table FAT sector ids
-- @param fat_sector_count integer declared FAT sector count
-- @return table|nil fat
-- @return string|nil err
local function read_fat(data, sector_size, sector_count, difat, fat_sector_count)
	if #difat < fat_sector_count then
		return nil, "CFB file does not contain all declared FAT sectors"
	end

	local fat = {}
	local index = 0
	local visited = {}

	for i = 1, fat_sector_count do
		local sector = difat[i]

		if sector >= sector_count then
			return nil, "CFB FAT sector is outside the file"
		end
		if visited[sector] then
			return nil, "CFB FAT sector is referenced more than once"
		end
		visited[sector] = true

		local pos = sector_offset(sector, sector_size)
		for j = 0, (sector_size // 4) - 1 do
			fat[index] = u32le(data, pos + (j * 4))
			index = index + 1
		end
	end

	return fat
end

local function read_directory(archive, first_directory_sector)
	local directory, err = archive:read_chain(first_directory_sector)
	if not directory then
		return nil, err
	end

	local entries = {}
	local streams = {}

	for offset = 1, #directory - 127, 128 do
		local name_len = u16le(directory, offset + 0x40) or 0
		local entry_type = directory:byte(offset + 0x42) or 0
		local start_sector = u32le(directory, offset + 0x74) or END_OF_CHAIN
		local size = u32le(directory, offset + 0x78) or 0
		local name = utf16le_name(directory, offset, name_len)

		if name ~= "" and entry_type ~= 0 then
			local entry = {
				name = name,
				type = entry_type,
				start_sector = start_sector,
				size = size,
			}
			entries[#entries + 1] = entry
			streams[name] = entry
			streams[name:upper()] = entry
		end
	end

	archive.entries = entries
	archive.streams = streams
	archive.root = streams["Root Entry"] or streams["ROOT ENTRY"]
	return true
end

--- Read the miniFAT and its root mini stream
-- @param archive table open CFB archive
-- @param first_minifat_sector integer first regular sector
-- @param minifat_sector_count integer declared regular-sector count
-- @return boolean|nil success
-- @return string|nil error_message
local function read_minifat(archive, first_minifat_sector, minifat_sector_count)
	local minifat_stream, minifat_error = archive:read_chain(
		first_minifat_sector,
		minifat_sector_count * archive.sector_size
	)

	if not minifat_stream then
		return nil, minifat_error
	end

	local minifat = {}

	for offset = 1, #minifat_stream - 3, 4 do
		minifat[(offset - 1) // 4] = u32le(minifat_stream, offset)
	end

	archive.minifat = minifat

	if archive.root then
		local mini_stream, mini_stream_error = archive:read_chain(
			archive.root.start_sector,
			archive.root.size
		)

		if not mini_stream then
			return nil, mini_stream_error
		end

		archive.mini_stream = mini_stream
	else
		archive.mini_stream = ""
	end

	return true
end

function cfb.open(filename)
	local data, err = read_all(filename)
	if not data then
		return nil, err
	end
	if data:sub(1, 8) ~= MAGIC then
		return nil, "not an OLE compound file"
	end
	if #data < 512 then
		return nil, "truncated CFB header"
	end

	local sector_shift = u16le(data, 0x1e + 1) or 9
	local mini_sector_shift = u16le(data, 0x20 + 1) or 6

	if (sector_shift ~= 9 and sector_shift ~= 12) or mini_sector_shift ~= 6 then
		return nil, "invalid CFB sector size"
	end

	local sector_size = 1 << sector_shift
	local sector_count = (#data - 512) // sector_size
	local fat_sector_count = u32le(data, 0x2c + 1) or 0
	local first_directory_sector = u32le(data, 0x30 + 1) or END_OF_CHAIN
	local mini_cutoff = u32le(data, 0x38 + 1) or 4096
	local first_minifat_sector = u32le(data, 0x3c + 1) or END_OF_CHAIN
	local minifat_sector_count = u32le(data, 0x40 + 1) or 0
	local first_difat_sector = u32le(data, 0x44 + 1) or END_OF_CHAIN
	local difat_sector_count = u32le(data, 0x48 + 1) or 0

	if fat_sector_count > sector_count or difat_sector_count > sector_count then
		return nil, "CFB allocation table count exceeds the file size"
	end

	local difat = {}
	for i = 0, 108 do
		local id = u32le(data, 0x4c + 1 + (i * 4))
		if id and id ~= FREE_SECTOR then
			difat[#difat + 1] = id
		end
	end
	local difat_ok, difat_err = read_difat(
		data,
		sector_size,
		first_difat_sector,
		difat_sector_count,
		sector_count,
		fat_sector_count,
		difat
	)

	if not difat_ok then
		return nil, difat_err
	end

	local fat, fat_err = read_fat(data, sector_size, sector_count, difat, fat_sector_count)

	if not fat then
		return nil, fat_err
	end

	local archive = setmetatable({
		data = data,
		sector_size = sector_size,
		sector_count = sector_count,
		mini_sector_size = 1 << mini_sector_shift,
		mini_cutoff = mini_cutoff,
		fat = fat,
		minifat = {},
		mini_stream = "",
	}, Archive)

	local ok, read_err = read_directory(archive, first_directory_sector)
	if not ok then
		return nil, read_err
	end
	local minifat_ok, minifat_err = read_minifat(archive, first_minifat_sector, minifat_sector_count)

	if not minifat_ok then
		return nil, minifat_err
	end

	return archive
end

return cfb
