1/*
2 * Copyright 2007, Haiku, Inc. All Rights Reserved.
3 * Distributed under the terms of the MIT License.
4 *
5 * Authors:
6 *		Ingo Weinhold <bonefish@cs.tu-berlin.de>
7 */
8
9#include <ctype.h>
10#include <stdio.h>
11#include <stdlib.h>
12#include <string.h>
13
14#include <DiskDevice.h>
15#include <DiskDeviceRoster.h>
16#include <DiskDeviceVisitor.h>
17#include <DiskSystem.h>
18#include <PartitioningInfo.h>
19#include <Path.h>
20#include <String.h>
21
22#include <ObjectList.h>
23
24
25extern "C" const char* __progname;
26static const char* kProgramName = __progname;
27
28static const char* kUsage =
29"Usage: %s <options> <device>\n"
30"\n"
31"Options:\n"
32"  -h, --help   - print this help text\n"
33;
34
35
36static void
37print_usage(bool error)
38{
39	fprintf(error ? stderr : stdout, kUsage, kProgramName);
40}
41
42
43static void
44print_usage_and_exit(bool error)
45{
46	print_usage(error);
47	exit(error ? 1 : 0);
48}
49
50
51static void
52get_size_string(off_t _size, BString& string)
53{
54	const char* suffixes[] = {
55		"", "K", "M", "G", "T", NULL
56	};
57
58	double size = _size;
59	int index = 0;
60	while (size > 1024 && suffixes[index + 1]) {
61		size /= 1024;
62		index++;
63	}
64
65	char buffer[128];
66	snprintf(buffer, sizeof(buffer), "%.2f%s", size, suffixes[index]);
67
68	string = buffer;
69}
70
71
72// PrintLongVisitor
73class PrintLongVisitor : public BDiskDeviceVisitor {
74public:
75	virtual bool Visit(BDiskDevice *device)
76	{
77		BPath path;
78		status_t error = device->GetPath(&path);
79		const char *pathString = NULL;
80		if (error == B_OK)
81			pathString = path.Path();
82		else
83			pathString = strerror(error);
84		printf("device %" B_PRId32 ": \"%s\"\n", device->ID(), pathString);
85		printf("  has media:      %d\n", device->HasMedia());
86		printf("  removable:      %d\n", device->IsRemovableMedia());
87		printf("  read only:      %d\n", device->IsReadOnlyMedia());
88		printf("  write once:     %d\n", device->IsWriteOnceMedia());
89		printf("  ---\n");
90		Visit(device, 0);
91		return false;
92	}
93
94	virtual bool Visit(BPartition *partition, int32 level)
95	{
96		char prefix[128];
97		sprintf(prefix, "%*s", 2 * (int)level, "");
98		if (level > 0) {
99			BPath path;
100			status_t error = partition->GetPath(&path);
101			const char *pathString = NULL;
102			if (error == B_OK)
103				pathString = path.Path();
104			else
105				pathString = strerror(error);
106			printf("%spartition %" B_PRId32 ": \"%s\"\n", prefix, partition->ID(),
107				   pathString);
108		}
109		printf("%s  offset:         %" B_PRId64 "\n", prefix, partition->Offset());
110		printf("%s  size:           %" B_PRId64 "\n", prefix, partition->Size());
111		printf("%s  block size:     %" B_PRIu32 "\n", prefix, partition->BlockSize());
112		printf("%s  index:          %" B_PRId32 "\n", prefix, partition->Index());
113		printf("%s  status:         %" B_PRIu32 "\n", prefix, partition->Status());
114		printf("%s  file system:    %d\n", prefix,
115			   partition->ContainsFileSystem());
116		printf("%s  part. system:   %d\n", prefix,
117			   partition->ContainsPartitioningSystem());
118		printf("%s  device:         %d\n", prefix, partition->IsDevice());
119		printf("%s  read only:      %d\n", prefix, partition->IsReadOnly());
120		printf("%s  mounted:        %d\n", prefix, partition->IsMounted());
121		printf("%s  flags:          %" B_PRIx32 "\n", prefix, partition->Flags());
122		printf("%s  name:           \"%s\"\n", prefix, partition->Name());
123		printf("%s  content name:   \"%s\"\n", prefix,
124			partition->ContentName().String());
125		printf("%s  type:           \"%s\"\n", prefix, partition->Type());
126		printf("%s  content type:   \"%s\"\n", prefix,
127			partition->ContentType());
128		printf("%s  params:         \"%s\"\n", prefix, partition->Parameters());
129		printf("%s  content params: \"%s\"\n", prefix,
130			partition->ContentParameters());
131		// volume, icon,...
132		return false;
133	}
134};
135
136
137static void
138print_partition_table_header()
139{
140	printf("   Index     Start      Size                Content Type      "
141		"Content Name\n");
142	printf("--------------------------------------------------------------"
143		"------------\n");
144}
145
146
147static void
148print_partition(BPartition* partition, int level, int index)
149{
150	BString offset, size;
151	get_size_string(partition->Offset(), offset);
152	get_size_string(partition->Size(), size);
153
154	printf("%*s%02d%*s  %8s  %8s  %26.26s  %16.16s\n", 2 * level, "",
155		index,
156		2 * max_c(3 - level, 0), "",
157		offset.String(), size.String(),
158		(partition->ContentType() ? partition->ContentType() : "-"),
159		partition->ContentName().String());
160}
161
162
163// PrintShortVisitor
164class PrintShortVisitor : public BDiskDeviceVisitor {
165public:
166	virtual bool Visit(BDiskDevice* device)
167	{
168		fIndex = 0;
169
170		// get path
171		BPath path;
172		status_t error = device->GetPath(&path);
173		const char *pathString = NULL;
174		if (error == B_OK)
175			pathString = path.Path();
176		else
177			pathString = strerror(error);
178
179		// check media present; if so which read/write abilities
180		const char* media;
181		const char* readWrite;
182		if (device->HasMedia()) {
183			if (device->IsRemovableMedia()) {
184				media = "removable media";
185			} else {
186				media = "fixed media";
187			}
188
189			if (device->IsReadOnlyMedia()) {
190				if (device->IsWriteOnceMedia())
191					readWrite = ", write once";
192				else
193					readWrite = ", read-only";
194			} else
195				readWrite = ", read/write";
196		} else {
197			media = "no media";
198			readWrite = "";
199		}
200
201		printf("\ndevice %" B_PRId32 ": \"%s\": %s%s\n\n", device->ID(), pathString,
202			media, readWrite);
203		print_partition_table_header();
204
205		Visit(device, 0);
206		return false;
207	}
208
209	virtual bool Visit(BPartition* partition, int32 level)
210	{
211		print_partition(partition, level, fIndex++);
212		return false;
213	}
214
215private:
216	int	fIndex;
217};
218
219
220// FindPartitionByIndexVisitor
221class FindPartitionByIndexVisitor : public BDiskDeviceVisitor {
222public:
223	FindPartitionByIndexVisitor(int index)
224		: fIndex(index)
225	{
226	}
227
228	virtual bool Visit(BDiskDevice *device)
229	{
230		return Visit(device, 0);
231	}
232
233	virtual bool Visit(BPartition *partition, int32 level)
234	{
235		return (fIndex-- == 0);
236	}
237
238private:
239	int	fIndex;
240};
241
242
243// Partitioner
244class Partitioner {
245public:
246	Partitioner(BDiskDevice* device)
247		: fDevice(device),
248		  fPrepared(false)
249	{
250	}
251
252	void Run()
253	{
254		// prepare device modifications
255		if (fDevice->IsReadOnly()) {
256			printf("Device is read-only. Can't change anything.\n");
257		} else {
258			status_t error = fDevice->PrepareModifications();
259			fPrepared = (error == B_OK);
260			if (error != B_OK) {
261				printf("Error: Failed to prepare device for modifications: "
262					"%s\n", strerror(error));
263			}
264		}
265
266		// main input loop
267		while (true) {
268			BString line;
269			if (!_ReadLine("party> ", line))
270				return;
271
272			if (line == "") {
273				// do nothing
274			} else if (line == "h" || line == "help") {
275				_PrintHelp();
276			} else if (line == "i") {
277				_InitializePartition();
278			} else if (line == "l") {
279				_PrintPartitionsShort();
280			} else if (line == "ll") {
281				_PrintPartitionsLong();
282			} else if (line == "n") {
283				_NewPartition();
284			} else if (line == "q" || line == "quit") {
285				return;
286			} else if (line == "w") {
287				_WriteChanges();
288			} else {
289				printf("Invalid command \"%s\", type \"help\" for help\n",
290					line.String());
291			}
292		}
293	}
294
295private:
296	void _PrintHelp()
297	{
298		printf("Valid commands:\n"
299			"  h, help  - prints this help text\n"
300			"  i        - initialize a partition\n"
301			"  l        - lists the device's partitions recursively\n"
302			"  ll       - lists the device's partitions recursively, long "
303				"format\n"
304			"  l        - create a new child partition\n"
305			"  q, quit  - quits the program, changes are discarded\n"
306			"  w        - write changes to disk\n");
307	}
308
309	void _PrintPartitionsShort()
310	{
311		PrintShortVisitor visitor;
312		fDevice->VisitEachDescendant(&visitor);
313	}
314
315	void _PrintPartitionsLong()
316	{
317		PrintLongVisitor visitor;
318		fDevice->VisitEachDescendant(&visitor);
319	}
320
321	void _InitializePartition()
322	{
323		if (!fPrepared) {
324			if (fDevice->IsReadOnly())
325				printf("Device is read-only!\n");
326			else
327				printf("Sorry, not prepared for modifications!\n");
328			return;
329		}
330
331		// get the partition
332		int32 partitionIndex;
333		BPartition* partition = NULL;
334		if (!_SelectPartition("partition index [-1 to abort]: ", partition,
335				partitionIndex)) {
336			return;
337		}
338
339		printf("\nselected partition:\n\n");
340		print_partition_table_header();
341		print_partition(partition, 0, partitionIndex);
342
343		// get available disk systems
344		BObjectList<BDiskSystem> diskSystems(20, true);
345		BDiskDeviceRoster roster;
346		{
347			BDiskSystem diskSystem;
348			while (roster.GetNextDiskSystem(&diskSystem) == B_OK) {
349				if (partition->CanInitialize(diskSystem.PrettyName()))
350					diskSystems.AddItem(new BDiskSystem(diskSystem));
351			}
352		}
353
354		if (diskSystems.IsEmpty()) {
355			printf("There are no disk systems, that can initialize this "
356				"partition.\n");
357			return;
358		}
359
360		// print the available disk systems
361		printf("\ndisk systems that can initialize the selected partition:\n");
362		for (int32 i = 0; BDiskSystem* diskSystem = diskSystems.ItemAt(i); i++)
363			printf("%2" B_PRId32 "  %s\n", i, diskSystem->PrettyName());
364
365		printf("\n");
366
367		// get the disk system
368		int64 diskSystemIndex;
369		if (!_ReadNumber("disk system index [-1 to abort]: ", 0,
370				diskSystems.CountItems() - 1, -1, "invalid index",
371				diskSystemIndex)) {
372			return;
373		}
374		BDiskSystem* diskSystem = diskSystems.ItemAt(diskSystemIndex);
375
376		bool supportsName = diskSystem->SupportsContentName();
377		BString name;
378		BString parameters;
379		while (true) {
380			// let the user enter name and parameters
381			if ((supportsName && !_ReadLine("partition name: ", name))
382				|| !_ReadLine("partition parameters: ", parameters)) {
383				return;
384			}
385
386			// validate parameters
387			BString validatedName(name);
388			if (partition->ValidateInitialize(diskSystem->PrettyName(),
389					supportsName ? &validatedName : NULL, parameters.String())
390					!= B_OK) {
391				printf("Validation of the given values failed. Sorry, can't "
392					"continue.\n");
393				return;
394			}
395
396			// did the disk system change the name?
397			if (!supportsName || name == validatedName) {
398				printf("Everything looks dandy.\n");
399			} else {
400				printf("The disk system adjusted the file name to \"%s\".\n",
401					validatedName.String());
402				name = validatedName;
403			}
404
405			// let the user decide whether to continue, change parameters, or
406			// abort
407			bool changeParameters = false;
408			while (true) {
409				BString line;
410				_ReadLine("[c]ontinue, change [p]arameters, or [a]bort? ", line);
411				if (line == "a")
412					return;
413				if (line == "p") {
414					changeParameters = true;
415					break;
416				}
417				if (line == "c")
418					break;
419
420				printf("invalid input\n");
421			}
422
423			if (!changeParameters)
424				break;
425		}
426
427		// initialize
428		status_t error = partition->Initialize(diskSystem->PrettyName(),
429			supportsName ? name.String() : NULL, parameters.String());
430		if (error != B_OK)
431			printf("Initialization failed: %s\n", strerror(error));
432	}
433
434	void _NewPartition()
435	{
436		if (!fPrepared) {
437			if (fDevice->IsReadOnly())
438				printf("Device is read-only!\n");
439			else
440				printf("Sorry, not prepared for modifications!\n");
441			return;
442		}
443
444		// get the parent partition
445		BPartition* partition = NULL;
446		int32 partitionIndex;
447		if (!_SelectPartition("parent partition index [-1 to abort]: ",
448				partition, partitionIndex)) {
449			return;
450		}
451
452		printf("\nselected partition:\n\n");
453		print_partition_table_header();
454		print_partition(partition, 0, partitionIndex);
455
456		if (!partition->ContainsPartitioningSystem()) {
457			printf("The selected partition does not contain a partitioning "
458				"system.\n");
459			return;
460		}
461
462		// get supported types
463		BObjectList<BString> supportedTypes(20, true);
464		BString typeBuffer;
465		int32 cookie = 0;
466		while (partition->GetNextSupportedChildType(&cookie, &typeBuffer)
467				== B_OK) {
468			supportedTypes.AddItem(new BString(typeBuffer));
469		}
470
471		if (supportedTypes.IsEmpty()) {
472			printf("The partitioning system is not able to create any "
473				"child partition (no supported types).\n");
474			return;
475		}
476
477		// get partitioning info
478		BPartitioningInfo partitioningInfo;
479		status_t error = partition->GetPartitioningInfo(&partitioningInfo);
480		if (error != B_OK) {
481			printf("Failed to get partitioning info for partition: %s\n",
482				strerror(error));
483			return;
484		}
485
486		int32 spacesCount = partitioningInfo.CountPartitionableSpaces();
487		if (spacesCount == 0) {
488			printf("There's no space on the partition where a child partition "
489				"could be created\n");
490			return;
491		}
492
493		// let the user select the partition type, if there's more than one
494		int64 typeIndex = 0;
495		int32 supportedTypesCount = supportedTypes.CountItems();
496		if (supportedTypesCount > 1) {
497			// list them
498			printf("Possible partition types:\n");
499			for (int32 i = 0; i < supportedTypesCount; i++)
500				printf("%2" B_PRId32 "  %s\n", i, supportedTypes.ItemAt(i)->String());
501
502			if (!_ReadNumber("supported type index [-1 to abort]: ", 0,
503					supportedTypesCount - 1, -1, "invalid index", typeIndex)) {
504				return;
505			}
506		}
507
508		const char* type = supportedTypes.ItemAt(typeIndex)->String();
509
510		// list partitionable spaces
511		printf("Unused regions where the new partition could be created:\n");
512		for (int32 i = 0; i < spacesCount; i++) {
513			off_t _offset;
514			off_t _size;
515			partitioningInfo.GetPartitionableSpaceAt(i, &_offset, &_size);
516			BString offset, size;
517			get_size_string(_offset, offset);
518			get_size_string(_size, size);
519			printf("%2" B_PRId32 "  start: %8s,  size:  %8s\n", i, offset.String(),
520				size.String());
521		}
522
523		// let the user select the partitionable space, if there's more than one
524		int64 spaceIndex = 0;
525		if (spacesCount > 1) {
526			if (!_ReadNumber("unused region index [-1 to abort]: ", 0,
527					spacesCount - 1, -1, "invalid index", spaceIndex)) {
528				return;
529			}
530		}
531
532		off_t spaceOffset;
533		off_t spaceSize;
534		partitioningInfo.GetPartitionableSpaceAt(spaceIndex, &spaceOffset,
535			&spaceSize);
536
537		off_t start;
538		off_t size;
539		BString parameters;
540		while (true) {
541			// let the user enter start, size, and parameters
542
543			// start
544			while (true) {
545				BString spaceOffsetString;
546				get_size_string(spaceOffset, spaceOffsetString);
547				BString prompt("partition start [default: ");
548				prompt << spaceOffsetString << "]: ";
549				if (!_ReadSize(prompt.String(), spaceOffset, start))
550					return;
551
552				if (start >= spaceOffset && start <= spaceOffset + spaceSize)
553					break;
554
555				printf("invalid partition start\n");
556			}
557
558			// size
559			off_t maxSize = spaceOffset + spaceSize - start;
560			while (true) {
561				BString maxSizeString;
562				get_size_string(maxSize, maxSizeString);
563				BString prompt("partition size [default: ");
564				prompt << maxSizeString << "]: ";
565				if (!_ReadSize(prompt.String(), maxSize, size))
566					return;
567
568				if (size >= 0 && start + size <= spaceOffset + spaceSize)
569					break;
570
571				printf("invalid partition size\n");
572			}
573
574			// parameters
575			if (!_ReadLine("partition parameters: ", parameters))
576				return;
577
578			// validate parameters
579			off_t validatedStart = start;
580			off_t validatedSize = size;
581// TODO: Support the name parameter!
582			if (partition->ValidateCreateChild(&start, &size, type, NULL,
583					parameters.String()) != B_OK) {
584				printf("Validation of the given values failed. Sorry, can't "
585					"continue.\n");
586				return;
587			}
588
589			// did the disk system change offset or size?
590			if (validatedStart == start && validatedSize == size) {
591				printf("Everything looks dandy.\n");
592			} else {
593				BString startString, sizeString;
594				get_size_string(validatedStart, startString);
595				get_size_string(validatedSize, sizeString);
596				printf("The disk system adjusted the partition start and "
597					"size to %" B_PRIdOFF " (%s) and %" B_PRIdOFF " (%s).\n",
598					validatedStart, startString.String(), validatedSize,
599					sizeString.String());
600				start = validatedStart;
601				size = validatedSize;
602			}
603
604			// let the user decide whether to continue, change parameters, or
605			// abort
606			bool changeParameters = false;
607			while (true) {
608				BString line;
609				_ReadLine("[c]ontinue, change [p]arameters, or [a]bort? ", line);
610				if (line == "a")
611					return;
612				if (line == "p") {
613					changeParameters = true;
614					break;
615				}
616				if (line == "c")
617					break;
618
619				printf("invalid input\n");
620			}
621
622			if (!changeParameters)
623				break;
624		}
625
626		// create child
627		error = partition->CreateChild(start, size, type, NULL,
628			parameters.String());
629		if (error != B_OK)
630			printf("Creating the partition failed: %s\n", strerror(error));
631	}
632
633	void _WriteChanges()
634	{
635		if (!fPrepared || !fDevice->IsModified()) {
636			printf("No changes have been made!\n");
637			return;
638		}
639
640		printf("Writing changes to disk. This can take a while...\n");
641
642		// commit
643		status_t error = fDevice->CommitModifications();
644		if (error == B_OK) {
645			printf("All changes have been written successfully!\n");
646		} else {
647			printf("Failed to write all changes: %s\n", strerror(error));
648		}
649
650		// prepare again
651		error = fDevice->PrepareModifications();
652		fPrepared = (error == B_OK);
653		if (error != B_OK) {
654			printf("Error: Failed to prepare device for modifications: "
655				"%s\n", strerror(error));
656		}
657	}
658
659	bool _SelectPartition(const char* prompt, BPartition*& partition,
660		int32& _partitionIndex)
661	{
662		// if the disk device has no children, we select it without asking
663		if (fDevice->CountChildren() == 0) {
664			partition = fDevice;
665			_partitionIndex = 0;
666			return true;
667		}
668
669		// otherwise let the user select
670		_PrintPartitionsShort();
671
672		partition = NULL;
673		int64 partitionIndex;
674		while (true) {
675			if (!_ReadNumber(prompt, partitionIndex) || partitionIndex < 0)
676				return false;
677
678			FindPartitionByIndexVisitor visitor(partitionIndex);
679			partition = fDevice->VisitEachDescendant(&visitor);
680			if (partition) {
681				_partitionIndex = partitionIndex;
682				return true;
683			}
684
685			printf("invalid partition index\n");
686		}
687	}
688
689	bool _ReadLine(const char* prompt, BString& _line)
690	{
691		// prompt
692		printf(prompt);
693		fflush(stdout);
694
695		// read line
696		char line[256];
697		if (!fgets(line, sizeof(line), stdin))
698			return false;
699
700		// remove trailing '\n'
701		if (char* trailingNL = strchr(line, '\n'))
702			*trailingNL = '\0';
703
704		_line = line;
705		return true;
706	}
707
708	bool _ReadNumber(const char* prompt, int64& number)
709	{
710		while (true) {
711			BString line;
712			if (!_ReadLine(prompt, line))
713				return false;
714
715			char buffer[256];
716			if (sscanf(line.String(), "%" B_PRId64 "%s", &number, buffer) == 1)
717				return true;
718
719			printf("invalid input\n");
720		}
721	}
722
723	bool _ReadNumber(const char* prompt, int64 minNumber, int64 maxNumber,
724		int64 abortNumber, const char* invalidNumberMessage, int64& number)
725	{
726		while (true) {
727			BString line;
728			if (!_ReadLine(prompt, line))
729				return false;
730
731			char buffer[256];
732			if (sscanf(line.String(), "%" B_PRId64 "%s", &number, buffer) != 1) {
733				printf("invalid input\n");
734				continue;
735			}
736
737			if (number == abortNumber)
738				return false;
739
740			if (number >= minNumber && number <= maxNumber)
741				return true;
742
743			puts(invalidNumberMessage);
744		}
745	}
746
747	bool _ReadSize(const char* prompt, off_t defaultValue, off_t& size)
748	{
749		while (true) {
750			BString _line;
751			if (!_ReadLine(prompt, _line))
752				return false;
753			const char* line = _line.String();
754
755			// skip whitespace
756			while (isspace(*line))
757				line++;
758
759			if (*line == '\0') {
760				size = defaultValue;
761				return true;
762			}
763
764			// get the number
765			int32 endIndex = 0;
766			while (isdigit(line[endIndex]))
767				endIndex++;
768
769			if (endIndex == 0) {
770				printf("invalid input\n");
771				continue;
772			}
773
774			size = atoll(BString(line, endIndex).String());
775
776			// skip whitespace
777			line += endIndex;
778			while (isspace(*line))
779				line++;
780
781			// get the size modifier
782			if (*line != '\0') {
783				switch (*line) {
784					case 'K':
785						size *= 1024;
786						break;
787					case 'M':
788						size *= 1024 * 1024;
789						break;
790					case 'G':
791						size *= 1024LL * 1024 * 1024;
792						break;
793					case 'T':
794						size *= 1024LL * 1024 * 1024 * 1024;
795						break;
796					default:
797						line--;
798				}
799
800				line++;
801			}
802
803			if (*line == 'B')
804				line++;
805
806			// skip whitespace
807			while (isspace(*line))
808				line++;
809
810			if (*line == '\0')
811				return true;
812
813			printf("invalid input\n");
814		}
815	}
816
817private:
818	BDiskDevice*	fDevice;
819	bool			fPrepared;
820};
821
822
823int
824main(int argc, const char* const* argv)
825{
826	// parse arguments
827	int argi = 1;
828	for (; argi < argc; argi++) {
829		const char* arg = argv[argi];
830		if (arg[0] == '-') {
831			if (arg[1] == '-') {
832				// a double '-' option
833				if (strcmp(arg, "--help") == 0) {
834					print_usage_and_exit(false);
835				} else
836					print_usage_and_exit(true);
837			} else {
838				// single '-' options
839				for (int i = 1; arg[i] != '\0'; i++) {
840					switch (arg[i]) {
841						case 'h':
842							print_usage_and_exit(false);
843						default:
844							print_usage_and_exit(true);
845					}
846				}
847			}
848		} else
849			break;
850	}
851
852	// only the device path should remain
853	if (argi != argc - 1)
854		print_usage_and_exit(true);
855	const char* devicePath = argv[argi];
856
857	// get the disk device
858	BDiskDeviceRoster roster;
859	BDiskDevice device;
860	status_t error = roster.GetDeviceForPath(devicePath, &device);
861	if (error != B_OK) {
862		fprintf(stderr, "Error: Failed to get disk device for path \"%s\": "
863			"%s\n", devicePath, strerror(error));
864	}
865
866	Partitioner partitioner(&device);
867	partitioner.Run();
868
869	return 0;
870}
871