/* $NetBSD: gptsubr.c,v 1.2 2025/02/24 15:42:05 martin Exp $ */

/*
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS
 * OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
 * SUCH DAMAGE.
 */

#include <sys/cdefs.h>
#ifndef lint
__RCSID("$NetBSD: gptsubr.c,v 1.2 2025/02/24 15:42:05 martin Exp $");
#endif /* not lint */

#include <sys/param.h>

#include <sys/disk.h>
#include <sys/disklabel.h>
#include <sys/disklabel_gpt.h>
#include <sys/dkio.h>
#include <sys/ioctl.h>
#include <sys/stat.h>
#include <sys/statvfs.h>

#include <err.h>
#include <fcntl.h>
#include <libgen.h>
#include <paths.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <uuid.h>

#include "defs.h"
#include "gpt.h"
#include "gptsubr.h"
#include "map.h"

#define DEBUG

#ifdef DEBUG
#define DPRINTF(fmt, ...)	printf(fmt, __VA_ARGS__)
#else
#define DPRINTF(fmt, ...)	do { } while (/*CONSTCOND*/0)
#endif

static const char *
map_type_name(int map_type)
{
	const char *xtbl[] = {
		"UNUSED",	// 0
		"MBR",		// 1
		"MBR_PART",	// 2
		"PRI_GPT_HDR",	// 3
		"SEC_GPT_HDR",	// 4
		"PRI_GPT_TBL",	// 5
		"SEC_GPT_TBL",	// 6
		"GPT_PART",	// 7
		"PMBR",		// 8
	};

	if (map_type >= (int)__arraycount(xtbl))
		return "UNKNOWN";

	return xtbl[map_type];
}

#if 0
struct map {
	off_t	map_start;
	off_t	map_size;
	struct map *map_next;
	struct map *map_prev;
	int	map_type;
#define	MAP_TYPE_UNUSED		0
#define	MAP_TYPE_MBR		1
#define	MAP_TYPE_MBR_PART	2
#define	MAP_TYPE_PRI_GPT_HDR	3
#define	MAP_TYPE_SEC_GPT_HDR	4
#define	MAP_TYPE_PRI_GPT_TBL	5
#define	MAP_TYPE_SEC_GPT_TBL	6
#define	MAP_TYPE_GPT_PART	7
#define	MAP_TYPE_PMBR		8
	unsigned int map_index;
	void 	*map_data;
	int	map_alloc;
};

struct gpt_ent {
	uint8_t		ent_type[16];	/* partition type GUID */
	uint8_t		ent_guid[16];	/* unique partition GUID */
	uint64_t	ent_lba_start;	/* start of partition */
	uint64_t	ent_lba_end;	/* end of partition */
	uint64_t	ent_attr;	/* partition attributes */
	uint16_t	ent_name[36];	/* partition name in UCS-2 */
};

struct gpt_hdr {
	int8_t		hdr_sig[8];	/* identifies GUID Partition Table */
	uint32_t	hdr_revision;	/* GPT specification revision */
	uint32_t	hdr_size;	/* size of GPT Header */
	uint32_t	hdr_crc_self;	/* CRC32 of GPT Header */
	uint32_t	hdr__rsvd0;	/* must be zero */
	uint64_t	hdr_lba_self;	/* LBA that contains this Header */
	uint64_t	hdr_lba_alt;	/* LBA of backup GPT Header */
	uint64_t	hdr_lba_start;	/* first LBA usable for partitions */
	uint64_t	hdr_lba_end;	/* last LBA usable for partitions */
	uint8_t		hdr_guid[16];	/* GUID to identify the disk */
	uint64_t	hdr_lba_table;	/* first LBA of GPE array */
	uint32_t	hdr_entries;	/* number of entries in GPE array */
	uint32_t	hdr_entsz;	/* size of each GPE */
	uint32_t	hdr_crc_table;	/* CRC32 of GPE array */
	/*
	 * The remainder of the block that contains the GPT Header
	 * is reserved by EFI for future GPT Header expansion, and
	 * must be zero.
	 */
};
#endif

struct map_widths {
	uint start;
	uint size;
	uint type_name;
};

