mop3/sdutil/sdutil.c

#include "limine-bios-hdd.h"
#include <cmdline_parser.h>
#include <devices.h>
#include <endianess.h>
#include <fs_types.h>
#include <in_input.h>
#include <libfat.h>
#include <malloc.h>
#include <mprintf.h>
#include <process.h>
#include <status.h>
#include <str_status.h>
#include <strconv.h>
#include <string.h>
#include <system.h>
#include <tcursor.h>
#include <tscreen.h>

static bool cmdline_part_dos = false;
static bool cmdline_list_part_dos = false;
static bool cmdline_format_fat32 = false;
static bool cmdline_format_fat16 = false;
static bool cmdline_install_limine_stage2 = false;
static bool cmdline_part_rescan = false;
static char cmdline_dev[CMDLINE_OPT_VALUE_MAX];

static struct cmdline_opt cmdline_opts[] = {
    CMDLINE_OPT("pd", "part-dos", CMDLINE_OPT_VALUE_BOOL, false, &cmdline_part_dos),
    CMDLINE_OPT("lpd", "list-part-dos", CMDLINE_OPT_VALUE_BOOL, false, &cmdline_list_part_dos),
    CMDLINE_OPT("ff32", "format-fat32", CMDLINE_OPT_VALUE_BOOL, false, &cmdline_format_fat32),
    CMDLINE_OPT("ff16", "format-fat16", CMDLINE_OPT_VALUE_BOOL, false, &cmdline_format_fat16),
    CMDLINE_OPT("ils2", "install-limine-stage2", CMDLINE_OPT_VALUE_BOOL, false,
                &cmdline_install_limine_stage2),
    CMDLINE_OPT("pr", "part-rescan", CMDLINE_OPT_VALUE_BOOL, false, &cmdline_part_rescan),
    CMDLINE_OPT("d", "device", CMDLINE_OPT_VALUE_STRING, true, (char*)cmdline_dev),
    CMDLINE_END(),
};

struct dos_pte {
  uint8_t drive_attrs;
  uint8_t chs_start_addr[3];
  uint8_t part_type;
  uint8_t chs_last_sect_addr[3];
  uint32_t start_lba;
  uint32_t sector_count;
} __attribute__((packed));

struct dos_mbr {
  uint8_t boot_code[440];
  uint8_t signature[4];
  uint8_t resv[2];
  struct dos_pte ptes[4];
  uint8_t valid_sign[2];
} __attribute__((packed));

static void lba_to_chs(uint32_t lba, uint8_t chs[3]) {
  uint32_t sectors_per_track = 63;
  uint32_t heads = 255;

  uint32_t cylinder = lba / (heads * sectors_per_track);
  uint32_t head = (lba / sectors_per_track) % heads;
  uint32_t sector = (lba % sectors_per_track) + 1;

  if (cylinder > 1023) {
    chs[0] = 254;
    chs[1] = 0xFF;
    chs[2] = 0xFF;
  } else {
    chs[0] = (uint8_t)head;
    chs[1] = (uint8_t)((sector & 0x3F) | ((cylinder >> 2) & 0xC0));
    chs[2] = (uint8_t)(cylinder & 0xFF);
  }
}

static void ask_and_setup_dos_partition(struct dos_pte* pte, int part_num) {
  char part_start_lba_str[12], part_sector_count_str[12];
  uint32_t part_start_lba, part_sector_count;
  char part_isboot_str[4];
  bool part_isboot;

  mprintf("\nConfigure parition %d:\n", part_num);

  in_stream_read_line("Start LBA: ", part_start_lba_str, sizeof(part_start_lba_str));
  in_stream_read_line("Sector count: ", part_sector_count_str, sizeof(part_sector_count_str));
  in_stream_read_line("Mark as bootable? [yes/no]: ", part_isboot_str, sizeof(part_isboot_str));

  part_start_lba = str_to_uint32(part_start_lba_str);
  part_sector_count = str_to_uint32(part_sector_count_str);
  part_isboot = (strcmp(part_isboot_str, "yes") == 0);

  if (part_sector_count == 0)
    return;

  pte->start_lba = part_start_lba;

  pte->sector_count = part_sector_count;

  lba_to_chs(part_start_lba, pte->chs_start_addr);

  lba_to_chs(part_start_lba + (part_sector_count > 1 ? part_sector_count : 1) - 1,
             pte->chs_last_sect_addr);

  pte->part_type = 0x0C;

  if (part_isboot)
    pte->drive_attrs |= (1 << 7);
}

