/*
 * Copyright 2023 Richard Hughes <richard@hughsie.com>
 *
 * SPDX-License-Identifier: LGPL-2.1-or-later
 */

#define G_LOG_DOMAIN "FuFirmware"

#include "config.h"

#include "fu-byte-array.h"
#include "fu-bytes.h"
#include "fu-common.h"
#include "fu-composite-input-stream.h"
#include "fu-coswid-firmware.h"
#include "fu-csv-firmware.h"
#include "fu-input-stream.h"
#include "fu-linear-firmware.h"
#include "fu-partial-input-stream.h"
#include "fu-pefile-firmware.h"
#include "fu-pefile-struct.h"
#include "fu-sbatlevel-section.h"
#include "fu-string.h"

/**
 * FuPefileFirmware:
 *
 * A PE file consists of a Microsoft MS-DOS stub, the PE signature, the COFF file header, and an
 * optional header, followed by section data.
 *
 * Documented:
 * https://learn.microsoft.com/en-gb/windows/win32/debug/pe-format
 */

typedef struct {
	gchar *authenticode_hash;
	guint16 subsystem_id;
} FuPefileFirmwarePrivate;

G_DEFINE_TYPE_WITH_PRIVATE(FuPefileFirmware, fu_pefile_firmware, FU_TYPE_FIRMWARE)
#define GET_PRIVATE(o) (fu_pefile_firmware_get_instance_private(o))

#define FU_PEFILE_SECTION_ID_STRTAB_SIZE 16

static void
fu_pefile_firmware_export(FuFirmware *firmware, FuFirmwareExportFlags flags, XbBuilderNode *bn)
{
	FuPefileFirmware *self = FU_PEFILE_FIRMWARE(firmware);
	FuPefileFirmwarePrivate *priv = GET_PRIVATE(self);
	fu_xmlb_builder_insert_kv(bn, "authenticode_hash", priv->authenticode_hash);
	fu_xmlb_builder_insert_kv(bn, "subsystem", fu_coff_subsystem_to_string(priv->subsystem_id));
}

static gboolean
fu_pefile_firmware_validate(FuFirmware *firmware,
			    GInputStream *stream,
			    gsize offset,
			    GError **error)
{
	return fu_struct_pe_dos_header_validate_stream(stream, offset, error);
}

typedef struct {
	gsize offset;
	gsize size;
	gchar *name;
} FuPefileFirmwareRegion;

static void
fu_pefile_firmware_add_region(GPtrArray *regions, const gchar *name, gsize offset, gsize size)
{
	FuPefileFirmwareRegion *r = g_new0(FuPefileFirmwareRegion, 1);
	r->name = g_strdup(name);
	r->offset = offset;
	r->size = size;
	g_ptr_array_add(regions, r);
}

static void
fu_pefile_firmware_region_free(FuPefileFirmwareRegion *r)
{
	g_free(r->name);
	g_free(r);
}

static gint
fu_pefile_firmware_region_sort_cb(gconstpointer a, gconstpointer b)
{
	const FuPefileFirmwareRegion *r1 = *((const FuPefileFirmwareRegion **)a);
	const FuPefileFirmwareRegion *r2 = *((const FuPefileFirmwareRegion **)b);
	if (r1->offset < r2->offset)
		return -1;
	if (r1->offset > r2->offset)
		return 1;
	return 0;
}