static struct map_widths
get_map_widths(gpt_t gpt)
{
	struct map_widths w;
	off_t max_start = 0;
	off_t max_size = 0;
	map_t m;

	w.type_name = 0;
	for (m = map_first(gpt); m != NULL; m = m->map_next) {
		uint n;

		if (max_start < m->map_start)
			max_start = m->map_start;
		if (max_size < m->map_size)
			max_size = m->map_size;
		if (m->map_type == MAP_TYPE_GPT_PART) {
			struct gpt_ent *ent = m->map_data;
			gpt_uuid_t gpt_uuid;
			char ent_type[128];

			memcpy(&gpt_uuid, ent->ent_type, sizeof(gpt_uuid));
			n = (uint)gpt_uuid_snprintf(ent_type, sizeof(ent_type), "%s", gpt_uuid);
			if (w.type_name < n)
				w.type_name = n;
		}

	}
	w.start = (uint)snprintf(NULL, 0, "%" PRIx64, max_start);
	w.size = (uint)snprintf(NULL, 0, "%" PRIx64, max_size);

#define MIN_WIDTH_OF(s)	(sizeof(s) - 1)
	if (w.type_name < MIN_WIDTH_OF("TYPE_NAME"))
		w.type_name = MIN_WIDTH_OF("TYPE_NAME");
	if (w.start < MIN_WIDTH_OF("START"))
		w.start = MIN_WIDTH_OF("START");
	if (w.size < MIN_WIDTH_OF("SIZE"))
		w.size = MIN_WIDTH_OF("SIZE");
#undef MIN_WIDTH_OF

	return w;
}

static void
show_map_banner(struct map_widths w)
{

	printf("IDX %*s %*s   ATTR            PARTITION_GUID"
	    "                         TYPE_UUID       %*s "
	    "TYPE_NAME ENTRY_NAME  DESCRIPTION\n",
	    w.start, "START",
	    w.size, "SIZE",
	    w.type_name, "");
}

static void
show_map(map_t m, struct map_widths w)
{
	struct gpt_ent *ent;
	struct gpt_hdr *hdr;
	gpt_uuid_t gpt_uuid;
	uuid_t uuid;
	uint32_t status;
	char *type_uuid;
	char *part_guid;
	uint64_t ent_attr;
	uint8_t  ent_desc[128];
	char     ent_type[128];

	ent_desc[0] = '\0';
	ent_type[0] = '\0';
	ent_attr = 0;

	switch (m->map_type) {
	case MAP_TYPE_PRI_GPT_HDR:
	case MAP_TYPE_SEC_GPT_HDR:
		hdr = m->map_data;
		memcpy(&uuid, hdr->hdr_guid, sizeof(uuid));
		uuid_to_string(&uuid, &part_guid, &status);
		type_uuid = estrdup("");
		break;
		
	case MAP_TYPE_GPT_PART:
		ent = m->map_data;

		memcpy(&uuid, ent->ent_type, sizeof(uuid));
		uuid_to_string(&uuid, &type_uuid, &status);
			
		memcpy(&uuid, ent->ent_guid, sizeof(uuid));
		uuid_to_string(&uuid, &part_guid, &status);
		
		ent_attr = ent->ent_attr;

		memcpy(&gpt_uuid, ent->ent_type, sizeof(uuid));
		gpt_uuid_snprintf(ent_type, sizeof(ent_type), "%s", gpt_uuid);
		
		/*
		 * Use the gpt.c code here rather than our
		 * ucs2_to_utf8() as we are in their world.
		 */
		utf16_to_utf8(ent->ent_name, sizeof(ent->ent_name),
		    ent_desc, sizeof(ent_desc));

		break;
		
	case MAP_TYPE_MBR_PART:
		part_guid = estrdup("");
		type_uuid = estrdup("");
		break;

	case MAP_TYPE_MBR: {
		struct mbr *mbr = m->map_data;

		easprintf(&part_guid, "%02x%02x%02x%02x-0000-0000-0000-000000000000",
		    mbr->mbr_code[440],
		    mbr->mbr_code[441],
		    mbr->mbr_code[442],
		    mbr->mbr_code[443]);
		uuid_to_string(&(uuid_t)GPT_ENT_TYPE_MBR, &type_uuid, &status);
		break;
	}
	default:
		part_guid = estrdup("");
		type_uuid = estrdup("");
		break;
	}

	printf("%2u: %*" PRIx64 " %*" PRIx64 " %08" PRIx64
	    " %36s %36s %*s %-11s %s\n",
	    m->map_index,
	    w.start,
	    m->map_start,
	    w.size,
	    m->map_size,
	    ent_attr,
	    part_guid,
	    type_uuid,
	    w.type_name,
	    ent_type,
	    map_type_name(m->map_type),
	    ent_desc);
	    
	free(part_guid);
	free(type_uuid);
}