static void start_part_dos(const char* dev_name) {
  struct dos_mbr mbr;
  memset(&mbr, 0, sizeof(mbr));
  char confirm_buf[4];
  memset(confirm_buf, 0, sizeof(confirm_buf));
  bool ok = false, write = false;

  ask_and_setup_dos_partition(&mbr.ptes[0], 0);
  ask_and_setup_dos_partition(&mbr.ptes[1], 1);
  ask_and_setup_dos_partition(&mbr.ptes[2], 2);
  ask_and_setup_dos_partition(&mbr.ptes[3], 3);

  struct dos_mbr* limine_mbr = (struct dos_mbr*)&binary_limine_hdd_bin_data[0];
  memcpy(mbr.boot_code, limine_mbr->boot_code, sizeof(mbr.boot_code));
  memcpy(mbr.signature, limine_mbr->signature, sizeof(mbr.signature));
  memcpy(mbr.valid_sign, limine_mbr->valid_sign, sizeof(mbr.valid_sign));

  /* clang-format off */
  mprintf ("\n\n-------------------------------------------------------\n");
  mprintf ("Current configuration (%s):\n", dev_name);
  mprintf ("Partition 0: start LBA=%u, sectors=%zu\n",mbr.ptes[0].start_lba, mbr.ptes[0].sector_count);
  mprintf ("Partition 1: start LBA=%u, sectors=%zu\n",mbr.ptes[1].start_lba, mbr.ptes[1].sector_count);
  mprintf ("Partition 2: start LBA=%u, sectors=%zu\n",mbr.ptes[2].start_lba, mbr.ptes[2].sector_count);
  mprintf ("Partition 3: start LBA=%u, sectors=%zu\n",mbr.ptes[3].start_lba, mbr.ptes[3].sector_count);
  /* clang-format on */

  in_stream_read_line("\nDoes this look plausible? [yes/no]: ", confirm_buf, sizeof(confirm_buf));

  if (strcmp(confirm_buf, "yes") == 0) {
    in_stream_read_line("\nWrite to device? There's no going back! [yes/no]: ", confirm_buf,
                        sizeof(confirm_buf));

    ok = true;

    if (strcmp(confirm_buf, "yes") == 0) {
      write = true;
    } else {
      mprintf("OK. Canceling\n");
    }
  } else {
    mprintf("OK. Canceling\n");
  }

  if (ok && write) {
    size_t sector = 0;
    size_t sector_count = 1;
    int r = device_do(dev_name, XDRV_WRITE, &sector, &sector_count, &mbr, NULL);
    mprintf("Finished writing. Result: %s (%d)\n", str_status[r < 0 ? -r : r], r);
  }
}

static void list_part_dos(const char* dev_name) {
  struct dos_mbr mbr;
  memset(&mbr, 0, sizeof(mbr));

  size_t sector = 0;
  size_t sector_count = 1;
  int r = device_do(dev_name, XDRV_READ, &sector, &sector_count, &mbr, NULL);
  mprintf("Finished reading. Result: %s (%d)\n", str_status[r < 0 ? -r : r], r);

  if (!(mbr.valid_sign[0] == 0x55 && mbr.valid_sign[1] == 0xAA)) {
    mprintf("ERROR invalid Master Boot Record!\n");
    return;
  }

  /* clang-format off */
  mprintf ("\n\n-------------------------------------------------------\n");
  mprintf ("Summary (%s):\n", dev_name);
  mprintf ("Partition 0: start LBA=%u, sectors=%zu\n",mbr.ptes[0].start_lba, mbr.ptes[0].sector_count);
  mprintf ("Partition 1: start LBA=%u, sectors=%zu\n",mbr.ptes[1].start_lba, mbr.ptes[1].sector_count);
  mprintf ("Partition 2: start LBA=%u, sectors=%zu\n",mbr.ptes[2].start_lba, mbr.ptes[2].sector_count);
  mprintf ("Partition 3: start LBA=%u, sectors=%zu\n",mbr.ptes[3].start_lba, mbr.ptes[3].sector_count);
  /* clang-format on */
}

