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 21 using namespace std; 22 23 const 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 30 int 31 main(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