PUBLIC map_t
find_gpt_map(const char *dev, uint idx)
{
	gpt_t gpt;
	map_t m;

	gpt = gpt_open(dev, GPT_READONLY, 0, 0, 0, 0);

	if (gpt == NULL)
		err(EXIT_FAILURE, "gpt_open");

	if (map_find(gpt, MAP_TYPE_PRI_GPT_HDR) == NULL)
		printf("GPT not found, displaying data from MBR.\n");

	for (m = map_first(gpt); m != NULL; m = m->map_next) {
		if (m->map_index == idx)
			break;
	}
	gpt_close(gpt);
	return m;
}

PUBLIC char *
parent_of_fname(const char *fname)
{
	struct dkwedge_info dkinfo;
	struct statvfs vfsb;
	struct stat sb;
	const char *d, *b;
	char *p;
	size_t n;
	int fd, rv;

	rv = stat(fname, &sb);
	if (rv == -1)
		err(EXIT_FAILURE, "stat: %s", fname);

	if (S_ISBLK(sb.st_mode) || S_ISCHR(sb.st_mode))
		return estrdup(fname);

	rv = statvfs(fname, &vfsb);
	if (rv == -1)
		err(EXIT_FAILURE, "statvfs: %s", fname);

	b = basename(vfsb.f_mntfromname);
	d = dirname(vfsb.f_mntfromname);
	easprintf(&p, "%s/r%s", d, b);
	
	fd = open(p, O_RDONLY);
	if (fd == -1)
		err(EXIT_FAILURE, "open");
		
	rv = ioctl(fd, DIOCGWEDGEINFO, &dkinfo);
	close(fd);

	if (rv != -1) {
		free(p);
		return estrdup(dkinfo.dkw_parent);
	}

	warn("ioctl: DIOCGWEDGEINFO");
			
	/*
	 * Hum.  No wedges?  Assume we have the old disklabel
	 * "/dev/rwd0x" syntax.  Convert it to "/dev/rwd0d".
	 * XXX: this probably won't work.
	 */
	n = strlen(p);
	p[n - 1] = 'd';

	return p;
}

PUBLIC int
mbr_sig_write(const char *fname, uint32_t new_sig, bool force, int verbose)
{
	gpt_t gpt;
	map_t m;
	struct mbr *mbr;
	const char *dev;
	uint32_t old_sig;

	if (fname == NULL)
		errx(EXIT_FAILURE, "please specify a device");

	dev = parent_of_fname(fname);
	if (dev == NULL) {
		warnx("unable to find parent device of %s\n", fname);
		return -1;
	}

	gpt = gpt_open(dev, GPT_MODIFIED, verbose, 0, 0, 0);

	if (gpt == NULL)
		err(EXIT_FAILURE, "gpt_open");

	m = map_find(gpt, MAP_TYPE_MBR);
	if (m == NULL)
		printf("No MBR partition found!\n");
	else {
		mbr = m->map_data;

		memcpy(&old_sig, &mbr->mbr_code[440], 4);

		if (old_sig == 0)
			force = true;

		if (force) {
			memcpy(&mbr->mbr_code[440], &new_sig, 4);
			if (gpt_write(gpt, m) == -1)
				warn("gpt_write");
			else if (verbose)
				printf("sig: 0x%08x -> 0x%08x\n",
				    old_sig, new_sig);
		}
		else if (verbose)
			printf("sig: 0x%08x (unchanged)\n", old_sig);
	}

	gpt_close(gpt);
	return 0;
}

PUBLIC int
show_gpt(const char *fname, int verbose)
{
	gpt_t gpt;
	map_t m;
	struct map_widths w;
	const char *dev;

	dev = parent_of_fname(fname);
	if (dev == NULL)
		return -1;

	gpt = gpt_open(dev, GPT_READONLY, verbose, 0, 0, 0);

	if (gpt == NULL)
		err(EXIT_FAILURE, "gpt_open");

	if (map_find(gpt, MAP_TYPE_PRI_GPT_HDR) == NULL)
		warnx("GPT not found, displaying data from MBR.");

	w = get_map_widths(gpt);
	show_map_banner(w);

	for (m = map_first(gpt); m != NULL; m = m->map_next)
		show_map(m, w);

	gpt_close(gpt);
	return 0;
}

#if 0 /* UNUSED */

PUBLIC char *
wedge_of_fname(const char *fname)
{
	struct statvfs vfsbuf;
	int rv;

	rv = statvfs(fname, &vfsbuf);
	if (rv) {
		warn("statvfs: %s", fname);
		return NULL;
	}

	return estrdup(vfsbuf.f_mntfromname);
}