static gboolean
fu_pefile_firmware_parse_section(FuPefileFirmware *self,
				 GInputStream *stream,
				 guint idx,
				 gsize hdr_offset,
				 gsize strtab_offset,
				 GPtrArray *regions,
				 FuFirmwareParseFlags flags,
				 GError **error)
{
	g_autofree gchar *sect_id = NULL;
	g_autofree gchar *sect_id_tmp = NULL;
	g_autoptr(FuFirmware) img = NULL;
	g_autoptr(FuStructPeCoffSection) st = NULL;

	st = fu_struct_pe_coff_section_parse_stream(stream, hdr_offset, error);
	if (st == NULL) {
		g_prefix_error_literal(error, "failed to read section: ");
		return FALSE;
	}
	sect_id_tmp = fu_struct_pe_coff_section_get_name(st);
	if (sect_id_tmp == NULL) {
		sect_id = g_strdup_printf(".nul%04x", idx);
	} else if (sect_id_tmp[0] == '/') {
		guint64 str_idx = 0x0;
		guint8 buf[FU_PEFILE_SECTION_ID_STRTAB_SIZE] = {0};

		if (!fu_strtoull(sect_id_tmp + 1,
				 &str_idx,
				 0,
				 G_MAXUINT32,
				 FU_INTEGER_BASE_10,
				 error)) {
			g_prefix_error(error, "failed to parse section ID '%s': ", sect_id_tmp + 1);
			return FALSE;
		}
		if (!fu_input_stream_read_safe(stream,
					       buf,
					       sizeof(buf),
					       0x0,
					       strtab_offset + str_idx, /* seek */
					       sizeof(buf),
					       error))
			return FALSE;
		sect_id = fu_strsafe((const gchar *)buf, sizeof(buf));
		if (sect_id == NULL) {
			g_set_error_literal(error,
					    FWUPD_ERROR,
					    FWUPD_ERROR_INVALID_DATA,
					    "no section name");
			return FALSE;
		}
	} else {
		sect_id = g_steal_pointer(&sect_id_tmp);
	}

	/* create new firmware */
	if (g_strcmp0(sect_id, ".sbom") == 0) {
		img = fu_linear_firmware_new(FU_TYPE_COSWID_FIRMWARE);
	} else if (g_strcmp0(sect_id, ".sbat") == 0 || g_strcmp0(sect_id, ".sbata") == 0 ||
		   g_strcmp0(sect_id, ".sbatl") == 0) {
		img = fu_csv_firmware_new();
		fu_csv_firmware_add_column_id(FU_CSV_FIRMWARE(img), "$id");
		fu_csv_firmware_add_column_id(FU_CSV_FIRMWARE(img), "$version_raw");
		fu_csv_firmware_add_column_id(FU_CSV_FIRMWARE(img), "vendor_name");
		fu_csv_firmware_add_column_id(FU_CSV_FIRMWARE(img), "vendor_package_name");
		fu_csv_firmware_add_column_id(FU_CSV_FIRMWARE(img), "$version");
		fu_csv_firmware_add_column_id(FU_CSV_FIRMWARE(img), "vendor_url");
		fu_csv_firmware_set_write_column_ids(FU_CSV_FIRMWARE(img), FALSE);
	} else if (g_strcmp0(sect_id, ".sbatlevel") == 0) {
		img = fu_sbatlevel_section_new();
	} else {
		img = fu_firmware_new();
	}
	fu_firmware_set_id(img, sect_id);
	fu_firmware_set_idx(img, idx);

	/* add data */
	if (fu_struct_pe_coff_section_get_virtual_size(st) > 0) {
		guint32 sect_offset = fu_struct_pe_coff_section_get_pointer_to_raw_data(st);
		guint32 sect_size = fu_struct_pe_coff_section_get_virtual_size(st);
		g_autoptr(GInputStream) img_stream = NULL;

		/* use the raw data size if the section is compressed */
		if (fu_struct_pe_coff_section_get_virtual_size(st) >
		    fu_struct_pe_coff_section_get_size_of_raw_data(st)) {
			g_debug("virtual size 0x%x bigger than raw data, truncating to 0x%x",
				sect_size,
				fu_struct_pe_coff_section_get_size_of_raw_data(st));
			sect_size = fu_struct_pe_coff_section_get_size_of_raw_data(st);
		}

		fu_firmware_set_offset(img, sect_offset);
		img_stream = fu_partial_input_stream_new(stream, sect_offset, sect_size, error);
		if (img_stream == NULL) {
			g_prefix_error_literal(error, "failed to cut raw PE data: ");
			return FALSE;
		}
		if (!fu_firmware_parse_stream(img, img_stream, 0x0, flags, error)) {
			g_prefix_error(error, "failed to parse raw data %s: ", sect_id);
			return FALSE;
		}

		/* add region for Authenticode checksum */
		fu_pefile_firmware_add_region(regions,
					      sect_id,
					      sect_offset,
					      fu_struct_pe_coff_section_get_size_of_raw_data(st));
	}

	/* success */
	return fu_firmware_add_image(FU_FIRMWARE(self), img, error);
}

