/*	$OpenBSD: cmd.c,v 1.167 2023/03/04 21:22:51 krw Exp $	*/

/*
 * Copyright (c) 1997 Tobias Weingartner
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

#include <sys/types.h>
#include <sys/disklabel.h>

#include <err.h>
#include <signal.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <uuid.h>

#include "part.h"
#include "disk.h"
#include "misc.h"
#include "mbr.h"
#include "gpt.h"
#include "user.h"
#include "cmd.h"

int		 gedit(const int);
int		 edit(const int, struct mbr *);
int		 gsetpid(const int);
int		 setpid(const int, struct mbr *);
int		 parsepn(const char *);

int		 ask_num(const char *, int, int, int);
int		 ask_pid(const int);
const struct uuid *ask_uuid(const struct uuid *);

extern const unsigned char	manpage[];
extern const int		manpage_sz;

int
Xreinit(const char *args, struct mbr *mbr)
{
	int			dogpt;

	dogpt = 0;

	if (strcasecmp(args, "gpt") == 0)
		dogpt = 1;
	else if (strcasecmp(args, "mbr") == 0)
		dogpt = 0;
	else if (strlen(args) > 0) {
		printf("Unrecognized modifier '%s'\n", args);
		return CMD_CONT;
	}

	if (dogpt) {
		GPT_init(GHANDGP);
		GPT_print("s", TERSE);
	} else {
		MBR_init(mbr);
		MBR_print(mbr, "s");
	}

	printf("Use 'write' to update disk.\n");

	return CMD_DIRTY;
}

int
Xswap(const char *args, struct mbr *mbr)
{
	char			 lbuf[LINEBUFSZ];
	char			*from, *to;
	int			 pf, pt;
	struct prt		 pp;
	struct gpt_partition	 gg;

	if (strlcpy(lbuf, args, sizeof(lbuf)) >= sizeof(lbuf)) {
		printf("argument string too long\n");
		return CMD_CONT;
	}

	to = lbuf;
	from = strsep(&to, WHITESPACE);

	pt = parsepn(to);
	if (pt == -1)
		return CMD_CONT;

	pf = parsepn(from);
	if (pf == -1)
		return CMD_CONT;

	if (pt == pf) {
		printf("%d same partition as %d, doing nothing.\n", pt, pf);
		return CMD_CONT;
	}

	if (gh.gh_sig == GPTSIGNATURE) {
		gg = gp[pt];
		gp[pt] = gp[pf];
		gp[pf] = gg;
	} else {
		pp = mbr->mbr_prt[pt];
		mbr->mbr_prt[pt] = mbr->mbr_prt[pf];
		mbr->mbr_prt[pf] = pp;
	}

	return CMD_DIRTY;
}

int
gedit(const int pn)
{
	struct uuid		 oldtype;

	oldtype = gp[pn].gp_type;

	if (gsetpid(pn))
		return -1;

	if (uuid_is_nil(&gp[pn].gp_type, NULL)) {
		if (uuid_is_nil(&oldtype, NULL) == 0) {
			memset(&gp[pn], 0, sizeof(gp[pn]));
			printf("Partition %d is disabled.\n", pn);
		}
		return 0;
	}

	if (GPT_get_lba_start(pn) == -1 ||
	    GPT_get_lba_end(pn) == -1 ||
	    GPT_get_name(pn) == -1) {
		return -1;
	}

	return 0;
}

int
parsepn(const char *pnstr)
{
	const char		*errstr;
	int			 maxpn, pn;

	if (pnstr == NULL) {
		printf("no partition number\n");
		return -1;
	}

	if (gh.gh_sig == GPTSIGNATURE)
		maxpn = gh.gh_part_num - 1;
	else
		maxpn = NDOSPART - 1;

	pn = strtonum(pnstr, 0, maxpn, &errstr);
	if (errstr) {
		printf("partition number is %s: %s\n", errstr, pnstr);
		return -1;
	}

	return pn;
}

int
edit(const int pn, struct mbr *mbr)
{
	struct chs		 start, end;
	struct prt		*pp;
	uint64_t		 track;
	unsigned char		 oldid;

	pp = &mbr->mbr_prt[pn];
	oldid = pp->prt_id;

	if (setpid(pn, mbr))
		return -1;

	if (pp->prt_id == DOSPTYP_UNUSED) {
		if (oldid != DOSPTYP_UNUSED) {
			memset(pp, 0, sizeof(*pp));
			printf("Partition %d is disabled.\n", pn);
		}
		return 0;
	}

	if (ask_yn("Do you wish to edit in CHS mode?")) {
		PRT_lba_to_chs(pp, &start, &end);
		start.chs_cyl = ask_num("BIOS Starting cylinder", start.chs_cyl,
		    0, disk.dk_cylinders - 1);
		start.chs_head = ask_num("BIOS Starting head", start.chs_head,
		    0, disk.dk_heads - 1);
		start.chs_sect = ask_num("BIOS Starting sector", start.chs_sect,
		    1, disk.dk_sectors);

		end.chs_cyl = ask_num("BIOS Ending cylinder", end.chs_cyl,
		    start.chs_cyl, disk.dk_cylinders - 1);
		end.chs_head = ask_num("BIOS Ending head", end.chs_head,
		    (start.chs_cyl == end.chs_cyl) ? start.chs_head : 0,
		    disk.dk_heads - 1);
		end.chs_sect = ask_num("BIOS Ending sector", end.chs_sect,
		    (start.chs_cyl == end.chs_cyl && start.chs_head ==
		    end.chs_head) ? start.chs_sect : 1, disk.dk_sectors);

		/* The ATA/ATAPI spec says LBA = (C × HPC + H) × SPT + (S − 1) */
		track = start.chs_cyl * disk.dk_heads + start.chs_head;
		pp->prt_bs = track * disk.dk_sectors + (start.chs_sect - 1);
		track = end.chs_cyl * disk.dk_heads + end.chs_head;
		pp->prt_ns = track * disk.dk_sectors + (end.chs_sect - 1) -
		    pp->prt_bs + 1;
	} else {
		pp->prt_bs = getuint64("Partition offset", pp->prt_bs, 0,
		    disk.dk_size - 1);
		pp->prt_ns = getuint64("Partition size",   pp->prt_ns, 1,
		    disk.dk_size - pp->prt_bs);
	}

	return 0;
}

