/* PE parser for POSIX
   Copyright (C) 2026 Alexander Larsson <alexl@redhat.com>

   This file is free software: you can redistribute it and/or modify
   it under the terms of the GNU Lesser General Public License as
   published by the Free Software Foundation; either version 2.1 of the
   License, or (at your option) any later version.

   This file is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU Lesser General Public License for more details.

   You should have received a copy of the GNU Lesser General Public License
   along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/
#include "config.h"

#include "pe_parser_posix.h"
#include "utils.h"

#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#ifdef HAVE_ENDIAN_H
#include <endian.h>
#endif
#ifdef HAVE_SYS_ENDIAN_H
#include <sys/endian.h>
#endif
#ifdef HAVE_MACHINE_ENDIAN_H
#include <machine/endian.h>
#endif

#define IMAGE_DOS_SIGNATURE 0x5A4D
#define IMAGE_NT_SIGNATURE  0x00004550

typedef struct attr_packed {
    uint16_t e_magic;
    uint16_t e_cblp;
    uint16_t e_cp;
    uint16_t e_crlc;
    uint16_t e_cparhdr;
    uint16_t e_minalloc;
    uint16_t e_maxalloc;
    uint16_t e_ss;
    uint16_t e_sp;
    uint16_t e_csum;
    uint16_t e_ip;
    uint16_t e_cs;
    uint16_t e_lfarlc;
    uint16_t e_ovno;
    uint16_t e_res[4];
    uint16_t e_oemid;
    uint16_t e_oeminfo;
    uint16_t e_res2[10];
    int32_t e_lfanew;
} ImageDosHeader;

typedef struct attr_packed {
    uint16_t Machine;
    uint16_t NumberOfSections;
    uint32_t TimeDateStamp;
    uint32_t PointerToSymbolTable;
    uint32_t NumberOfSymbols;
    uint16_t SizeOfOptionalHeader;
    uint16_t Characteristics;
} ImageFileHeader;

typedef struct attr_packed {
    uint32_t Signature;
    ImageFileHeader FileHeader;
} ImageNtHeaders;

typedef struct attr_packed {
    uint8_t Name[PE_SECTION_NAME_SIZE];
    union {
	uint32_t PhysicalAddress;
	uint32_t VirtualSize;
    } Misc;
    uint32_t VirtualAddress;
    uint32_t SizeOfRawData;
    uint32_t PointerToRawData;
    uint32_t PointerToRelocations;
    uint32_t PointerToLinenumbers;
    uint16_t NumberOfRelocations;
    uint16_t NumberOfLinenumbers;
    uint32_t Characteristics;
} ImageSectionHeader;

#define IMAGE_NT_OPTIONAL_HDR32_MAGIC 0x10b
#define IMAGE_NT_OPTIONAL_HDR64_MAGIC 0x20b

static bool
pe_read_headers(int fd,
		off_t base_offset,
		uint32_t *nt_offset_out,
		ImageNtHeaders *nt_out,
		uint16_t *opt_magic_out)
{
    ImageDosHeader dos_header;
    if (!full_pread(fd, &dos_header, sizeof(dos_header), base_offset))
	return false;

    if (le16toh(dos_header.e_magic) != IMAGE_DOS_SIGNATURE) {
	error("Invalid DOS header magic\n");
	return false;
    }

    uint32_t nt_offset = (uint32_t)le32toh(dos_header.e_lfanew);

    ImageNtHeaders nt_header;
    if (!full_pread(fd, &nt_header, sizeof(nt_header), base_offset + nt_offset))
	return false;

    if (le32toh(nt_header.Signature) != IMAGE_NT_SIGNATURE) {
	error("Invalid PE signature\n");
	return false;
    }

    uint16_t opt_magic;
    if (!full_pread(fd,
		    &opt_magic,
		    sizeof(opt_magic),
		    base_offset + nt_offset + sizeof(ImageNtHeaders)))
	return false;
    opt_magic = le16toh(opt_magic);
    if (opt_magic != IMAGE_NT_OPTIONAL_HDR32_MAGIC && opt_magic != IMAGE_NT_OPTIONAL_HDR64_MAGIC) {
	error("Unsupported PE optional header magic: 0x%x\n", opt_magic);
	return false;
    }

    *nt_offset_out = nt_offset;
    *nt_out = nt_header;
    *opt_magic_out = opt_magic;
    return true;
}

bool
pe_parse_sections(int fd, off_t base_offset, PeSectionInfo **sections_out, uint32_t *n_sections_out)
{
    uint32_t nt_offset;
    ImageNtHeaders nt_header;
    uint16_t opt_magic;
    if (!pe_read_headers(fd, base_offset, &nt_offset, &nt_header, &opt_magic))
	return false;

    uint16_t opt_header_size = le16toh(nt_header.FileHeader.SizeOfOptionalHeader);
    uint32_t num_sections = le16toh(nt_header.FileHeader.NumberOfSections);
    uint32_t sections_offset = nt_offset + sizeof(ImageNtHeaders) + opt_header_size;

    autofree ImageSectionHeader *raw_sections = xcalloc(num_sections, sizeof(ImageSectionHeader));

    if (!full_pread(fd,
		    raw_sections,
		    num_sections * sizeof(ImageSectionHeader),
		    base_offset + sections_offset))
	return false;

    autofree PeSectionInfo *sections = xcalloc(num_sections, sizeof(PeSectionInfo));

    for (uint32_t i = 0; i < num_sections; i++) {
	memcpy(sections[i].name, raw_sections[i].Name, PE_SECTION_NAME_SIZE);
	sections[i].name[PE_SECTION_NAME_SIZE] = '\0';
	sections[i].file_offset = le32toh(raw_sections[i].PointerToRawData);
	sections[i].size = le32toh(raw_sections[i].Misc.VirtualSize);
    }

    *sections_out = steal_pointer(&sections);
    *n_sections_out = num_sections;
    return true;
}

/* Authenticode signature verification */

