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