/* mbr.cc -- Functions for loading, saving, and manipulating legacy MBR partition

   data. */

/* Initial coding by Rod Smith, January to February, 2009 */

/* This program is copyright (c) 2009-2013 by Roderick W. Smith. It is distributed

  under the terms of the GNU GPL version 2, as detailed in the COPYING file. */

#define __STDC_LIMIT_MACROS

#define __STDC_CONSTANT_MACROS

#include <stdio.h>

#include <stdlib.h>

#include <stdint.h>

#include <fcntl.h>

#include <string.h>

#include <time.h>

#include <sys/stat.h>

#include <errno.h>

#include <iostream>

#include "mbr.h"

using namespace std;

/****************************************

 *                                      *

 * MBRData class and related structures *

 *                                      *

 ****************************************/

MBRData::~MBRData(void) {

} // MBRData destructor

// Assignment operator -- copy entire set of MBR data.

MBRData & MBRData::operator=(const BasicMBRData & orig) {

   BasicMBRData::operator=(orig);

   return *this;

} // MBRData::operator=()

/*****************************************************

 *                                                   *

 * Functions to create, delete, or change partitions *

 *                                                   *

 *****************************************************/

// Create a protective MBR. Clears the boot loader area if clearBoot > 0.

void MBRData::MakeProtectiveMBR(int clearBoot) {

   EmptyMBR(clearBoot);

   // Initialize variables

   nulls = 0;

   MBRSignature = MBR_SIGNATURE;

   diskSignature = UINT32_C(0);

   partitions[0].SetStatus(0); // Flag the protective part. as unbootable

   partitions[0].SetType(UINT8_C(0xEE));

   if (diskSize < UINT32_MAX) { // If the disk is under 2TiB

      partitions[0].SetLocation(UINT32_C(1), (uint32_t) diskSize - UINT32_C(1));

   } else { // disk is too big to represent, so fake it...

      partitions[0].SetLocation(UINT32_C(1), UINT32_MAX);

   } // if/else

   partitions[0].SetInclusion(PRIMARY);

   state = gpt;

} // MBRData::MakeProtectiveMBR()

// Optimizes the size of the 0xEE (EFI GPT) partition

void MBRData::OptimizeEESize(void) {

   int i, typeFlag = 0;

   uint64_t after;

   for (i = 0; i < 4; i++) {

      // Check for non-empty and non-0xEE partitions

      if ((partitions[i].GetType() != 0xEE) && (partitions[i].GetType() != 0x00))

         typeFlag++;

      if (partitions[i].GetType() == 0xEE) {

         // Blank space before this partition; fill it....

         if (SectorUsedAs(partitions[i].GetStartLBA() - 1, 4) == NONE) {

            partitions[i].SetStartLBA(FindFirstInFree(partitions[i].GetStartLBA() - 1));

         } // if

         // Blank space after this partition; fill it....

         after = partitions[i].GetStartLBA() + partitions[i].GetLengthLBA();

         if (SectorUsedAs(after, 4) == NONE) {

            partitions[i].SetLengthLBA(FindLastInFree(after) - partitions[i].GetStartLBA() + 1);

         } // if free space after

         if (after > diskSize) {

            if (diskSize < UINT32_MAX) { // If the disk is under 2TiB

               partitions[i].SetLengthLBA((uint32_t) diskSize - partitions[i].GetStartLBA());

            } else { // disk is too big to represent, so fake it...

               partitions[i].SetLengthLBA(UINT32_MAX - partitions[i].GetStartLBA());

            } // if/else

         } // if protective partition is too big

         RecomputeCHS(i);

      } // if partition is 0xEE

   } // for partition loop

   if (typeFlag == 0) { // No non-hybrid partitions found

      MakeProtectiveMBR(); // ensure it's a fully compliant protective MBR.

   } // if

} // MBRData::OptimizeEESize()

// Delete a partition if one exists at the specified location.

// Returns 1 if a partition was deleted, 0 otherwise....

// Used to help keep GPT & hybrid MBR partitions in sync....

int MBRData::DeleteByLocation(uint64_t start64, uint64_t length64) {

   uint32_t start32, length32;

   int i, deleted = 0;

   if ((start64 < UINT32_MAX) && (length64 < UINT32_MAX)) {

      start32 = (uint32_t) start64;

      length32 = (uint32_t) length64;

      for (i = 0; i < MAX_MBR_PARTS; i++) {

         if ((partitions[i].GetType() != 0xEE) && (partitions[i].GetStartLBA() == start32)

             && (partitions[i].GetLengthLBA() == length32)) {

            DeletePartition(i);

         if (state == hybrid)

            OptimizeEESize();

         deleted = 1;

         } // if (match found)

      } // for i (partition scan)

   } // if (hybrid & GPT partition < 2TiB)

   return deleted;

} // MBRData::DeleteByLocation()

/******************************************************

 *                                                    *

 * Functions that extract data on specific partitions *

 *                                                    *

 ******************************************************/

// Return the MBR data as a GPT partition....

GPTPart MBRData::AsGPT(int i) {

   MBRPart* origPart;

   GPTPart newPart;

   uint8_t origType;

   uint64_t firstSector, lastSector;

   newPart.BlankPartition();

   origPart = GetPartition(i);

   if (origPart != NULL) {

      origType = origPart->GetType();

      // don't convert extended, hybrid protective, or null (non-existent)

      // partitions (Note similar protection is in GPTData::XFormPartitions(),

      // but I want it here too in case I call this function in another

      // context in the future....)

      if ((origType != 0x05) && (origType != 0x0f) && (origType != 0x85) &&

          (origType != 0x00) && (origType != 0xEE)) {

         firstSector = (uint64_t) origPart->GetStartLBA();

         newPart.SetFirstLBA(firstSector);

         lastSector = (uint64_t) origPart->GetLastLBA();

         newPart.SetLastLBA(lastSector);

         newPart.SetType(((uint16_t) origType) * 0x0100);

         newPart.RandomizeUniqueGUID();

         newPart.SetAttributes(0);

         newPart.SetName(newPart.GetTypeName());

      } // if not extended, protective, or non-existent

   } // if (origPart != NULL)

   return newPart;

} // MBRData::AsGPT()