int
Xedit(const char *args, struct mbr *mbr)
{
	struct gpt_partition	 oldgg;
	struct prt		 oldprt;
	int			 pn;

	pn = parsepn(args);
	if (pn == -1)
		return CMD_CONT;

	if (gh.gh_sig == GPTSIGNATURE) {
		oldgg = gp[pn];
		if (gedit(pn))
			gp[pn] = oldgg;
		else if (memcmp(&gp[pn], &oldgg, sizeof(gp[pn])))
			return CMD_DIRTY;
	} else {
		oldprt = mbr->mbr_prt[pn];
		if (edit(pn, mbr))
			mbr->mbr_prt[pn] = oldprt;
		else if (memcmp(&mbr->mbr_prt[pn], &oldprt, sizeof(oldprt)))
			return CMD_DIRTY;
	}

	return CMD_CONT;
}

int
gsetpid(const int pn)
{
	uint32_t		 status;

	GPT_print_parthdr(TERSE);
	GPT_print_part(pn, "s", TERSE);

	if (PRT_protected_guid(&gp[pn].gp_type)) {
		printf("can't edit partition type %s\n",
		    PRT_uuid_to_sname(&gp[pn].gp_type));
		return -1;
	}

	gp[pn].gp_type = *ask_uuid(&gp[pn].gp_type);
	if (PRT_protected_guid(&gp[pn].gp_type)) {
		printf("can't change partition type to %s\n",
		    PRT_uuid_to_sname(&gp[pn].gp_type));
		return -1;
	}

	if (uuid_is_nil(&gp[pn].gp_guid, NULL)) {
		uuid_create(&gp[pn].gp_guid, &status);
		if (status != uuid_s_ok) {
			printf("could not create guid for partition\n");
			return -1;
		}
	}

	return 0;
}

int
setpid(const int pn, struct mbr *mbr)
{
	struct prt		*pp;

	pp = &mbr->mbr_prt[pn];

	PRT_print_parthdr();
	PRT_print_part(pn, pp, "s");

	pp->prt_id = ask_pid(pp->prt_id);

	return 0;
}

int
Xsetpid(const char *args, struct mbr *mbr)
{
	struct gpt_partition	oldgg;
	struct prt		oldprt;
	int			pn;

	pn = parsepn(args);
	if (pn == -1)
		return CMD_CONT;

	if (gh.gh_sig == GPTSIGNATURE) {
		oldgg = gp[pn];
		if (gsetpid(pn))
			gp[pn] = oldgg;
		else if (memcmp(&gp[pn], &oldgg, sizeof(gp[pn])))
			return CMD_DIRTY;
	} else {
		oldprt = mbr->mbr_prt[pn];
		if (setpid(pn, mbr))
			mbr->mbr_prt[pn] = oldprt;
		else if (memcmp(&mbr->mbr_prt[pn], &oldprt, sizeof(oldprt)))
			return CMD_DIRTY;
	}

	return CMD_CONT;
}