static gboolean
fu_pefile_firmware_parse(FuFirmware *firmware,
			 GInputStream *stream,
			 FuFirmwareParseFlags flags,
			 GError **error)
{
	FuPefileFirmware *self = FU_PEFILE_FIRMWARE(firmware);
	FuPefileFirmwarePrivate *priv = GET_PRIVATE(self);
	guint32 cert_table_sz = 0;
	gsize offset = 0;
	gsize streamsz = 0;
	gsize strtab_offset;
	guint32 nr_sections;
	g_autoptr(FuStructPeCoffFileHeader) st_coff = NULL;
	g_autoptr(FuStructPeDosHeader) st_doshdr = NULL;
	g_autoptr(GPtrArray) regions = NULL;
	g_autoptr(GInputStream) composite_stream = fu_composite_input_stream_new();

	/* get size */
	if (!fu_input_stream_size(stream, &streamsz, error))
		return FALSE;

	/* parse the DOS header to get the COFF header */
	st_doshdr = fu_struct_pe_dos_header_parse_stream(stream, offset, error);
	if (st_doshdr == NULL) {
		g_prefix_error_literal(error, "failed to read DOS header: ");
		return FALSE;
	}
	offset += fu_struct_pe_dos_header_get_lfanew(st_doshdr);
	st_coff = fu_struct_pe_coff_file_header_parse_stream(stream, offset, error);
	if (st_coff == NULL) {
		g_prefix_error_literal(error, "failed to read COFF header: ");
		return FALSE;
	}
	offset += st_coff->buf->len;

	regions = g_ptr_array_new_with_free_func((GDestroyNotify)fu_pefile_firmware_region_free);

	/* 1st Authenticode region */
	fu_pefile_firmware_add_region(regions,
				      "pre-cksum",
				      0x0,
				      offset + FU_STRUCT_PE_COFF_OPTIONAL_HEADER64_OFFSET_CHECKSUM);

	if (!fu_input_stream_read_safe(
		stream,
		(guint8 *)&priv->subsystem_id,
		sizeof(priv->subsystem_id),
		0x0,
		offset + FU_STRUCT_PE_COFF_OPTIONAL_HEADER64_OFFSET_SUBSYSTEM, /* seek */
		sizeof(priv->subsystem_id),
		error))
		return FALSE;

	/* 2nd Authenticode region */
	fu_pefile_firmware_add_region(
	    regions,
	    "chksum->cert-table",
	    offset + FU_STRUCT_PE_COFF_OPTIONAL_HEADER64_OFFSET_SUBSYSTEM,
	    FU_STRUCT_PE_COFF_OPTIONAL_HEADER64_OFFSET_CERTIFICATE_TABLE -
		FU_STRUCT_PE_COFF_OPTIONAL_HEADER64_OFFSET_SUBSYSTEM); /* end */

	/* verify optional extra header */
	if (fu_struct_pe_coff_file_header_get_size_of_optional_header(st_coff) > 0) {
		g_autoptr(FuStructPeCoffOptionalHeader64) st_opt =
		    fu_struct_pe_coff_optional_header64_parse_stream(stream, offset, error);
		if (st_opt == NULL) {
			g_prefix_error_literal(error, "failed to read optional header: ");
			return FALSE;
		}

		/* 3rd Authenticode region */
		if (fu_struct_pe_coff_optional_header64_get_size_of_headers(st_opt) > 0) {
			fu_pefile_firmware_add_region(
			    regions,
			    "cert-table->end-of-headers",
			    offset + FU_STRUCT_PE_COFF_OPTIONAL_HEADER64_OFFSET_DEBUG_TABLE,
			    fu_struct_pe_coff_optional_header64_get_size_of_headers(st_opt) -
				(offset + FU_STRUCT_PE_COFF_OPTIONAL_HEADER64_OFFSET_DEBUG_TABLE));
		}

		/* 4th Authenticode region */
		cert_table_sz =
		    fu_struct_pe_coff_optional_header64_get_size_of_certificate_table(st_opt);

		offset += fu_struct_pe_coff_file_header_get_size_of_optional_header(st_coff);
	}

	/* read number of sections */
	nr_sections = fu_struct_pe_coff_file_header_get_number_of_sections(st_coff);
	if (nr_sections == 0) {
		g_set_error_literal(error,
				    FWUPD_ERROR,
				    FWUPD_ERROR_INVALID_FILE,
				    "invalid number of sections");
		return FALSE;
	}
	strtab_offset = fu_struct_pe_coff_file_header_get_pointer_to_symbol_table(st_coff) +
			fu_struct_pe_coff_file_header_get_number_of_symbols(st_coff) *
			    FU_STRUCT_PE_COFF_SYMBOL_SIZE;

	/* read out each section */
	for (guint idx = 0; idx < nr_sections; idx++) {
		if (!fu_pefile_firmware_parse_section(self,
						      stream,
						      idx,
						      offset,
						      strtab_offset,
						      regions,
						      flags,
						      error)) {
			g_prefix_error(error, "failed to read section 0x%x: ", idx);
			return FALSE;
		}
		offset += FU_STRUCT_PE_COFF_SECTION_SIZE;
	}

	/* make sure ordered by address */
	g_ptr_array_sort(regions, fu_pefile_firmware_region_sort_cb);

	/* for the data at the end of the image */
	if (regions->len > 0) {
		FuPefileFirmwareRegion *r = g_ptr_array_index(regions, regions->len - 1);
		gsize offset_end = r->offset + r->size;
		fu_pefile_firmware_add_region(regions,
					      "tabledata->cert-table",
					      offset_end,
					      streamsz - (offset_end + cert_table_sz));
	}

	/* calculate the checksum we would find in the dbx */
	for (guint i = 0; i < regions->len; i++) {
		FuPefileFirmwareRegion *r = g_ptr_array_index(regions, i);
		g_autoptr(GInputStream) partial_stream = NULL;

		if (r->size == 0)
			continue;
		g_debug("authenticode region %s: 0x%04x -> 0x%04x [0x%04x]",
			r->name,
			(guint)r->offset,
			(guint)(r->offset + r->size),
			(guint)r->size);
		partial_stream = fu_partial_input_stream_new(stream, r->offset, r->size, error);
		if (partial_stream == NULL) {
			g_prefix_error_literal(error, "failed to cut Authenticode region: ");
			return FALSE;
		}
		fu_composite_input_stream_add_partial_stream(
		    FU_COMPOSITE_INPUT_STREAM(composite_stream),
		    FU_PARTIAL_INPUT_STREAM(partial_stream));
	}
	priv->authenticode_hash =
	    fu_input_stream_compute_checksum(composite_stream, G_CHECKSUM_SHA256, error);
	if (priv->authenticode_hash == NULL)
		return FALSE;

	/* success */
	return TRUE;
}