#ifdef HAVE_OPENSSL

#include <openssl/bio.h>
#include <openssl/crypto.h>
#include <openssl/err.h>
#include <openssl/evp.h>
#include <openssl/objects.h>
#include <openssl/pem.h>
#include <openssl/pkcs7.h>
#include <openssl/x509.h>

#define WIN_CERT_REVISION_2_0 0x0200
#define WIN_CERT_TYPE_PKCS7   0x0002
#define HASH_BUF_SIZE	      (64 * 1024)

static inline void
cleanup_x509_storep(X509_STORE **p)
{
    if (*p)
	X509_STORE_free(*p);
}
static inline void
cleanup_x509_store_ctxp(X509_STORE_CTX **p)
{
    if (*p)
	X509_STORE_CTX_free(*p);
}
static inline void
cleanup_evp_md_ctxp(EVP_MD_CTX **p)
{
    if (*p)
	EVP_MD_CTX_free(*p);
}
static inline void
cleanup_pkcs7p(PKCS7 **p)
{
    if (*p)
	PKCS7_free(*p);
}

#define auto_x509_store	    __attribute__((cleanup(cleanup_x509_storep)))
#define auto_x509_store_ctx __attribute__((cleanup(cleanup_x509_store_ctxp)))
#define auto_evp_md_ctx	    __attribute__((cleanup(cleanup_evp_md_ctxp)))
#define auto_pkcs7	    __attribute__((cleanup(cleanup_pkcs7p)))

typedef struct attr_packed {
    uint32_t dwLength;
    uint16_t wRevision;
    uint16_t wCertificateType;
} WinCertificate;

typedef struct {
    off_t base_offset;
    off_t opt_header_offset;
    bool is_pe32plus;
    off_t checksum_offset;
    off_t security_dir_offset;
    uint32_t cert_table_rva;
    uint32_t cert_table_size;
    uint32_t size_of_headers;
    uint16_t num_sections;
    off_t sections_offset;
} PeLayout;