int
Xselect(const char *args, struct mbr *mbr)
{
	static uint64_t		lba_firstembr = 0;
	uint64_t		lba_self;
	int			pn;

	pn = parsepn(args);
	if (pn == -1)
		return CMD_CONT;

	lba_self = mbr->mbr_prt[pn].prt_bs;

	if ((mbr->mbr_prt[pn].prt_id != DOSPTYP_EXTEND) &&
	    (mbr->mbr_prt[pn].prt_id != DOSPTYP_EXTENDL)) {
		printf("Partition %d is not an extended partition.\n", pn);
		return CMD_CONT;
	}

	if (lba_firstembr == 0)
		lba_firstembr = lba_self;

	if (lba_self == 0) {
		printf("Loop to MBR (sector 0)! Not selected.\n");
		return CMD_CONT;
	} else {
		printf("Selected extended partition %d\n", pn);
		printf("New EMBR at offset %llu.\n", lba_self);
	}

	USER_edit(lba_self, lba_firstembr);

	return CMD_CONT;
}

int
Xprint(const char *args, struct mbr *mbr)
{
	if (gh.gh_sig == GPTSIGNATURE)
		GPT_print(args, VERBOSE);
	else if (MBR_valid_prt(mbr))
		MBR_print(mbr, args);
	else {
		DISK_printgeometry("s");
		printf("Offset: %d\tSignature: 0x%X.\tNo MBR or GPT.\n",
		    DOSBBSECTOR, (int)mbr->mbr_signature);
	}

	return CMD_CONT;
}

int
Xwrite(const char *args, struct mbr *mbr)
{
	unsigned int		i, n;

	for (i = 0, n = 0; i < nitems(mbr->mbr_prt); i++)
		if (mbr->mbr_prt[i].prt_id == DOSPTYP_OPENBSD)
			n++;
	if (n > 1) {
		warnx("MBR contains more than one OpenBSD partition!");
		if (ask_yn("Write MBR anyway?") == 0)
			return CMD_CONT;
	}

	if (gh.gh_sig == GPTSIGNATURE) {
		printf("Writing GPT.\n");
		if (GPT_write() == -1) {
			warnx("error writing GPT");
			return CMD_CONT;
		}
	} else {
		printf("Writing MBR at offset %llu.\n", mbr->mbr_lba_self);
		if (MBR_write(mbr) == -1) {
			warnx("error writing MBR");
			return CMD_CONT;
		}
		GPT_zap_headers();
	}

	return CMD_CLEAN;
}

int
Xquit(const char *args, struct mbr *mbr)
{
	return CMD_QUIT;
}

int
Xabort(const char *args, struct mbr *mbr)
{
	exit(0);
}

int
Xexit(const char *args, struct mbr *mbr)
{
	return CMD_EXIT;
}

int
Xhelp(const char *args, struct mbr *mbr)
{
	USER_help(mbr);

	return CMD_CONT;
}

int
Xupdate(const char *args, struct mbr *mbr)
{
	memcpy(mbr->mbr_code, default_dmbr.dmbr_boot, sizeof(mbr->mbr_code));
	mbr->mbr_signature = DOSMBR_SIGNATURE;
	printf("Machine code updated.\n");
	return CMD_DIRTY;
}

int
Xflag(const char *args, struct mbr *mbr)
{
	char			 lbuf[LINEBUFSZ];
	const char		*errstr;
	char			*part, *flag;
	long long		 val = -1;
	int			 i, pn;

	if (strlcpy(lbuf, args, sizeof(lbuf)) >= sizeof(lbuf)) {
		printf("argument string too long\n");
		return CMD_CONT;
	}

	flag = lbuf;
	part = strsep(&flag, WHITESPACE);

	pn = parsepn(part);
	if (pn == -1)
		return CMD_CONT;

	if (flag != NULL) {
		if (gh.gh_sig == GPTSIGNATURE)
			val = strtonum(flag, 0, INT64_MAX, &errstr);
		else
			val = strtonum(flag, 0, 0xff, &errstr);
		if (errstr) {
			printf("flag value is %s: %s.\n", errstr, flag);
			return CMD_CONT;
		}
		if (gh.gh_sig == GPTSIGNATURE)
			gp[pn].gp_attrs = val;
		else
			mbr->mbr_prt[pn].prt_flag = val;
		printf("Partition %d flag value set to 0x%llx.\n", pn, val);
	} else {
		if (gh.gh_sig == GPTSIGNATURE) {
			for (i = 0; i < gh.gh_part_num; i++) {
				if (i == pn)
					gp[i].gp_attrs = GPTPARTATTR_BOOTABLE;
				else
					gp[i].gp_attrs = 0;
			}
		} else {
			for (i = 0; i < nitems(mbr->mbr_prt); i++) {
				if (i == pn)
					mbr->mbr_prt[i].prt_flag = DOSACTIVE;
				else
					mbr->mbr_prt[i].prt_flag = 0x00;
			}
		}
		printf("Partition %d marked active.\n", pn);
	}

	return CMD_DIRTY;
}