typedef struct {
	GBytes *blob;
	gchar *id;
	gsize offset;
	gsize blobsz_aligned;
} FuPefileSection;

static void
fu_pefile_firmware_section_free(FuPefileSection *section)
{
	if (section->blob != NULL)
		g_bytes_unref(section->blob);
	g_free(section->id);
	g_free(section);
}

G_DEFINE_AUTOPTR_CLEANUP_FUNC(FuPefileSection, fu_pefile_firmware_section_free)

static GByteArray *
fu_pefile_firmware_write(FuFirmware *firmware, GError **error)
{
	gsize offset = 0;
	g_autoptr(GPtrArray) imgs = fu_firmware_get_images(firmware);
	g_autoptr(FuStructPeDosHeader) st = fu_struct_pe_dos_header_new();
	g_autoptr(FuStructPeCoffFileHeader) st_hdr = fu_struct_pe_coff_file_header_new();
	g_autoptr(FuStructPeCoffOptionalHeader64) st_opt =
	    fu_struct_pe_coff_optional_header64_new();
	g_autoptr(GByteArray) strtab = g_byte_array_new();
	g_autoptr(GPtrArray) sections =
	    g_ptr_array_new_with_free_func((GDestroyNotify)fu_pefile_firmware_section_free);

	/* calculate the offset for each of the sections */
	offset += st->buf->len + st_hdr->buf->len + st_opt->buf->len;
	offset += FU_STRUCT_PE_COFF_SECTION_SIZE * imgs->len;
	for (guint i = 0; i < imgs->len; i++) {
		g_autoptr(FuPefileSection) section = g_new0(FuPefileSection, 1);
		FuFirmware *img = g_ptr_array_index(imgs, i);

		section->offset = offset;
		section->blob = fu_firmware_write(img, error);
		if (section->blob == NULL)
			return NULL;
		if (g_bytes_get_size(section->blob) == 0) {
			g_debug("skipping zero length section %u", i);
			continue;
		}
		section->id = g_strdup(fu_firmware_get_id(img));
		section->blobsz_aligned = fu_common_align_up(g_bytes_get_size(section->blob), 4);
		offset += section->blobsz_aligned;
		g_ptr_array_add(sections, g_steal_pointer(&section));
	}

	/* export_table -> architecture_table */
	fu_struct_pe_coff_optional_header64_set_number_of_rva_and_sizes(st_opt, 7);

	/* COFF file header */
	fu_struct_pe_coff_file_header_set_size_of_optional_header(st_hdr, st_opt->buf->len);
	fu_struct_pe_coff_file_header_set_number_of_sections(st_hdr, sections->len);
	fu_struct_pe_coff_file_header_set_pointer_to_symbol_table(st_hdr, offset);
	fu_byte_array_append_array(st->buf, st_hdr->buf);
	fu_byte_array_append_array(st->buf, st_opt->buf);

	/* add sections */
	for (guint i = 0; i < sections->len; i++) {
		FuPefileSection *section = g_ptr_array_index(sections, i);
		g_autoptr(FuStructPeCoffSection) st_sect = fu_struct_pe_coff_section_new();

		fu_struct_pe_coff_section_set_size_of_raw_data(st_sect,
							       g_bytes_get_size(section->blob));
		fu_struct_pe_coff_section_set_virtual_address(st_sect, 0x0);
		fu_struct_pe_coff_section_set_virtual_size(st_sect, section->blobsz_aligned);
		fu_struct_pe_coff_section_set_pointer_to_raw_data(st_sect, section->offset);

		/* set the name directly, or add to the string table */
		if (section->id == NULL) {
			g_set_error(error,
				    FWUPD_ERROR,
				    FWUPD_ERROR_INVALID_DATA,
				    "image %u has no ID",
				    i);
			return NULL;
		}
		if (strlen(section->id) <= 8) {
			if (!fu_struct_pe_coff_section_set_name(st_sect, section->id, error))
				return NULL;
		} else {
			g_autofree gchar *name_tmp = g_strdup_printf("/%u", strtab->len);
			g_autoptr(GByteArray) strtab_buf = g_byte_array_new();

			if (!fu_struct_pe_coff_section_set_name(st_sect, name_tmp, error))
				return NULL;

			/* create a byte buffer of exactly the correct chunk size */
			g_byte_array_append(strtab_buf,
					    (const guint8 *)section->id,
					    strlen(section->id));
			if (strtab_buf->len > FU_PEFILE_SECTION_ID_STRTAB_SIZE) {
				g_set_error(error,
					    FWUPD_ERROR,
					    FWUPD_ERROR_INVALID_DATA,
					    "image ID %s is too long",
					    section->id);
				return NULL;
			}
			fu_byte_array_set_size(strtab_buf, FU_PEFILE_SECTION_ID_STRTAB_SIZE, 0x0);
			g_byte_array_append(strtab, strtab_buf->data, strtab_buf->len);
		}
		fu_byte_array_append_array(st->buf, st_sect->buf);
	}

	/* add the section data itself */
	for (guint i = 0; i < sections->len; i++) {
		FuPefileSection *section = g_ptr_array_index(sections, i);
		g_autoptr(GBytes) blob_aligned =
		    fu_bytes_pad(section->blob, section->blobsz_aligned, 0xFF);
		fu_byte_array_append_bytes(st->buf, blob_aligned);
	}

	/* string table comes last */
	g_byte_array_append(st->buf, strtab->data, strtab->len);

	/* success */
	return g_steal_pointer(&st->buf);
}

