src/bin/writembr/writembr.cpp

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