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
22 using namespace std;
23
24 const 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
31 int
main(int argc,char ** argv)32 main(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