static const char* spinner = "-\\|/";
static size_t spinner_state = 0;

static void format_update(size_t sector, size_t sector_count) {
  (void)sector;

  char spinner_char = spinner[(spinner_state++) % 4];
  mprintf(ANSIQ_CUR_SET_COL(0) ANSIQ_SCR_CLR2LEND "%c", spinner_char);
}

static void format_fat32(const char* dev_name) {
  mprintf("Formatting device %s!\n", dev_name);

  struct fatfs_ctx ctx;
  memset(&ctx, 0, sizeof(ctx));
  ctx.update_cb = &format_update;

  int r = fat_format_drive(&ctx, dev_name, FS_FAT32);

  mprintf("\nFormatting done: %s (%d)\n", str_status[r < 0 ? -r : r], r);
}

static void format_fat16(const char* dev_name) {
  mprintf("Formatting device %s!\n", dev_name);

  struct fatfs_ctx ctx;
  memset(&ctx, 0, sizeof(ctx));
  ctx.update_cb = &format_update;

  int r = fat_format_drive(&ctx, dev_name, FS_FAT16);

  mprintf("\nFormatting done: %s (%d)\n", str_status[r < 0 ? -r : r], r);
}

static void install_limine_stage2(const char* dev_name) {
  int r;
  struct dos_mbr mbr;
  memset(&mbr, 0, sizeof(mbr));

  size_t stage2_size = sizeof(binary_limine_hdd_bin_data) - 512;
  size_t stage2_sectors = (stage2_size + 511) / 512;
  size_t mbr_lba = 0, mbr_sector_count = 1;

  uint64_t stage2_lba = 1;
  uint8_t* stage2_data = (uint8_t*)((uintptr_t)&binary_limine_hdd_bin_data[0] + 512);

  size_t sector = 0;
  size_t sector_count = 1;
  r = device_do(dev_name, XDRV_READ, &sector, &sector_count, &mbr, NULL);

  if (r < 0) {
    mprintf("ERROR could not read MBR\n");
    return;
  }

  size_t stage2_end = stage2_lba + stage2_sectors;

  for (size_t i = 0; i < 4; i++) {
    if (mbr.ptes[i].sector_count == 0)
      continue;

    if (mbr.ptes[i].start_lba < stage2_end) {
      mprintf("ERROR partition overlaps with Limine stage 2 area!\n");
      return;
    }
  }

  uint64_t stage2_loc = stage2_lba * 512;
  memcpy((uint8_t*)&mbr + 0x1A4, &stage2_loc, sizeof(stage2_loc));
  r = device_do(dev_name, XDRV_WRITE, &mbr_lba, &mbr_sector_count, (void*)&mbr, NULL);

  if (r < 0) {
    mprintf("ERROR could not write patched MBR\n");
    return;
  }

  r = device_do(dev_name, XDRV_WRITE, &stage2_lba, &stage2_sectors, (void*)stage2_data, NULL);
  mprintf("Finished writing Limine Stage 2. Result: %s (%d)\n", str_status[r < 0 ? -r : r], r);
}

void app_main(void) {
  if (cmdline_parse(get_cmdline(), cmdline_opts) < 0) {
    mprintf("Failed to parse commandline arguments\n");
    return;
  }

  if (cmdline_part_dos) {
    start_part_dos(cmdline_dev);
  } else if (cmdline_list_part_dos) {
    list_part_dos(cmdline_dev);
  } else if (cmdline_format_fat32) {
    format_fat32(cmdline_dev);
  } else if (cmdline_format_fat16) {
    format_fat16(cmdline_dev);
  } else if (cmdline_install_limine_stage2) {
    install_limine_stage2(cmdline_dev);
  } else if (cmdline_part_rescan) {
    mprintf("WARNING Make sure all filesystems were unmounted from this device.\n");
    mprintf("Otherwise a rescan may result in broken filesystems and crashes!\n");

    device_do(cmdline_dev, XDRV_PARTITION_RESCAN, NULL, NULL, NULL, NULL);
  }
}