static bool
pe_read_layout(int fd, off_t base_offset, PeLayout *layout)
{
    uint32_t nt_offset;
    ImageNtHeaders nt_header;
    uint16_t opt_magic;
    if (!pe_read_headers(fd, base_offset, &nt_offset, &nt_header, &opt_magic))
	return false;

    layout->base_offset = base_offset;
    layout->opt_header_offset = base_offset + nt_offset + sizeof(ImageNtHeaders);
    layout->num_sections = le16toh(nt_header.FileHeader.NumberOfSections);
    uint16_t opt_size = le16toh(nt_header.FileHeader.SizeOfOptionalHeader);
    layout->sections_offset = layout->opt_header_offset + opt_size;
    layout->is_pe32plus = (opt_magic == IMAGE_NT_OPTIONAL_HDR64_MAGIC);
    layout->checksum_offset = layout->opt_header_offset + 64;

    uint32_t security_dir_table_offset = layout->is_pe32plus ? 144 : 128;
    layout->security_dir_offset = layout->opt_header_offset + security_dir_table_offset;

    uint32_t sec_dir[2];
    if (!full_pread(fd, sec_dir, sizeof(sec_dir), layout->security_dir_offset))
	return false;
    layout->cert_table_rva = le32toh(sec_dir[0]);
    layout->cert_table_size = le32toh(sec_dir[1]);

    uint32_t size_of_headers;
    if (!full_pread(fd, &size_of_headers, sizeof(size_of_headers), layout->opt_header_offset + 60))
	return false;
    layout->size_of_headers = le32toh(size_of_headers);

    return true;
}

static int
section_raw_ptr_cmp(const void *a, const void *b)
{
    const ImageSectionHeader *sa = a;
    const ImageSectionHeader *sb = b;
    uint32_t pa = le32toh(sa->PointerToRawData);
    uint32_t pb = le32toh(sb->PointerToRawData);
    return (pa > pb) - (pa < pb);
}

static bool
hash_file_range(EVP_MD_CTX *ctx, int fd, off_t offset, size_t length)
{
    autofree uint8_t *buf = xmalloc(HASH_BUF_SIZE);

    while (length > 0) {
	size_t chunk = length < HASH_BUF_SIZE ? length : HASH_BUF_SIZE;
	if (!full_pread(fd, buf, chunk, offset))
	    return false;
	if (!EVP_DigestUpdate(ctx, buf, chunk))
	    return false;
	offset += chunk;
	length -= chunk;
    }
    return true;
}

static bool
compute_authenticode_hash(int fd,
			  PeLayout *layout,
			  const EVP_MD *md,
			  uint8_t *hash_out,
			  unsigned int *hash_len)
{
    auto_evp_md_ctx EVP_MD_CTX *ctx = EVP_MD_CTX_new();
    if (!ctx)
	return false;

    if (!EVP_DigestInit_ex(ctx, md, NULL))
	return false;

    off_t base = layout->base_offset;
    off_t checksum_off = layout->checksum_offset;
    off_t secdir_off = layout->security_dir_offset;
    off_t after_checksum = checksum_off + 4;
    off_t after_secdir = secdir_off + 8;
    off_t headers_end = base + layout->size_of_headers;

    /* Hash headers in three ranges, skipping the CheckSum field
       (4 bytes) and the Security Directory entry (8 bytes). */
    if (!hash_file_range(ctx, fd, base, checksum_off - base))
	return false;

    if (!hash_file_range(ctx, fd, after_checksum, secdir_off - after_checksum))
	return false;

    if (after_secdir < headers_end) {
	if (!hash_file_range(ctx, fd, after_secdir, headers_end - after_secdir))
	    return false;
    }

    autofree ImageSectionHeader *shdrs = xcalloc(layout->num_sections, sizeof(ImageSectionHeader));
    if (!full_pread(fd,
		    shdrs,
		    layout->num_sections * sizeof(ImageSectionHeader),
		    layout->sections_offset))
	return false;

    qsort(shdrs, layout->num_sections, sizeof(ImageSectionHeader), section_raw_ptr_cmp);

    off_t last_section_end = 0;
    for (uint16_t i = 0; i < layout->num_sections; i++) {
	uint32_t raw_ptr = le32toh(shdrs[i].PointerToRawData);
	uint32_t raw_size = le32toh(shdrs[i].SizeOfRawData);
	if (raw_size == 0)
	    continue;
	if (!hash_file_range(ctx, fd, base + raw_ptr, raw_size))
	    return false;
	off_t sec_end = raw_ptr + raw_size;
	if (sec_end > last_section_end)
	    last_section_end = sec_end;
    }

    if (layout->cert_table_rva > 0 && layout->cert_table_rva > last_section_end) {
	if (!hash_file_range(ctx,
			     fd,
			     base + last_section_end,
			     layout->cert_table_rva - last_section_end))
	    return false;
    }

    if (!EVP_DigestFinal_ex(ctx, hash_out, hash_len))
	return false;

    return true;
}