PUBLIC int
find_partition_idx(const char *fname, int verbose)
{
	struct dkwedge_info dkinfo;
	struct stat sb;
	struct statvfs vfsbuf;
	gpt_t gpt;
	map_t m;
	off_t offset;
	size_t size;
	const char *parent;
	char *b, *d;
	int rv;

	/* the following are for gpt_open() */
	off_t mediasz = 0;
	u_int secsz = 0;
	time_t timestamp = 0;

	rv = stat(fname, &sb);
	if (rv == -1)
		err(EXIT_FAILURE, "stat: %s", fname);

	if (S_ISBLK(sb.st_mode) || S_ISCHR(sb.st_mode)) {
		parent = fname;
		offset = 0;
		size = 0;
		goto doit;
	}

	rv = statvfs(fname, &vfsbuf);
	if (rv == -1)
		err(EXIT_FAILURE, "statvfs: %s", fname);

	b = basename(vfsbuf.f_mntfromname);
	d = dirname(vfsbuf.f_mntfromname);

	{
		char *p;

		easprintf(&p, "%s/r%s", d, b);

		int fd = open(p, O_RDONLY);
		if (fd == -1)
			err(EXIT_FAILURE, "open");
		
		rv = ioctl(fd, DIOCGWEDGEINFO, &dkinfo);
		if (rv != -1) {
			parent = dkinfo.dkw_parent;
			offset = dkinfo.dkw_offset;
			size   = dkinfo.dkw_size;
		}
		else {
			struct disklabel dl;
			
			warn("ioctl: DIOCGWEDGEINFO");
			
			rv = ioctl(fd, DIOCGDINFO, &dl);
			if (rv == -1)
				err(EXIT_FAILURE, "ioctl: DIOCGDINFO");
			
			size_t n = strlen(p);
			
			int pnum = p[n - 1] - 'a';
			p[n - 1] = 'd';
			
			printf("num_parts: %u\n", dl.d_npartitions);
			printf("partition %d\n", pnum);
			printf("  offset = %u (%#x)\n", dl.d_partitions[pnum].p_offset, dl.d_partitions[pnum].p_offset);
			printf("  size = %u (%#x)\n",   dl.d_partitions[pnum].p_size,   dl.d_partitions[pnum].p_size);
			
			parent = p;	// vfsbuf.f_mntfromname;
			offset = dl.d_partitions[pnum].p_offset;
			size   = dl.d_partitions[pnum].p_size;
		}
		
		close(fd);
		free(p);
	}

 doit:
	DPRINTF("parent: %s\n", parent);
	DPRINTF("offset: 0x%lx\n", offset);
	DPRINTF("size: 0x%lx\n", size);

	gpt = gpt_open(parent, GPT_READONLY,
	    verbose, mediasz, secsz, timestamp);

	if (gpt == NULL)
		err(EXIT_FAILURE, "gpt_open");

	if (map_find(gpt, MAP_TYPE_PRI_GPT_HDR) == NULL)
		printf("GPT not found, displaying data from MBR.\n");

	int index = -1;
	struct map_widths w;
	if (verbose) {
		w = get_map_widths(gpt);
		show_map_banner(w);
	}
	for (m = map_first(gpt); m != NULL; m = m->map_next) {
		if (verbose)
			show_map(m, w);

		if (m->map_type != MAP_TYPE_GPT_PART || m->map_index < 1) {
			continue;
		}

		if ((off_t)offset == m->map_start &&
		    (off_t)size   == m->map_size) {
			if (index != -1)
				printf("match: %u\n", index);
			index = (int)m->map_index;
		}
	}
	return index;
}

PUBLIC char *
find_partition_pathname(const char *fname)
{
	char *pname;	/* partition name */
	char *rname;	/* real name */
	struct statvfs vfsbuf;
	int i, rv;

	rname = realpath(fname, NULL);

	DPRINTF("fname: %s\n", fname);
	DPRINTF("rname: %s\n", rname);

	rv = statvfs(rname, &vfsbuf);
	if (rv) {
		warn("statvfs: %s", rname);
		free(rname);
		return NULL;
	}

	DPRINTF("mount: %s\n", vfsbuf.f_mntonname);

	i = 0;
	while (vfsbuf.f_mntonname[i] == rname[i] && vfsbuf.f_mntonname[i] != '\0')
		i++;

	if (vfsbuf.f_mntonname[i] != '\0')
		errx(EXIT_FAILURE, "mntonname mismatch: %s",
		    vfsbuf.f_mntonname + i);
	
	pname = estrdup(rname + i);
	free(rname);

	return pname;
}

#endif