1/*
2 * Copyright 2020-2021 Haiku, Inc. All rights reserved.
3 * All rights reserved. Distributed under the terms of the MIT License.
4 *
5 * Authors:
6 * 		Alexander von Gluck IV <kallisti5@unixzen.com>
7 */
8
9
10#include <errno.h>
11#include <getopt.h>
12#include <fcntl.h>
13#include <unistd.h>
14#include <stdint.h>
15#include <stdio.h>
16#include <stdlib.h>
17#include <string.h>
18#include <sys/stat.h>
19
20#define EXIT_FAILURE 1
21
22
23//#define DEBUG_MBRTOOL
24#ifdef DEBUG_MBRTOOL
25#   define TRACE(x...) printf("mbrtool: " x)
26#else
27#   define TRACE(x...) ;
28#endif
29
30#define INFO(x...) printf("mbrtool: " x)
31
32
33// Disk sector size, 512 assumed!
34static const size_t kSectorSize = 512;
35
36static void
37print_usage(bool error)
38{
39	printf("\n");
40	printf("usage: mbrtool (options) <diskImage> <id> <type> <start> <len>\n");
41	printf("       <diskImage>            Disk image to operate on\n");
42	printf("       <id>                   Partition ID (0-3)\n");
43	printf("       <type>                 Partition type id (hex)\n");
44	printf("       <start>                Partition start offset (KiB)\n");
45	printf("       <len>                  Partition length (KiB)\n\n");
46	printf("  Options:\n");
47	printf("       -a, --active           Partition boot flag\n");
48	printf("\nWarning: This tool requires precision!\n");
49	printf("         Inputs are only lightly validated!\n\n");
50	exit(error ? EXIT_FAILURE : 0);
51}
52
53
54static void
55checkError(bool failed, const char *message)
56{
57	if (!failed)
58		return;
59
60	if (errno != 0)
61		INFO("Error: %s: %s\n", message, strerror(errno));
62	else
63		INFO("Error: %s\n", message);
64
65	exit(1);
66}
67
68
69static ssize_t
70mbrWipe(int handle)
71{
72	// Blow away MBR while avoiding bootloader
73	uint8_t emptyMBR[66] = {};
74	// End of MBR marker
75	emptyMBR[64] = 0x55;
76	emptyMBR[65] = 0xAA;
77	return pwrite(handle, emptyMBR, 66, 0x1BE);
78}
79
80
81static bool
82mbrValid(int handle)
83{
84	// TODO: this is a really basic check and ignores invalid
85	// partition table entries
86	uint8_t mbrBytes[66] = {};
87	ssize_t read = pread(handle, mbrBytes, 66, 0x1BE);
88	checkError(read < 0, "failed to read MBR for validation");
89	return (mbrBytes[64] == 0x55 && mbrBytes[65] == 0xAA);
90}
91
92
93static void
94createPartition(int handle, int index, bool active, uint8_t type,
95	uint64_t offset, uint64_t size)
96{
97	uint8_t bootable = active ? 0x80 : 0x0;
98	uint8_t partition[16] = {
99		bootable,		// bootable
100		0xff, 0xff, 0xff,	// CHS first block (default to LBA)
101		type,			// partition type
102		0xff, 0xff, 0xff,	// CHS last block (default to LBA)
103		0x00, 0x00, 0x00, 0x00,	// imageOffset in blocks (written below)
104		0x00, 0x00, 0x00, 0x00	// imageSize in blocks (written below)
105	};
106
107	// fill in LBA values
108	uint32_t partitionOffset = (uint32_t)(offset / kSectorSize);
109	((uint32_t *)partition)[2] = partitionOffset;
110	((uint32_t *)partition)[3] = (uint32_t)(size / kSectorSize);
111
112	TRACE("%s: #%d %c bytes: %u-%u, sectors: %u-%u\n", __func__, index,
113		active ? 'b' : '-', offset, offset + size, partitionOffset,
114		partitionOffset + uint32_t(size / kSectorSize));
115
116	ssize_t written = pwrite(handle, partition, 16, 512 - 2 - 16 * (4 - index));
117	checkError(written != 16, "failed to write partition entry");
118
119	if (active) {
120		// make it bootable
121		written = pwrite(handle, &partitionOffset, 4, offset + 512 - 2 - 4);
122		checkError(written != 4, "failed to make partition bootable");
123	}
124	return;
125}
126
127
128int
129main(int argc, char *argv[])
130{
131	const char *imageFile = NULL;
132
133	int partType = -1;
134	int partIndex = -1;
135	int64_t partStartOffset = -1;
136	int64_t partLength = -1;
137	bool partBootable = false;
138
139	while (1) {
140		int c;
141		static struct option long_options[] = {
142			{"active", no_argument, 0, 'a'},
143			{"help", no_argument, 0, 'h'},
144			{0, 0, 0, 0}
145		};
146
147		opterr = 1; /* don't print errors */
148		c = getopt_long(argc, argv, "+ha", long_options, NULL);
149
150		if (c == -1)
151			break;
152
153		switch (c) {
154			case 's':
155				print_usage(false);
156				break;
157			case 'a':
158				partBootable = true;
159				break;
160			default:
161				print_usage(true);
162		}
163	}
164
165	if ((argc - optind) != 5)
166		print_usage(true);
167
168	for (int index = optind; index < argc; index++) {
169		if (imageFile == NULL)
170			imageFile = argv[index];
171		else if (partIndex < 0)
172			partIndex = atoi(argv[index]);
173		else if (partType < 0)
174			partType = (int)strtol(argv[index], NULL, 0);
175		else if (partStartOffset < 0 ) {
176			partStartOffset = (int64_t)strtol(argv[index], NULL, 10);
177			partStartOffset *= 1024;
178		} else if (partLength < 0) {
179			partLength = (int64_t)strtol(argv[index], NULL, 10);
180			partLength *= 1024;
181		}
182	}
183
184	checkError(partIndex < 0 || partIndex > 3,
185		"invalid partition index, valid range is 0-3");
186	checkError(partType < 0,  "Incorrect Partition Type!");
187
188	checkError((partStartOffset + partLength) > 2089072000000,
189		"partitions beyond 2TiB are not accepted!");
190
191	int imageFileHandle = open(imageFile, O_RDWR);
192	checkError(imageFileHandle < 0, "failed to open disk image file");
193
194	struct stat stat;
195	int result = fstat(imageFileHandle, &stat);
196	checkError(result != 0, "failed to stat image file");
197	off_t imageSize = stat.st_size;
198
199	if (!mbrValid(imageFileHandle)) {
200		INFO("MBR of image is invalid, creating a fresh one.\n");
201		mbrWipe(imageFileHandle);
202	}
203
204	// Just a warning. This is technically valid since MBR partition
205	// definitions are purely within the first 512 bytes
206	if (partStartOffset + partLength > imageSize)
207		INFO("Warning: Partition extends beyond end of file!\n");
208
209	createPartition(imageFileHandle, partIndex, partBootable, partType,
210		partStartOffset, partLength);
211	return 0;
212}
213