static gchar *
fu_pefile_firmware_get_checksum(FuFirmware *firmware, GChecksumType csum_kind, GError **error)
{
	FuPefileFirmware *self = FU_PEFILE_FIRMWARE(firmware);
	FuPefileFirmwarePrivate *priv = GET_PRIVATE(self);
	if (csum_kind != G_CHECKSUM_SHA256) {
		g_set_error_literal(error,
				    FWUPD_ERROR,
				    FWUPD_ERROR_NOT_SUPPORTED,
				    "Authenticode only supports SHA256");
		return NULL;
	}
	if (priv->authenticode_hash == NULL) {
		g_set_error_literal(error,
				    FWUPD_ERROR,
				    FWUPD_ERROR_INVALID_DATA,
				    "Authenticode checksum not set");
		return NULL;
	}
	return g_strdup(priv->authenticode_hash);
}

static void
fu_pefile_firmware_init(FuPefileFirmware *self)
{
	fu_firmware_add_image_gtype(FU_FIRMWARE(self), FU_TYPE_FIRMWARE);
	fu_firmware_add_image_gtype(FU_FIRMWARE(self), FU_TYPE_CSV_FIRMWARE);
	fu_firmware_add_image_gtype(FU_FIRMWARE(self), FU_TYPE_SBATLEVEL_SECTION);
	fu_firmware_set_images_max(FU_FIRMWARE(self), 100);
}

