1/*
2 * Copyright 2012 Aleksas Pantechovskis, <alexp.frl@gmail.com>
3 * All rights reserved. Distributed under the terms of the MIT License.
4 */
5
6#include <fstream>
7#include <iostream>
8#include <sstream>
9#include <string.h>
10#include <string>
11
12#include <DiskDevice.h>
13#include <DiskDeviceRoster.h>
14#include <Path.h>
15#include <Volume.h>
16#include <VolumeRoster.h>
17
18#include "MBR.h"
19
20
21using namespace std;
22
23const char* kUsageMessage = \
24	"Usage: writembr [ device ] \n"
25	"#\tRewrites the MBR for the specified device.\n"
26	"#\tIf no device is specified, the boot device is used.\n"
27	"#\t--help shows this usage message\n";
28
29
30int
31main(int argc, char** argv)
32{
33	if ((argc == 2 && strcmp(argv[1], "--help") == 0) || argc > 2) {
34		cerr << kUsageMessage;
35		return B_ERROR;
36	}
37
38	BPath device;
39
40	if (argc == 2)
41		// user specified device for rewriting
42		device.SetTo(argv[1]);
43
44	else if (argc == 1) {
45		// no parameters specified, rewrite boot device
46		BVolumeRoster volumeRoster;
47		BVolume bootVolume;
48		if (volumeRoster.GetBootVolume(&bootVolume) != B_OK) {
49			cerr << "Can not find boot device" << endl;
50			return B_ERROR;
51		}
52
53		BDiskDeviceRoster roster;
54		BDiskDevice bootDevice;
55		if(roster.FindPartitionByVolume(bootVolume, &bootDevice, NULL) != B_OK) {
56			cerr << "Can not find boot device" << endl;
57			return B_ERROR;
58		}
59
60		bootDevice.GetPath(&device);
61	}
62
63
64	if (strcmp(device.Leaf(), "raw") != 0) {
65		cerr << device.Path() << " is not a raw device" << endl;
66		return B_ERROR;
67	}
68
69	fstream fs;
70	fs.open(device.Path(), fstream::in | fstream::out | fstream::binary);
71	if (!fs.is_open()) {
72		cerr << "Can't open " << device.Path() << endl;
73		return B_ERROR;
74	}
75
76	unsigned char MBR[kMBRSize];
77	fs.read((char*)MBR, kMBRSize);
78	if (fs.fail() || fs.gcount() < kMBRSize ) {
79		cerr << "Cannot read " << kMBRSize
80			<< " bytes from " << device.Path() << endl;
81		fs.close();
82		return B_ERROR;
83	}
84
85	// update only the code area and the MBR signature
86	memcpy(MBR, kMBR, 0x1be);
87	MBR[0x1FE] = kMBR[0x1FE];
88	MBR[0x1FF] = kMBR[0x1FF];
89
90	cerr << "About to overwrite the MBR boot code on " << device.Path()
91		<< "\nThis may disable any partition managers you have installed.\n"
92		<< "Are you sure you want to continue?\nyes/[no]: ";
93
94	string choice;
95	getline(cin, choice, '\n');
96	if (choice == "no" || choice == "" || choice != "yes") {
97		cerr << "MBR was NOT written" << endl;
98		fs.close();
99		return B_ERROR;
100	}
101
102	cerr << "Rewriting MBR for " << device.Path() << endl;
103
104	fs.seekg(0, ios::beg);
105	fs.write((char*)MBR, kMBRSize);
106	if (fs.fail()) {
107		cerr << "Cannot write " << kMBRSize
108			<< " bytes to " << device.Path() << endl;
109		fs.close();
110		return B_ERROR;
111	}
112
113	fs.close();
114
115	cerr << "MBR was written OK" << endl;
116	return B_OK;
117}
118
119