git.mutantstargoat.com Git - instimg/blob - src/main.c

   1 #include <stdio.h>
   2 #include <stdlib.h>
   3 #include <ctype.h>
   4 #include <malloc.h>
   5 #include <windows.h>
   6 #include <winioctl.h>
   7 #include "widgets.h"
   8 #include "rawdisk.h"
   9
  10 #define IMG_FNAME       "disk.img"
  11
  12 #define WIN_XPAD        10
  13 #define WIN_YPAD        15
  14
  15 static int instimg(const char *devpath);
  16 static void update_disks(void);
  17 static void onclick(struct wgt_widget *w);
  18 static void onmodify(struct wgt_widget *w);
  19 static void ck_usbonly_handler(struct wgt_widget *w);
  20
  21
  22 static struct wgt_window *win;
  23 static struct wgt_widget *lb_instto, *bn_inst, *bn_cancel, *cb_devs, *ck_usbonly, *bn_cancel_inst;
  24
  25 static struct rawdisk_device rawdev[64];
  26 static const char *items[64];
  27 static int num_rawdev;
  28
  29 static FILE *imgfile;
  30 static long imgfile_size;
  31 static const char *err_noimg_fmt =
  32         "Failed to open 256boss disk image: " IMG_FNAME "\n"
  33         "%s\n"
  34         "\n"
  35         "Please run this installer from the 256boss distribution directory";
  36 static const char *fail_err_str;
  37
  38 int WINAPI WinMain(HINSTANCE hinst, HINSTANCE hprevinst, char *cmdline, int showcmd)
  39 {
  40         int x, y, width, height, max_width;
  41         MSG msg;
  42         char *msgbuf, *errmsg;
  43
  44         if(!(imgfile = fopen(IMG_FNAME, "rb"))) {
  45                 errmsg = strerror(errno);
  46                 msgbuf = alloca(strlen(err_noimg_fmt) + strlen(errmsg) + 1);
  47                 sprintf(msgbuf, err_noimg_fmt, errmsg);
  48                 MessageBox(0, msgbuf, "Missing disk image!", MB_OK | MB_ICONSTOP | MB_APPLMODAL);
  49                 return 1;
  50         }
  51         fseek(imgfile, 0, SEEK_END);
  52         imgfile_size = ftell(imgfile);
  53
  54         if(!(win = wgt_window("256boss USB stick installer", 420, 170))) {
  55                 fclose(imgfile);
  56                 return 1;
  57         }
  58
  59         max_width = 0;
  60         y = WIN_YPAD;
  61         lb_instto = wgt_label(win, "Install to device:", WIN_XPAD, WIN_YPAD);
  62
  63         x = wgt_xpos_after(lb_instto, WGT_AUTO);
  64         cb_devs = wgt_combo(win, items, num_rawdev, 0, x, WIN_YPAD, 270, WGT_AUTO, onmodify);
  65         if((width = wgt_xpos_after(cb_devs, WIN_XPAD)) > max_width) max_width = width;
  66
  67         y = wgt_ypos_after(cb_devs, 16);
  68         ck_usbonly = wgt_checkbox(win, "only show USB devices", 1, x, y,
  69                         WGT_AUTO, WGT_AUTO, ck_usbonly_handler);
  70         if((width = wgt_xpos_after(ck_usbonly, WIN_XPAD)) > max_width) max_width = width;
  71
  72         y = wgt_ypos_after(ck_usbonly, 32);
  73         bn_inst = wgt_button(win, "Install", 70, y, 100, WGT_AUTO, onclick);
  74         bn_cancel = wgt_button(win, "Cancel", 230, y, 100, WGT_AUTO, onclick);
  75
  76         height = wgt_ypos_after(bn_inst, WIN_YPAD);
  77
  78         wgt_resize_window(win, max_width, height);
  79
  80         update_disks();
  81
  82         while(GetMessage(&msg, 0, 0, 0)) {
  83                 TranslateMessage(&msg);
  84                 DispatchMessage(&msg);
  85         }
  86
  87         fclose(imgfile);
  88         wgt_free_window(win);
  89         return 0;
  90 }
  91
  92 #define MIN_BS          (4 * 1024 * 1024)
  93
  94 static int instimg(const char *devpath)
  95 {
  96         MSG msg;
  97         HANDLE hdev;
  98         DWORD wrcount;
  99         DISK_GEOMETRY geom;
 100         int width, max_width, height, y, sz, blksz = 512, status = -1;
 101         char *blkbuf;
 102         struct wgt_window *win_inst;
 103         struct wgt_widget *w, *pbar, *ptext;
 104         struct wgt_rect *rect;
 105         long num_blk_done, total_blk;
 106         int progr;
 107
 108         if((hdev = CreateFile(devpath, GENERIC_WRITE, FILE_SHARE_READ | FILE_SHARE_WRITE,
 109                         0, OPEN_EXISTING, FILE_FLAG_WRITE_THROUGH, 0)) == INVALID_HANDLE_VALUE) {
 110                 DWORD err = GetLastError();
 111                 FormatMessage(FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM, 0, err, 0,
 112                                 &fail_err_str, 0, 0);
 113                 return -1;
 114         }
 115
 116         /*if(rawdisk_eject(hdev) == 0) {
 117                 rawdisk_load(hdev);
 118         }*/
 119
 120         if(DeviceIoControl(hdev, IOCTL_DISK_GET_DRIVE_GEOMETRY, 0, 0, &geom, sizeof geom, &wrcount, 0)) {
 121                 blksz = geom.BytesPerSector;
 122         }
 123
 124         while(blksz < MIN_BS) blksz <<= 1;
 125
 126         num_blk_done = 0;
 127         total_blk = (imgfile_size + blksz - 1) / blksz;
 128
 129         if(!(blkbuf = malloc(blksz))) {
 130                 CloseHandle(hdev);
 131                 fail_err_str = "failed to allocate block buffer";
 132                 return -1;
 133         }
 134
 135         max_width = 0;
 136
 137         win_inst = wgt_window("Installing ...", 320, 140);
 138         wgt_quit_on_close(win_inst, 0);
 139
 140         y = WIN_YPAD;
 141         sprintf(blkbuf, "Device sector size: %d bytes", (int)geom.BytesPerSector);
 142         w = wgt_label(win_inst, blkbuf, WIN_XPAD, y);
 143         if((width = wgt_xpos_after(w, WIN_XPAD)) > max_width) max_width = width;
 144
 145         y = wgt_ypos_after(w, 8);
 146         sprintf(blkbuf, "Copying in blocks of: %d bytes", blksz);
 147         w = wgt_label(win_inst, blkbuf, WIN_XPAD, y);
 148         if((width = wgt_xpos_after(w, WIN_XPAD)) > max_width) max_width = width;
 149
 150         y = wgt_ypos_after(w, 8);
 151         sprintf(blkbuf, "Image size: %ld bytes (%ld blocks)", imgfile_size, total_blk);
 152         w = wgt_label(win_inst, blkbuf, WIN_XPAD, y);
 153         if((width = wgt_xpos_after(w, WIN_XPAD)) > max_width) max_width = width;
 154
 155         y = wgt_ypos_after(w, 8);
 156         ptext = wgt_label(win_inst, "-/- (0%)", WIN_XPAD, y);
 157
 158         y = wgt_ypos_after(ptext, 8);
 159         pbar = wgt_progbar(win_inst, WIN_XPAD, y, width - WIN_XPAD * 2, WGT_AUTO, 0);
 160         if((width = wgt_xpos_after(pbar, WIN_XPAD)) > max_width) max_width = width;
 161
 162         y = wgt_ypos_after(pbar, 8);
 163         bn_cancel_inst = wgt_button(win_inst, "Cancel", 100, y, WGT_AUTO, WGT_AUTO, onclick);
 164         rect = wgt_widget_rect(bn_cancel_inst, 0);
 165         wgt_move_widget(bn_cancel_inst, WIN_XPAD + (width - rect->width) / 2, rect->y);
 166
 167         height = wgt_ypos_after(bn_cancel_inst, WIN_YPAD);
 168
 169         wgt_resize_window(win_inst, max_width, height);
 170
 171         /* go through any pending messages once first, to allow the window to get drawn correctly */
 172         while(PeekMessage(&msg, 0, 0, 0, PM_REMOVE)) {
 173                 if(msg.message == WM_CLOSE || msg.message == WM_DESTROY) {
 174                         goto end;
 175                 }
 176                 TranslateMessage(&msg);
 177                 DispatchMessage(&msg);
 178         }
 179
 180         rewind(imgfile);
 181         while((sz = fread(blkbuf, 1, blksz, imgfile)) > 0) {
 182                 if(sz < blksz) {
 183                         memset(blkbuf + sz, 0, blksz - sz);
 184                 }
 185                 SetFilePointer(hdev, num_blk_done * blksz, 0, FILE_BEGIN);
 186                 if(!WriteFile(hdev, blkbuf, blksz, &wrcount, 0)) {
 187                         DWORD err = GetLastError();
 188                         FormatMessage(FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM, 0, err, 0,
 189                                         &fail_err_str, 0, 0);
 190                         goto end;
 191                 }
 192
 193                 num_blk_done++;
 194                 progr = num_blk_done * 100 / total_blk;
 195                 sprintf(blkbuf, "%ld/%ld (%d%%)", num_blk_done, total_blk, progr);
 196                 wgt_set_text(ptext, blkbuf);
 197                 wgt_set_progress(pbar, progr);
 198
 199                 while(PeekMessage(&msg, 0, 0, 0, PM_REMOVE)) {
 200                         if(msg.message == WM_CLOSE || msg.message == WM_DESTROY) {
 201                                 goto end;
 202                         }
 203                         TranslateMessage(&msg);
 204                         DispatchMessage(&msg);
 205                 }
 206         }
 207
 208         rawdisk_refresh(hdev);
 209
 210         status = 0;
 211 end:
 212         wgt_free_window(win_inst);
 213
 214         free(blkbuf);
 215         CloseHandle(hdev);
 216         return status;
 217 }
 218
 219 static int isusbdev(struct rawdisk_device *dev)
 220 {
 221         char *ptr = dev->path;
 222         while(*ptr) {
 223                 int i;
 224                 char buf[3];
 225                 for(i=0; i<3; i++) {
 226                         buf[i] = tolower(ptr[i]);
 227                 }
 228                 if(memcmp(buf, "usb", 3) == 0) {
 229                         return 1;
 230                 }
 231                 ptr++;
 232         }
 233         return 0;
 234 }
 235
 236 static void update_disks(void)
 237 {
 238         int i;
 239
 240         for(i=0; i<num_rawdev; i++) {
 241                 free(rawdev[i].name);
 242                 free(rawdev[i].path);
 243         }
 244
 245         if((num_rawdev = rawdisk_detect(rawdev, sizeof rawdev / sizeof *rawdev)) == -1) {
 246                 MessageBox(0, "Failed to detect storage devices!", 0, MB_OK);
 247                 exit(1);
 248         }
 249
 250         if(wgt_checkbox_checked(ck_usbonly)) {
 251                 i = 0;
 252                 while(i < num_rawdev) {
 253                         if(!isusbdev(rawdev + i)) {
 254                                 free(rawdev[i].name);
 255                                 free(rawdev[i].path);
 256                                 rawdev[i] = rawdev[--num_rawdev];
 257                         } else {
 258                                 i++;
 259                         }
 260                 }
 261         }
 262
 263         for(i=0; i<num_rawdev; i++) {
 264                 items[i] = rawdev[i].name;
 265         }
 266
 267         wgt_combo_setitems(cb_devs, items, num_rawdev);
 268
 269         if(num_rawdev) {
 270                 wgt_enable_widget(cb_devs);
 271                 wgt_enable_widget(bn_inst);
 272         } else {
 273                 wgt_disable_widget(cb_devs);
 274                 wgt_disable_widget(bn_inst);
 275         }
 276 }
 277
 278 static const char *verify_fmt = \
 279         "You are about to install 256boss to: %s\n"
 280         "\n"
 281         "Device path:\n"
 282         "%s\n"
 283         "\n"
 284         "Any existing data in the selected device will be lost!\n"
 285         "Please verify you have selected the intended device.\n"
 286         "\n"
 287         "Are you certain you wish to proceed with the installation?\n";
 288
 289 static const char *success_str = \
 290         "Installation complete!\n"
 291         "\n"
 292         "If the 256boss partition doesn't show up immediately,\n"
 293         "please eject and re-attach the device\n";
 294
 295 static const char *fail_fmt = \
 296         "Installation failed!\n"
 297         "Error: %s\n";
 298
 299 static void onclick(struct wgt_widget *w)
 300 {
 301         int sel, len;
 302         unsigned int flags;
 303         char *msgbuf;
 304         static int inst_running;
 305
 306         if(w == bn_inst) {
 307                 if(inst_running) return;
 308
 309                 sel = wgt_combo_selected(cb_devs);
 310                 if(sel < 0 || sel >= num_rawdev) {
 311                         return;
 312                 }
 313
 314                 len = strlen(verify_fmt) + strlen(rawdev[sel].name) + strlen(rawdev[sel].path);
 315                 msgbuf = alloca(len + 1);
 316                 sprintf(msgbuf, verify_fmt, rawdev[sel].name, rawdev[sel].path);
 317
 318                 flags = MB_YESNO | MB_ICONWARNING | MB_APPLMODAL | MB_SETFOREGROUND;
 319                 if(MessageBox(0, msgbuf, "Proceed with installation?", flags) == IDYES) {
 320                         inst_running = 1;
 321                         if(instimg(rawdev[sel].path) == -1) {
 322                                 int prev_len = len;
 323                                 len = strlen(fail_fmt) + (fail_err_str ? strlen(fail_err_str) : strlen("unknown"));
 324                                 if(len > prev_len) {
 325                                         alloca(len - prev_len);
 326                                 }
 327                                 sprintf(msgbuf, fail_fmt, fail_err_str ? fail_err_str : "unknown");
 328                                 MessageBox(0, msgbuf, 0, MB_OK | MB_ICONSTOP | MB_APPLMODAL | MB_SETFOREGROUND);
 329                         } else {
 330                                 MessageBox(0, success_str, "Done", MB_OK | MB_ICONASTERISK | MB_APPLMODAL | MB_SETFOREGROUND);
 331                         }
 332                         inst_running = 0;
 333                 }
 334
 335         } else if(w == bn_cancel) {
 336                 PostQuitMessage(0);
 337
 338         } else if(w == bn_cancel_inst) {
 339                 struct wgt_window *win = wgt_widget_window(w);
 340                 PostMessage(wgt_window_handle(win), WM_CLOSE, 0, 0);
 341         }
 342 }
 343
 344 static void onmodify(struct wgt_widget *w)
 345 {
 346         int sel = wgt_combo_selected(w);
 347         if(sel >= 0) {
 348                 const char *selstr = wgt_get_combo_item(w, sel);
 349                 MessageBox(0, selstr, "selected item", MB_OK);
 350         }
 351 }
 352
 353 static void ck_usbonly_handler(struct wgt_widget *w)
 354 {
 355         update_disks();
 356 }
 357
 358 static int close_action(struct wgt_window *w)
 359 {
 360         if(w != win) return 0;
 361         return 1;
 362 }