/* ukibootimg
   Copyright (C) 2025 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 "bootctl.h"
#include "gpt_parser.h"
#include "pe_parser_posix.h"

#include <fcntl.h>
#include <getopt.h>
#include <inttypes.h>
#include <stdio.h>
#include <sys/stat.h>

#define OPT_HELP       256
#define OPT_VERSION    257
#define OPT_VERBOSE    258
#define OPT_IMAGE      259
#define OPT_BOOTCTL    260
#define OPT_BOOT_A     261
#define OPT_BOOT_B     262
#define OPT_SLOT       263
#define OPT_OUTPUT_DIR 264
#define OPT_CERT       265
#define OPT_VERIFY     266

static const char *opt_image = NULL;
static const char *opt_bootctl_path = NULL;
static const char *opt_boot_a = NULL;
static const char *opt_boot_b = NULL;
static const char *opt_output_dir = ".";
static const char *opt_cert = NULL;
static bool opt_verify = false;
static int opt_slot = -1;

static bool
read_section_content(int fd, off_t base_offset, PeSectionInfo *section, char **out, size_t *out_len)
{
    autofree char *buf = xmalloc(section->size + 1);

    if (!full_pread(fd, buf, section->size, base_offset + section->file_offset))
	return false;

    buf[section->size] = '\0';
    *out = steal_pointer(&buf);
    if (out_len)
	*out_len = section->size;
    return true;
}

static PeSectionInfo *
find_section(PeSectionInfo *sections, uint32_t n_sections, const char *name)
{
    for (uint32_t i = 0; i < n_sections; i++) {
	if (strcmp(sections[i].name, name) == 0)
	    return &sections[i];
    }
    return NULL;
}

static char *
find_osrel_value(const char *osrel, const char *key)
{
    size_t key_len = strlen(key);
    for (const char *line = osrel; line;) {
	if (strncmp(line, key, key_len) == 0 && line[key_len] == '=') {
	    const char *val = line + key_len + 1;
	    const char *end = strchr(val, '\n');
	    if (!end)
		end = val + strlen(val);
	    if (end > val && val[0] == '"' && end[-1] == '"')
		return strndup(val + 1, end - val - 2);
	    return strndup(val, end - val);
	}
	const char *next = strchr(line, '\n');
	line = next ? next + 1 : NULL;
    }
    return NULL;
}

static int
cmd_status(int argc,
	   arg_unused char **argv,
	   int uki_fd,
	   off_t uki_offset,
	   uint32_t active_slot,
	   const char *uki_path,
	   PeSectionInfo *sections,
	   uint32_t n_sections)
{
    if (argc > 0) {
	error("Unexpected arguments to status\n");
	return 1;
    }

    printf("Active slot: %c\n", active_slot == 0 ? 'a' : 'b');

    PeSectionInfo *osrel_sec = find_section(sections, n_sections, ".osrel");
    if (osrel_sec) {
	autofree char *osrel = NULL;
	if (read_section_content(uki_fd, uki_offset, osrel_sec, &osrel, NULL)) {
	    autofree char *name = find_osrel_value(osrel, "NAME");
	    autofree char *version = find_osrel_value(osrel, "VERSION");
	    if (name && version)
		printf("OS: %s %s\n", name, version);
	    else if (name)
		printf("OS: %s\n", name);
	}
    }

    PeSectionInfo *uname_sec = find_section(sections, n_sections, ".uname");
    if (uname_sec) {
	autofree char *uname = NULL;
	if (read_section_content(uki_fd, uki_offset, uname_sec, &uname, NULL))
	    printf("Kernel version: %s\n", uname);
    }

    PeSectionInfo *cmdline_sec = find_section(sections, n_sections, ".cmdline");
    if (cmdline_sec) {
	autofree char *cmdline = NULL;
	if (read_section_content(uki_fd, uki_offset, cmdline_sec, &cmdline, NULL))
	    printf("Command line: %s androidboot.slot_suffix=_%c\n",
		   cmdline,
		   active_slot == 0 ? 'a' : 'b');
    }

    PeSectionInfo *linux_sec = find_section(sections, n_sections, ".linux");
    if (linux_sec)
	printf("Kernel: %s offset=0x%" PRIx64 " size=%u\n",
	       uki_path,
	       (uint64_t)(uki_offset + linux_sec->file_offset),
	       linux_sec->size);

    PeSectionInfo *initrd_sec = find_section(sections, n_sections, ".initrd");
    if (initrd_sec)
	printf("Initrd: %s offset=0x%" PRIx64 " size=%u\n",
	       uki_path,
	       (uint64_t)(uki_offset + initrd_sec->file_offset),
	       initrd_sec->size);

    return 0;
}

static int
cmd_list(int argc,
	 arg_unused char **argv,
	 arg_unused int uki_fd,
	 arg_unused off_t uki_offset,
	 arg_unused uint32_t active_slot,
	 arg_unused const char *uki_path,
	 PeSectionInfo *sections,
	 uint32_t n_sections)
{
    if (argc > 0) {
	error("Unexpected arguments to list\n");
	return 1;
    }

    printf("%-16s %10s %10s\n", "Name", "Offset", "Size");
    for (uint32_t i = 0; i < n_sections; i++) {
	printf("%-16s 0x%08x %10u\n", sections[i].name, sections[i].file_offset, sections[i].size);
    }
    return 0;
}

#define COPY_BUF_SIZE (64 * 1024)

static bool
extract_section(int fd,
		off_t base_offset,
		PeSectionInfo *section,
		const char *output_dir,
		uint32_t active_slot)
{
    const char *name = section->name;
    if (name[0] == '.')
	name++;

    autofree char *path = NULL;
    if (asprintf(&path, "%s/%s", output_dir, name) < 0)
	oom();

    autoclose int out_fd = open(path, O_WRONLY | O_CREAT | O_TRUNC | O_CLOEXEC, 0644);
    if (out_fd < 0) {
	error("Failed to create %s: %s\n", path, strerror(errno));
	return false;
    }

    autofree uint8_t *buf = xmalloc(COPY_BUF_SIZE);

    uint32_t remaining = section->size;
    off_t offset = base_offset + section->file_offset;

    while (remaining > 0) {
	size_t chunk = remaining < COPY_BUF_SIZE ? remaining : COPY_BUF_SIZE;
	ssize_t n = pread(fd, buf, chunk, offset);
	if (n <= 0) {
	    error("Failed to read section %s: %s\n",
		  section->name,
		  n < 0 ? strerror(errno) : "unexpected end of file");
	    return false;
	}
	if (!full_write(out_fd, buf, n))
	    return false;
	remaining -= n;
	offset += n;
    }

    if (strcmp(section->name, ".cmdline") == 0) {
	const char *suffix = active_slot == 0 ? " androidboot.slot_suffix=_a"
					      : " androidboot.slot_suffix=_b";
	if (!full_write(out_fd, suffix, strlen(suffix)))
	    return false;
    }

    verbose("Extracted %s (%u bytes)\n", path, section->size);
    return true;
}

static int
cmd_verify(int argc,
	   arg_unused char **argv,
	   int uki_fd,
	   off_t uki_offset,
	   arg_unused uint32_t active_slot,
	   arg_unused const char *uki_path,
	   arg_unused PeSectionInfo *sections,
	   arg_unused uint32_t n_sections)
{
    if (argc > 0) {
	error("Unexpected arguments to verify\n");
	return 1;
    }

    if (!opt_cert)
	fatal("--cert is required for verify\n");

    if (!pe_verify_signature(uki_fd, uki_offset, opt_cert))
	return 1;

    printf("Signature verification successful\n");
    return 0;
}

static int
cmd_extract(int argc,
	    char **argv,
	    int uki_fd,
	    off_t uki_offset,
	    uint32_t active_slot,
	    arg_unused const char *uki_path,
	    PeSectionInfo *sections,
	    uint32_t n_sections)
{
    if (opt_verify) {
	if (!opt_cert)
	    fatal("--cert is required with --verify\n");
	if (!pe_verify_signature(uki_fd, uki_offset, opt_cert))
	    fatal("Signature verification failed, aborting extract\n");
	verbose("Signature verified\n");
    }

    bool extract_all = (argc == 0);

    for (uint32_t i = 0; i < n_sections; i++) {
	bool should_extract = extract_all;

	if (!extract_all) {
	    for (int a = 0; a < argc; a++) {
		const char *want = argv[a];
		if (strcmp(sections[i].name, want) == 0 ||
		    (want[0] != '.' && sections[i].name[0] == '.' &&
		     strcmp(sections[i].name + 1, want) == 0)) {
		    should_extract = true;
		    break;
		}
	    }
	}

	if (should_extract) {
	    if (!extract_section(uki_fd, uki_offset, &sections[i], opt_output_dir, active_slot))
		return 1;
	}
    }

    return 0;
}

static bool
resolve_active_slot(int bootctl_fd, off_t bootctl_offset, uint32_t *slot_out)
{
    UkiBootCtrl ctrl;
    ssize_t n = pread(bootctl_fd, &ctrl, sizeof(ctrl), bootctl_offset);
    if (n < (ssize_t)sizeof(ctrl)) {
	error("Failed to read bootctl: %s\n", n < 0 ? strerror(errno) : "short read");
	return false;
    }

    if (!bootctl_validate(&ctrl)) {
	verbose("Bootctl invalid, using default initialization\n");
	bootctl_init_default(&ctrl);
    }

    if (!bootctl_find_active_slot(&ctrl, slot_out)) {
	error("No valid active slot found\n");
	return false;
    }

    return true;
}

typedef int (*CmdFn)(int argc,
		     char **argv,
		     int uki_fd,
		     off_t uki_offset,
		     uint32_t active_slot,
		     const char *uki_path,
		     PeSectionInfo *sections,
		     uint32_t n_sections);

typedef struct {
    const char *name;
    CmdFn fn;
} CmdHandler;

static CmdHandler commands[] = {
    {"extract", cmd_extract},
    {"list", cmd_list},
    {"status", cmd_status},
    {"verify", cmd_verify},
};

static CmdHandler *
find_command_handler(const char *name)
{
    for (size_t i = 0; i < N_ELEMENTS(commands); i++) {
	if (strcmp(name, commands[i].name) == 0)
	    return &commands[i];
    }
    return NULL;
}

static void
usage(const char *argv0)
{
    const char *bin = gnu_basename(argv0);
    fprintf(stderr,
	    "Usage: %s [OPTIONS] COMMAND [...]\n"
	    "\n"
	    "Extract and inspect UKI sections from ukiboot partitions.\n"
	    "\n"
	    "Input modes:\n"
	    "  --image PATH       GPT disk image\n"
	    "  --bootctl PATH     Path to ukibootctl partition/file\n"
	    "  --boot-a PATH      Path to ukiboot_a partition/file\n"
	    "  --boot-b PATH      Path to ukiboot_b partition/file\n"
	    "\n"
	    "Options:\n"
	    "  --output-dir DIR   Output directory (default: current dir)\n"
	    "  --slot a|b         Force slot selection\n"
	    "  --cert PATH        Trusted certificate (PEM) for verification\n"
	    "  --verify           Verify signature before extracting\n"
	    "  --verbose          Verbose output\n"
	    "  --help             Show this help\n"
	    "  --version          Show version\n"
	    "\n"
	    "Commands:\n"
	    "  extract [SECTION...] Extract sections (all if none specified)\n"
	    "  list                 List PE sections with offsets and sizes\n"
	    "  status               Print active slot, kernel version, cmdline,\n"
	    "                       and kernel/initrd offset+size\n"
	    "  verify               Verify UKI Authenticode signature\n",
	    bin);
}

int
main(int argc, char **argv)
{
    const char *bin = argv[0];
    const struct option longopts[] = {
	{.name = "help", .has_arg = no_argument, .val = OPT_HELP},
	{.name = "version", .has_arg = no_argument, .val = OPT_VERSION},
	{.name = "verbose", .has_arg = no_argument, .val = OPT_VERBOSE},
	{.name = "image", .has_arg = required_argument, .val = OPT_IMAGE},
	{.name = "bootctl", .has_arg = required_argument, .val = OPT_BOOTCTL},
	{.name = "boot-a", .has_arg = required_argument, .val = OPT_BOOT_A},
	{.name = "boot-b", .has_arg = required_argument, .val = OPT_BOOT_B},
	{.name = "slot", .has_arg = required_argument, .val = OPT_SLOT},
	{.name = "output-dir", .has_arg = required_argument, .val = OPT_OUTPUT_DIR},
	{.name = "cert", .has_arg = required_argument, .val = OPT_CERT},
	{.name = "verify", .has_arg = no_argument, .val = OPT_VERIFY},
	{},
    };
    int opt;

    while ((opt = getopt_long(argc, argv, "", longopts, NULL)) != -1) {
	switch (opt) {
	case OPT_HELP:
	    usage(bin);
	    return 0;
	case OPT_VERSION:
	    printf("ukibootimg %s\n", PACKAGE_VERSION);
	    return 0;
	case OPT_VERBOSE:
	    opt_verbose = true;
	    break;
	case OPT_IMAGE:
	    opt_image = optarg;
	    break;
	case OPT_BOOTCTL:
	    opt_bootctl_path = optarg;
	    break;
	case OPT_BOOT_A:
	    opt_boot_a = optarg;
	    break;
	case OPT_BOOT_B:
	    opt_boot_b = optarg;
	    break;
	case OPT_SLOT:
	    if (strcmp(optarg, "a") == 0 || strcmp(optarg, "0") == 0)
		opt_slot = 0;
	    else if (strcmp(optarg, "b") == 0 || strcmp(optarg, "1") == 0)
		opt_slot = 1;
	    else
		fatal("Invalid slot '%s', expected 'a' or 'b'\n", optarg);
	    break;
	case OPT_OUTPUT_DIR:
	    opt_output_dir = optarg;
	    break;
	case OPT_CERT:
	    opt_cert = optarg;
	    break;
	case OPT_VERIFY:
	    opt_verify = true;
	    break;
	default:
	    usage(bin);
	    return 1;
	}
    }

    argv += optind;
    argc -= optind;

    if (argc < 1) {
	error("No command specified\n");
	usage(bin);
	return 1;
    }

    const char *command = argv[0];
    argv++;
    argc--;

    CmdHandler *handler = find_command_handler(command);
    if (!handler) {
	error("Unknown command '%s'\n", command);
	usage(bin);
	return 1;
    }

    if (opt_image && (opt_bootctl_path || opt_boot_a || opt_boot_b))
	fatal("Cannot combine --image with --bootctl/--boot-a/--boot-b\n");

    if (!opt_image && !opt_boot_a && !opt_boot_b)
	fatal("Must specify --image or at least one of --boot-a/--boot-b\n");

    autoclose int image_fd = -1;
    autoclose int bootctl_fd = -1;
    autoclose int uki_fd = -1;
    GptPartitionTable gpt = {0};
    uint32_t active_slot;

    if (opt_image) {
	image_fd = open(opt_image, O_RDONLY | O_CLOEXEC);
	if (image_fd < 0)
	    fatal("Failed to open %s: %s\n", opt_image, strerror(errno));

	if (!gpt_parse(image_fd, &gpt))
	    fatal("Failed to parse GPT from %s\n", opt_image);
    }

    if (opt_slot >= 0) {
	active_slot = opt_slot;
    } else {
	int bctl_fd;
	off_t bctl_offset = 0;

	if (opt_image) {
	    GptPartition *p = gpt_find_partition(&gpt, "ukibootctl");
	    if (!p)
		fatal("No ukibootctl partition found in image\n");
	    bctl_fd = image_fd;
	    bctl_offset = p->offset;
	} else if (opt_bootctl_path) {
	    bootctl_fd = open(opt_bootctl_path, O_RDONLY | O_CLOEXEC);
	    if (bootctl_fd < 0)
		fatal("Failed to open %s: %s\n", opt_bootctl_path, strerror(errno));
	    bctl_fd = bootctl_fd;
	} else {
	    verbose("No bootctl source, defaulting to slot a\n");
	    active_slot = 0;
	    goto have_slot;
	}

	if (!resolve_active_slot(bctl_fd, bctl_offset, &active_slot))
	    fatal("Failed to determine active slot\n");
    }
have_slot:

    verbose("Active slot: %c\n", active_slot == 0 ? 'a' : 'b');

    const char *slot_names[] = {"ukiboot_a", "ukiboot_b"};
    int u_fd;
    off_t u_offset = 0;
    const char *uki_path = NULL;

    if (opt_image) {
	GptPartition *p = gpt_find_partition(&gpt, slot_names[active_slot]);
	if (!p)
	    fatal("No %s partition found in image\n", slot_names[active_slot]);
	u_fd = image_fd;
	u_offset = p->offset;
	uki_path = opt_image;
    } else {
	const char *path = (active_slot == 0) ? opt_boot_a : opt_boot_b;
	if (!path)
	    fatal("No --%s specified for active slot\n", active_slot == 0 ? "boot-a" : "boot-b");
	uki_fd = open(path, O_RDONLY | O_CLOEXEC);
	if (uki_fd < 0)
	    fatal("Failed to open %s: %s\n", path, strerror(errno));
	u_fd = uki_fd;
	uki_path = path;
    }

    autofree PeSectionInfo *sections = NULL;
    uint32_t n_sections;

    if (!pe_parse_sections(u_fd, u_offset, &sections, &n_sections))
	fatal("Failed to parse PE sections\n");

    int ret = handler->fn(argc, argv, u_fd, u_offset, active_slot, uki_path, sections, n_sections);

    gpt_partition_table_free(&gpt);
    return ret;
}