/* Extract the file hash from a DER-encoded SpcIndirectDataContent:
   SEQUENCE { SpcAttributeTypeAndOptionalValue, DigestInfo }
   where DigestInfo = SEQUENCE { AlgorithmIdentifier, OCTET STRING } */
static bool
extract_spc_file_hash(const uint8_t *data,
		      int data_len,
		      uint8_t *hash_out,
		      unsigned int *hash_len_out,
		      unsigned int hash_max)
{
    const uint8_t *p = data;
    long len;
    int tag, cls, ret;

    /* Outer SEQUENCE */
    ret = ASN1_get_object(&p, &len, &tag, &cls, data_len);
    if (ret & 0x80 || tag != V_ASN1_SEQUENCE)
	return false;
    const uint8_t *seq_end = p + len;

    /* Skip SpcAttributeTypeAndOptionalValue SEQUENCE */
    ret = ASN1_get_object(&p, &len, &tag, &cls, seq_end - p);
    if (ret & 0x80 || tag != V_ASN1_SEQUENCE)
	return false;
    p += len;

    /* DigestInfo SEQUENCE */
    ret = ASN1_get_object(&p, &len, &tag, &cls, seq_end - p);
    if (ret & 0x80 || tag != V_ASN1_SEQUENCE)
	return false;
    const uint8_t *di_end = p + len;

    /* Skip AlgorithmIdentifier SEQUENCE */
    ret = ASN1_get_object(&p, &len, &tag, &cls, di_end - p);
    if (ret & 0x80 || tag != V_ASN1_SEQUENCE)
	return false;
    p += len;

    /* OCTET STRING: the file hash */
    ret = ASN1_get_object(&p, &len, &tag, &cls, di_end - p);
    if (ret & 0x80 || tag != V_ASN1_OCTET_STRING)
	return false;
    if (len <= 0 || (unsigned int)len > hash_max)
	return false;

    memcpy(hash_out, p, len);
    *hash_len_out = len;
    return true;
}

