genbootsector.c

来自「EFI BIOS是Intel提出的下一代的BIOS标准。这里上传的Edk源代码是」· C语言 代码 · 共 585 行 · 第 1/2 页

  CHAR        *DiskName,
  CHAR        *FileName,
  BOOL        WriteToDisk,
  PATCH_TYPE  PatchType,
  BOOL        ProcessMbr
  )
{
  BYTE    DiskPartition[0x200];
  BYTE    DiskPartitionBackup[0x200];
  HANDLE  DiskHandle;
  HANDLE  FileHandle;
  DWORD   BytesReturn;
  DWORD   DbrOffset;
  INT     DrvNumOffset;

  DiskHandle = CreateFile (
                 DiskName, 
                 GENERIC_READ | GENERIC_WRITE, 
                 FILE_SHARE_READ, 
                 NULL, 
                 OPEN_EXISTING, 
                 FILE_ATTRIBUTE_NORMAL, 
                 NULL
                 );
  if (DiskHandle == INVALID_HANDLE_VALUE) {
    return ErrorFileCreate;
  }

  FileHandle = CreateFile (
                 FileName,
                 GENERIC_READ | GENERIC_WRITE,
                 0,
                 NULL,
                 OPEN_ALWAYS,
                 FILE_ATTRIBUTE_NORMAL,
                 NULL
                 );
  if (FileHandle == INVALID_HANDLE_VALUE) {
    return ErrorFileCreate;
  }

  DbrOffset = 0;
  //
  // Skip potential MBR for Ide & USB disk
  //
  if ((PatchType == PatchTypeIde) || (PatchType == PatchTypeUsb)) {
    DbrOffset = GetBootSectorOffset (DiskHandle, WriteToDisk, PatchType);

    if (!ProcessMbr) {
      SetFilePointer (DiskHandle, DbrOffset * 0x200, NULL, FILE_BEGIN);
    }
    else if(DbrOffset == 0) {
      //
      // If user want to process Mbr, but no Mbr exists, simply return FALSE
      //
      return ErrorNoMbr;
    }
    else {
      SetFilePointer (DiskHandle, 0, NULL, FILE_BEGIN);
    }
  }

  //
  // [File Pointer is pointed to beginning of Mbr or Dbr]
  //
  if (WriteToDisk) {
    //
    // Write
    //
    if (!ReadFile (FileHandle, DiskPartition, 0x200, &BytesReturn, NULL)) {
      return ErrorFileReadWrite;
    }
    if (ProcessMbr) {
      //
      // Use original partition table
      //
      if (!ReadFile (DiskHandle, DiskPartitionBackup, 0x200, &BytesReturn, NULL)) {
        return ErrorFileReadWrite;
      }
      memcpy (DiskPartition + 0x1BE, DiskPartitionBackup + 0x1BE, 0x40);
      SetFilePointer (DiskHandle, 0, NULL, FILE_BEGIN);
    }

    if (!WriteFile (DiskHandle, DiskPartition, 0x200, &BytesReturn, NULL)) {
      return ErrorFileReadWrite;
    }

  } else {
    //
    // Read
    //
    if (!ReadFile (DiskHandle, DiskPartition, 0x200, &BytesReturn, NULL)) {
      return ErrorFileReadWrite;
    }

    if (PatchType == PatchTypeUsb) {
      // Manually set BS_DrvNum to 0x80 as window's format.exe has a bug which will clear this field discarding whether USB disk's MBR. 
      // offset of BS_DrvNum is 0x24 for FAT12/16
      //                        0x40 for FAT32
      //
      DrvNumOffset = GetDrvNumOffset (DiskPartition);
      if (DrvNumOffset == -1) {
        return ErrorFatType;
      }
      //
      // Some legacy BIOS require 0x80 discarding MBR.
      // Question left here: is it needed to check Mbr before set 0x80?
      //
      DiskPartition[DrvNumOffset] = ((DbrOffset > 0) ? 0x80 : 0);
  }


    if (PatchType == PatchTypeIde) {
      //
      // Patch LBAOffsetForBootSector
      //
      *(DWORD *)&DiskPartition [BOOT_SECTOR_LBA_OFFSET] = DbrOffset;
    }
    if (!WriteFile (FileHandle, DiskPartition, 0x200, &BytesReturn, NULL)) {
      return ErrorFileReadWrite;
    }
  }
  CloseHandle (FileHandle);
  CloseHandle (DiskHandle);
  return ErrorSuccess;
}