static void
fu_pefile_firmware_finalize(GObject *object)
{
	FuPefileFirmware *self = FU_PEFILE_FIRMWARE(object);
	FuPefileFirmwarePrivate *priv = GET_PRIVATE(self);
	g_free(priv->authenticode_hash);
	G_OBJECT_CLASS(fu_pefile_firmware_parent_class)->finalize(object);
}

static void
fu_pefile_firmware_class_init(FuPefileFirmwareClass *klass)
{
	FuFirmwareClass *firmware_class = FU_FIRMWARE_CLASS(klass);
	GObjectClass *object_class = G_OBJECT_CLASS(klass);
	object_class->finalize = fu_pefile_firmware_finalize;
	firmware_class->validate = fu_pefile_firmware_validate;
	firmware_class->parse = fu_pefile_firmware_parse;
	firmware_class->write = fu_pefile_firmware_write;
	firmware_class->export = fu_pefile_firmware_export;
	firmware_class->get_checksum = fu_pefile_firmware_get_checksum;
}

/**
 * fu_pefile_firmware_new:
 *
 * Creates a new #FuPefileFirmware
 *
 * Since: 1.8.10
 **/
FuFirmware *
fu_pefile_firmware_new(void)
{
	return FU_FIRMWARE(g_object_new(FU_TYPE_PEFILE_FIRMWARE, NULL));
}