static bool
verify_pkcs7_signature(PKCS7 *p7,
		       PKCS7_SIGNER_INFO *si,
		       const EVP_MD *md,
		       const uint8_t *content_data,
		       int content_len,
		       const char *cert_path)
{
    FILE *cert_file = fopen(cert_path, "r");
    if (!cert_file) {
	error("Failed to open certificate %s: %s\n", cert_path, strerror(errno));
	return false;
    }

    X509 *trusted = PEM_read_X509(cert_file, NULL, NULL, NULL);
    fclose(cert_file);
    if (!trusted) {
	error("Failed to parse PEM certificate %s\n", cert_path);
	return false;
    }

    auto_x509_store X509_STORE *store = X509_STORE_new();
    if (!store) {
	X509_free(trusted);
	return false;
    }
    X509_STORE_add_cert(store, trusted);
    X509_free(trusted);

    /* Match EFI firmware behavior: don't check cert expiry, since
       there is no reliable clock at early boot. */
    X509_STORE_set_flags(store, X509_V_FLAG_NO_CHECK_TIME);

    /* We can't use PKCS7_verify() because it doesn't handle the
       Authenticode SpcIndirectDataContent content type.  Instead
       we verify the cert chain, messageDigest, and signature manually. */

    /* Verify signer cert against trusted store */
    X509 *signer = PKCS7_cert_from_signer_info(p7, si);
    if (!signer) {
	error("Cannot find signer certificate in PKCS#7\n");
	return false;
    }

    auto_x509_store_ctx X509_STORE_CTX *store_ctx = X509_STORE_CTX_new();
    if (!store_ctx)
	return false;

    if (!X509_STORE_CTX_init(store_ctx, store, signer, p7->d.sign->cert))
	return false;

    if (X509_verify_cert(store_ctx) != 1) {
	error("Signer certificate verification failed\n");
	return false;
    }

    verbose("Signer certificate trusted\n");

    /* Verify the messageDigest attribute matches hash of content */
    ASN1_OCTET_STRING *msg_digest = PKCS7_digest_from_attributes(si->auth_attr);
    if (!msg_digest) {
	error("No messageDigest attribute in signer info\n");
	return false;
    }

    uint8_t computed_md[EVP_MAX_MD_SIZE];
    unsigned int computed_md_len;
    if (!EVP_Digest(content_data, content_len, computed_md, &computed_md_len, md, NULL)) {
	error("Failed to compute content digest\n");
	return false;
    }

    if (msg_digest->length != (int)computed_md_len ||
	CRYPTO_memcmp(msg_digest->data, computed_md, computed_md_len) != 0) {
	error("Content digest mismatch\n");
	return false;
    }

    /* Verify the RSA/ECDSA signature over the authenticated attributes */
    EVP_PKEY *pkey = X509_get0_pubkey(signer);
    if (!pkey) {
	error("Cannot get signer public key\n");
	return false;
    }

    int attr_der_len =
	ASN1_item_i2d((ASN1_VALUE *)si->auth_attr, NULL, ASN1_ITEM_rptr(PKCS7_ATTR_VERIFY));
    if (attr_der_len <= 0) {
	error("Failed to encode authenticated attributes\n");
	return false;
    }

    autofree uint8_t *attr_der = xmalloc(attr_der_len);
    uint8_t *attr_p = attr_der;
    ASN1_item_i2d((ASN1_VALUE *)si->auth_attr, &attr_p, ASN1_ITEM_rptr(PKCS7_ATTR_VERIFY));

    auto_evp_md_ctx EVP_MD_CTX *md_ctx = EVP_MD_CTX_new();
    if (!md_ctx)
	return false;

    if (!EVP_DigestVerifyInit(md_ctx, NULL, md, NULL, pkey) ||
	!EVP_DigestVerify(md_ctx,
			  si->enc_digest->data,
			  si->enc_digest->length,
			  attr_der,
			  attr_der_len)) {
	error("Signature verification failed\n");
	return false;
    }

    return true;
}