VOID
PrintUsage (
  CHAR* AppName
  )
{
  fprintf (
    stdout,
    "Usage: %s [OPTIONS]...\n"
    "Copy file content from/to bootsector.\n"
    "\n"
    "  -l        list disks\n"
    "  -if=FILE  specified an input, can be files or disks\n"
    "  -of=FILE  specified an output, can be files or disks\n"
    "  -mbr      process MBR also\n"
    "  -h        print this message\n"
    "\n"
    "FILE providing a volume plus a colon (X:), indicates a disk\n"
    "FILE providing other format, indicates a file\n",
    AppName
    );
}
 
INT
main (
  INT  argc,
  CHAR *argv[]
  )
{
  CHAR          *AppName;
  INT           Index;
  BOOL          ProcessMbr;
  CHAR          VolumeLetter;
  CHAR          *FilePath;
  BOOL          WriteToDisk;
  DRIVE_INFO    DriveInfo;
  PATCH_TYPE    PatchType;
  ERROR_STATUS  Status;

  CHAR        FloppyPathTemplate[] = "\\\\.\\%c:";
  CHAR        DiskPathTemplate[]   = "\\\\.\\PHYSICALDRIVE%u";
  CHAR        DiskPath[MAX_PATH];

  AppName = *argv;
  argv ++;
  argc --;
  
  ProcessMbr    = FALSE;
  WriteToDisk   = TRUE;
  FilePath      = NULL;
  VolumeLetter  = 0;

  //
  // Parse command line
  //
  for (Index = 0; Index < argc; Index ++) {
    if (_stricmp (argv[Index], "-l") == 0) {
      ListDrive ();
      return 0;
    }
    else if (_stricmp (argv[Index], "-mbr") == 0) {
      ProcessMbr = TRUE;
    }
    else if ((_strnicmp (argv[Index], "-if=", 4) == 0) ||
             (_strnicmp (argv[Index], "-of=", 4) == 0)
             ) {
      if (argv[Index][6] == '\0' && argv[Index][5] == ':' && IsLetter (argv[Index][4])) {
        VolumeLetter = argv[Index][4];
        if (_strnicmp (argv[Index], "-if=", 4) == 0) {
          WriteToDisk = FALSE;
        }
      }
      else {
        FilePath = &argv[Index][4];
      }
    }
    else {
      PrintUsage (AppName);
      return 1;
    }
  }

  //
  // Check parameter
  //
  if (VolumeLetter == 0) {
    fprintf (stderr, "ERROR: Volume isn't provided!\n");
    PrintUsage (AppName);
    return 1;
  }
  
  if (FilePath == NULL) {
    fprintf (stderr, "ERROR: File isn't pvovided!\n");
    PrintUsage (AppName);
    return 1;
  }
    
  PatchType = PatchTypeUnknown;

  if ((VolumeLetter == 'A') || (VolumeLetter == 'a') || 
      (VolumeLetter == 'B') || (VolumeLetter == 'b') 
      ) {
    //
    // Floppy
    //
    sprintf (DiskPath, FloppyPathTemplate, VolumeLetter);
    PatchType = PatchTypeFloppy;
  }
  else {
    //
    // Hard/USB disk
    //
    if (!GetDriveInfo (VolumeLetter, &DriveInfo)) {
      fprintf (stderr, "ERROR: GetDriveInfo - 0x%x\n", GetLastError ());
      return 1;
    }

    //
    // Shouldn't patch my own hard disk, but can read it.
    // very safe then:)
    //
    if (DriveInfo.DriveType->Type == DRIVE_FIXED && WriteToDisk) {
      fprintf (stderr, "ERROR: Write to local harddisk - permission denied!\n");
      return 1;
    }
    
    sprintf (DiskPath, DiskPathTemplate, DriveInfo.DiskNumber);
    if (DriveInfo.DriveType->Type == DRIVE_REMOVABLE) {
      PatchType = PatchTypeUsb;
    }
    else if (DriveInfo.DriveType->Type == DRIVE_FIXED) {
      PatchType = PatchTypeIde;
    }
  }

  if (PatchType == PatchTypeUnknown) {
    fprintf (stderr, "ERROR: PatchType unknown!\n");
    return 1;
  }

  //
  // Process DBR (Patch or Read)
  //
  Status = ProcessBootSector (DiskPath, FilePath, WriteToDisk, PatchType, ProcessMbr);
  if (Status == ErrorSuccess) {
    fprintf (
      stdout, 
      "%s %s successfully!\n", 
      WriteToDisk ? "Write" : "Read", 
      ProcessMbr ? "MBR" : "DBR"
      );
    return 0;
  }
  else {
    fprintf (
      stderr, 
      "ERROR: %s %s failed - %s (LastError: 0x%x)!\n", 
      WriteToDisk ? "Write" : "Read", 
      ProcessMbr ? "MBR" : "DBR", 
      ErrorStatusDesc[Status],
      GetLastError ()
      );
    return 1;
  }
}