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 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