int
Xmanual(const char *args, struct mbr *mbr)
{
	char			*pager = "/usr/bin/less";
	char			*p;
	FILE			*f;
	sig_t			 opipe;

	opipe = signal(SIGPIPE, SIG_IGN);
	if ((p = getenv("PAGER")) != NULL && (*p != '\0'))
		pager = p;
	if (asprintf(&p, "gunzip -qc|%s", pager) != -1) {
		f = popen(p, "w");
		if (f) {
			fwrite(manpage, manpage_sz, 1, f);
			pclose(f);
		}
		free(p);
	}

	signal(SIGPIPE, opipe);

	return CMD_CONT;
}

int
ask_num(const char *str, int dflt, int low, int high)
{
	char			 lbuf[LINEBUFSZ];
	const char		*errstr;
	int			 num;

	if (dflt < low)
		dflt = low;
	else if (dflt > high)
		dflt = high;

	do {
		printf("%s [%d - %d]: [%d] ", str, low, high, dflt);
		string_from_line(lbuf, sizeof(lbuf), TRIMMED);

		if (lbuf[0] == '\0') {
			num = dflt;
			errstr = NULL;
		} else {
			num = (int)strtonum(lbuf, low, high, &errstr);
			if (errstr)
				printf("%s is %s: %s.\n", str, errstr, lbuf);
		}
	} while (errstr);

	return num;
}

int
ask_pid(const int dflt)
{
	char			lbuf[LINEBUFSZ];
	int			num;

	for (;;) {
		printf("Partition id ('0' to disable) [01 - FF]: [%02X] ", dflt);
		printf("(? for help) ");
		string_from_line(lbuf, sizeof(lbuf), TRIMMED);

		if (strlen(lbuf) == 0)
			return dflt;
		if (strcmp(lbuf, "?") == 0) {
			PRT_print_mbrtypes();
			continue;
		}

		num = hex_octet(lbuf);
		if (num != -1)
			return num;

		printf("'%s' is not a valid partition id.\n", lbuf);
	}
}

const struct uuid *
ask_uuid(const struct uuid *olduuid)
{
	static struct uuid	 uuid;
	char			 lbuf[LINEBUFSZ];
	char			*dflt = NULL;
	uint32_t		 status;
	int			 num = 0;

	uuid = *olduuid;
	if (uuid_is_nil(&uuid, NULL) == 0) {
		num = PRT_uuid_to_type(&uuid);
		if (num == 0) {
			uuid_to_string(&uuid, &dflt, &status);
			if (status != uuid_s_ok) {
				printf("uuid_to_string() failed\n");
				goto done;
			}
		}
	}
	if (dflt == NULL) {
		if (asprintf(&dflt, "%X", num) == -1) {
			warn("asprintf()");
			goto done;
		}
	}

	for (;;) {
		printf("Partition id ('0' to disable) [01 - FF, <uuid>]: [%s] ",
		    dflt);
		printf("(? for help) ");
		string_from_line(lbuf, sizeof(lbuf), TRIMMED);

		if (strcmp(lbuf, "?") == 0) {
			PRT_print_gpttypes();
			continue;
		} else if (strlen(lbuf) == 0) {
			uuid = *olduuid;
			goto done;
		}

		uuid_from_string(lbuf, &uuid, &status);
		if (status == uuid_s_ok)
			goto done;

		num = hex_octet(lbuf);
		switch (num) {
		case -1:
			printf("'%s' is not a valid partition id\n", lbuf);
			break;
		case 0:
			uuid_create_nil(&uuid, NULL);
			goto done;
		default:
			uuid = *PRT_type_to_guid(num);
			if (uuid_is_nil(&uuid, NULL) == 0)
				goto done;
			printf("'%s' has no associated UUID\n", lbuf);
			break;
		}
	}

 done:
	free(dflt);
	return &uuid;
}