bool
pe_verify_signature(int fd, off_t base_offset, const char *cert_path)
{
    PeLayout layout;
    if (!pe_read_layout(fd, base_offset, &layout))
	return false;

    if (layout.cert_table_size == 0) {
	error("PE file has no embedded signature\n");
	return false;
    }

    off_t cert_offset = base_offset + layout.cert_table_rva;
    WinCertificate win_cert;

    if (!full_pread(fd, &win_cert, sizeof(win_cert), cert_offset))
	return false;

    if (le16toh(win_cert.wRevision) != WIN_CERT_REVISION_2_0) {
	error("Unsupported certificate revision: 0x%04x\n", le16toh(win_cert.wRevision));
	return false;
    }
    if (le16toh(win_cert.wCertificateType) != WIN_CERT_TYPE_PKCS7) {
	error("Unsupported certificate type: 0x%04x\n", le16toh(win_cert.wCertificateType));
	return false;
    }

    uint32_t dwLength = le32toh(win_cert.dwLength);
    if (dwLength < sizeof(WinCertificate) || dwLength > layout.cert_table_size) {
	error("Invalid certificate length: %u\n", dwLength);
	return false;
    }

    uint32_t pkcs7_size = dwLength - sizeof(WinCertificate);
    autofree uint8_t *pkcs7_buf = xmalloc(pkcs7_size);
    if (!full_pread(fd, pkcs7_buf, pkcs7_size, cert_offset + sizeof(win_cert)))
	return false;

    BIO *bio = BIO_new_mem_buf(pkcs7_buf, pkcs7_size);
    if (!bio) {
	error("Failed to create BIO\n");
	return false;
    }

    auto_pkcs7 PKCS7 *p7 = d2i_PKCS7_bio(bio, NULL);
    BIO_free(bio);
    if (!p7) {
	error("Failed to parse PKCS#7 signature\n");
	return false;
    }

    if (!PKCS7_type_is_signed(p7)) {
	error("PKCS#7 is not SignedData\n");
	return false;
    }

    STACK_OF(PKCS7_SIGNER_INFO) *sinfos = PKCS7_get_signer_info(p7);
    if (!sinfos || sk_PKCS7_SIGNER_INFO_num(sinfos) != 1) {
	error("Expected exactly one signer in PKCS#7\n");
	return false;
    }
    PKCS7_SIGNER_INFO *si = sk_PKCS7_SIGNER_INFO_value(sinfos, 0);

    const EVP_MD *md = EVP_get_digestbyobj(si->digest_alg->algorithm);
    if (!md) {
	error("Unknown digest algorithm in signature\n");
	return false;
    }

    verbose("Signature digest algorithm: %s\n", EVP_MD_get0_name(md));

    /* Extract expected file hash from SpcIndirectDataContent */
    ASN1_TYPE *content_val = p7->d.sign->contents->d.other;
    if (!content_val || content_val->type != V_ASN1_SEQUENCE) {
	error("Missing or invalid SpcIndirectDataContent\n");
	return false;
    }
    ASN1_OCTET_STRING *content_os = content_val->value.sequence;
    if (!content_os || !content_os->data) {
	error("Empty SpcIndirectDataContent\n");
	return false;
    }
    uint8_t expected_hash[EVP_MAX_MD_SIZE];
    unsigned int expected_hash_len;
    if (!extract_spc_file_hash(content_os->data,
			       content_os->length,
			       expected_hash,
			       &expected_hash_len,
			       sizeof(expected_hash))) {
	error("Failed to extract file hash from signature\n");
	return false;
    }

    /* Compute Authenticode hash and compare */
    uint8_t computed_hash[EVP_MAX_MD_SIZE];
    unsigned int computed_hash_len;
    if (!compute_authenticode_hash(fd, &layout, md, computed_hash, &computed_hash_len))
	return false;

    if (computed_hash_len != expected_hash_len ||
	CRYPTO_memcmp(computed_hash, expected_hash, computed_hash_len) != 0) {
	error("Authenticode hash mismatch\n");
	return false;
    }

    verbose("Authenticode hash matches\n");

    /* Verify the PKCS#7 signature and certificate chain.
       The messageDigest attribute hashes the content value bytes
       (inside the outer SEQUENCE tag, not including it). */
    const uint8_t *content_inner = content_os->data;
    int content_inner_len = content_os->length;
    long inner_len;
    int tag, cls;
    int ret = ASN1_get_object(&content_inner, &inner_len, &tag, &cls, content_inner_len);
    if (ret & 0x80 || tag != V_ASN1_SEQUENCE) {
	error("Invalid SpcIndirectDataContent encoding\n");
	return false;
    }

    return verify_pkcs7_signature(p7, si, md, content_inner, inner_len, cert_path);
}

#else /* !HAVE_OPENSSL */

bool
pe_verify_signature(arg_unused int fd,
		    arg_unused off_t base_offset,
		    arg_unused const char *cert_path)
{
    error("Signature verification not available (built without OpenSSL)\n");
    return false;
}

#endif /* HAVE_OPENSSL */
