VtoyTool/vtoyexpand.c

   1 /******************************************************************************
   2  * vtoyexpand.c  ---- ventoy auto install script variable expansion
   3  *
   4  * Copyright (c) 2022, longpanda <admin@ventoy.net>
   5  *
   6  * This program is free software; you can redistribute it and/or
   7  * modify it under the terms of the GNU General Public License as
   8  * published by the Free Software Foundation; either version 3 of the
   9  * License, or (at your option) any later version.
  10  *
  11  * This program is distributed in the hope that it will be useful, but
  12  * WITHOUT ANY WARRANTY; without even the implied warranty of
  13  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  14  * General Public License for more details.
  15  *
  16  * You should have received a copy of the GNU General Public License
  17  * along with this program; if not, see <http://www.gnu.org/licenses/>.
  18  *
  19  */
  20
  21 #include <stdio.h>
  22 #include <stdlib.h>
  23 #include <string.h>
  24 #include <stdarg.h>
  25 #include <errno.h>
  26 #include <unistd.h>
  27 #include <fcntl.h>
  28 #include <sys/types.h>
  29 #include <sys/mman.h>
  30 #include <sys/ioctl.h>
  31 #include <sys/stat.h>
  32 #include <sys/types.h>
  33 #include <linux/fs.h>
  34 #include <dirent.h>
  35 #include "vtoytool.h"
  36
  37 #ifndef O_BINARY
  38 #define O_BINARY 0
  39 #endif
  40
  41 #define TMP_FILE    "/ventoy/tmp_var_expansion"
  42
  43 #define SIZE_1MB    (1024 * 1024)
  44 #define ulong unsigned long
  45 #define ulonglong  unsigned long long
  46
  47 typedef struct disk_info
  48 {
  49     char name[128];
  50     ulonglong size;
  51     int isUSB;
  52     int isSDX;
  53 }disk_info;
  54
  55 static disk_info *g_disk_list = NULL;
  56 static int g_disk_num = 0;
  57 static const char *g_vtoy_disk_name = NULL;
  58
  59 static void vlog(const char *fmt, ...)
  60 {
  61     int n = 0;
  62     va_list arg;
  63     FILE *fp = NULL;
  64     char log[1024];
  65
  66     fp = fopen("/ventoy/autoinstall.log", "a+");
  67     if (fp)
  68     {
  69         va_start(arg, fmt);
  70         n += vsnprintf(log, sizeof(log) - 1, fmt, arg);
  71         va_end(arg);
  72
  73         fwrite(log, 1, n, fp);
  74         fclose(fp);
  75     }
  76 }
  77
  78 static int copy_file(const char *file1, const char *file2)
  79 {
  80     int n;
  81     int size;
  82     int ret = 1;
  83     FILE *fp1 = NULL;
  84     FILE *fp2 = NULL;
  85     char *buf = NULL;
  86
  87     fp1 = fopen(file1, "rb");
  88     if (!fp1)
  89     {
  90         vlog("Failed to read file <%s>\n", file1);
  91         goto end;
  92     }
  93
  94     fp2 = fopen(file2, "wb+");
  95     if (!fp2)
  96     {
  97         vlog("Failed to create file <%s>\n", file2);
  98         goto end;
  99     }
 100
 101     fseek(fp1, 0, SEEK_END);
 102     size = (int)ftell(fp1);
 103     fseek(fp1, 0, SEEK_SET);
 104
 105     buf = malloc(size);
 106     if (!buf)
 107     {
 108         vlog("Failed to malloc buf\n");
 109         goto end;
 110     }
 111
 112     n = fread(buf, 1, size, fp1);
 113     if (n != size)
 114     {
 115         vlog("Failed to read <%s> %d %d\n", file1, n, size);
 116         goto end;
 117     }
 118
 119     n = fwrite(buf, 1, size, fp2);
 120     if (n != size)
 121     {
 122         vlog("Failed to write <%s> %d %d\n", file2, n, size);
 123         goto end;
 124     }
 125
 126     ret = 0;
 127
 128 end:
 129
 130     if (fp1)
 131         fclose(fp1);
 132     if (fp2)
 133         fclose(fp2);
 134     if (buf)
 135         free(buf);
 136
 137     return ret;
 138 }
 139
 140 static int vtoy_is_possible_blkdev(const char *name)
 141 {
 142     if (name[0] == '.')
 143     {
 144         return 0;
 145     }
 146
 147     /* /dev/ramX */
 148     if (name[0] == 'r' && name[1] == 'a' && name[2] == 'm')
 149     {
 150         return 0;
 151     }
 152
 153     /* /dev/loopX */
 154     if (name[0] == 'l' && name[1] == 'o' && name[2] == 'o' && name[3] == 'p')
 155     {
 156         return 0;
 157     }
 158
 159     /* /dev/dm-X */
 160     if (name[0] == 'd' && name[1] == 'm' && name[2] == '-' && IS_DIGIT(name[3]))
 161     {
 162         return 0;
 163     }
 164
 165     /* /dev/srX */
 166     if (name[0] == 's' && name[1] == 'r' && (name[2] >= '0' && name[2] <= '9'))
 167     {
 168         return 0;
 169     }
 170
 171     return 1;
 172 }
 173
 174 static ulonglong vtoy_get_disk_size_in_byte(const char *disk)
 175 {
 176     int fd;
 177     int rc;
 178     unsigned long long size = 0;
 179     char diskpath[256] = {0};
 180     char sizebuf[64] = {0};
 181
 182     // Try 1: get size from sysfs
 183     snprintf(diskpath, sizeof(diskpath) - 1, "/sys/block/%s/size", disk);
 184     if (access(diskpath, F_OK) >= 0)
 185     {
 186         vlog("get disk size from sysfs for %s\n", disk);
 187
 188         fd = open(diskpath, O_RDONLY | O_BINARY);
 189         if (fd >= 0)
 190         {
 191             read(fd, sizebuf, sizeof(sizebuf));
 192             size = strtoull(sizebuf, NULL, 10);
 193             close(fd);
 194             return (size * 512);
 195         }
 196     }
 197     else
 198     {
 199         vlog("%s not exist \n", diskpath);
 200     }
 201
 202     // Try 2: get size from ioctl
 203     snprintf(diskpath, sizeof(diskpath) - 1, "/dev/%s", disk);
 204     fd = open(diskpath, O_RDONLY);
 205     if (fd >= 0)
 206     {
 207         vlog("get disk size from ioctl for %s\n", disk);
 208         rc = ioctl(fd, BLKGETSIZE64, &size);
 209         if (rc == -1)
 210         {
 211             size = 0;
 212             vlog("failed to ioctl %d\n", rc);
 213         }
 214         close(fd);
 215     }
 216     else
 217     {
 218         vlog("failed to open %s %d\n", diskpath, errno);
 219     }
 220
 221     vlog("disk %s size %llu bytes\n", disk, (ulonglong)size);
 222     return size;
 223 }
 224
 225 static int get_disk_num(void)
 226 {
 227     int n = 0;
 228     DIR* dir = NULL;
 229     struct dirent* p = NULL;
 230
 231     dir = opendir("/sys/block");
 232     if (!dir)
 233     {
 234         return 0;
 235     }
 236
 237     while ((p = readdir(dir)) != NULL)
 238     {
 239         n++;
 240     }
 241     closedir(dir);
 242
 243     return n;
 244 }
 245
 246 static int is_usb_disk(const char *diskname)
 247 {
 248     int rc;
 249     char dstpath[1024] = { 0 };
 250     char syspath[1024] = { 0 };
 251
 252     snprintf(syspath, sizeof(syspath), "/sys/block/%s", diskname);
 253     rc = readlink(syspath, dstpath, sizeof(dstpath) - 1);
 254     if (rc > 0 && strstr(dstpath, "/usb"))
 255     {
 256         return 1;
 257     }
 258
 259     return 0;
 260 }
 261
 262 static int get_all_disk(void)
 263 {
 264     int i = 0;
 265     int j = 0;
 266     int num = 0;
 267     ulonglong cursize = 0;
 268     DIR* dir = NULL;
 269     struct dirent* p = NULL;
 270     disk_info *node = NULL;
 271     disk_info tmpnode;
 272
 273     num = get_disk_num();
 274     if (num <= 0)
 275     {
 276         return 1;
 277     }
 278
 279     g_disk_list = malloc(num * sizeof(disk_info));
 280     if (!g_disk_list)
 281     {
 282         return 1;
 283     }
 284     memset(g_disk_list, 0, num * sizeof(disk_info));
 285
 286     dir = opendir("/sys/block");
 287     if (!dir)
 288     {
 289         return 0;
 290     }
 291
 292     while (((p = readdir(dir)) != NULL) && g_disk_num < num)
 293     {
 294         if (!vtoy_is_possible_blkdev(p->d_name))
 295         {
 296             vlog("disk %s is filted by name\n", p->d_name);
 297             continue;
 298         }
 299
 300         cursize = vtoy_get_disk_size_in_byte(p->d_name);
 301
 302         node = g_disk_list + g_disk_num;
 303         g_disk_num++;
 304
 305
 306         snprintf(node->name, sizeof(node->name), p->d_name);
 307         node->size = cursize;
 308         node->isUSB = is_usb_disk(p->d_name);
 309         if (strncmp(node->name, "sd", 2) == 0)
 310         {
 311             node->isSDX = 1;
 312         }
 313     }
 314     closedir(dir);
 315
 316     /* sort disks */
 317     for (i = 0; i < g_disk_num; i++)
 318     {
 319         for (j = i + 1; j < g_disk_num; j++)
 320         {
 321             if (g_disk_list[i].isSDX && g_disk_list[j].isSDX)
 322             {
 323                 if (strcmp(g_disk_list[i].name, g_disk_list[j].name) > 0)
 324                 {
 325                     memcpy(&tmpnode, g_disk_list + i, sizeof(tmpnode));
 326                     memcpy(g_disk_list + i, g_disk_list + j, sizeof(tmpnode));
 327                     memcpy(g_disk_list + j, &tmpnode, sizeof(tmpnode));
 328                 }
 329             }
 330         }
 331     }
 332
 333     vlog("============ DISK DUMP BEGIN ===========\n");
 334     for (i = 0; i < g_disk_num; i++)
 335     {
 336         node = g_disk_list + i;
 337         vlog("[%d] %s %dGB(%llu) USB:%d\n", i, node->name,
 338             node->size / 1024 / 1024 / 1024, node->size, node->isUSB);
 339     }
 340     vlog("============ DISK DUMP END ===========\n");
 341
 342     return 0;
 343 }
 344
 345 static int expand_var(const char *var, char *value, int len)
 346 {
 347     int i;
 348     int index = -1;
 349     ulonglong uiDst = 0;
 350     ulonglong delta = 0;
 351     ulonglong maxsize = 0;
 352     ulonglong maxdelta = 0xFFFFFFFFFFFFFFFFULL;
 353     disk_info *node = NULL;
 354     value[0] = 0;
 355
 356     if (strcmp(var, "VT_LINUX_DISK_SDX_1ST_NONVTOY") == 0)
 357     {
 358         for (i = 0; i < g_disk_num; i++)
 359         {
 360             node = g_disk_list + i;
 361             if (node->size > 0 && node->isSDX && strcmp(node->name, g_vtoy_disk_name) != 0)
 362             {
 363                 vlog("%s=<%s>\n", var, node->name);
 364                 snprintf(value, len, "%s", node->name);
 365                 return 0;
 366             }
 367         }
 368
 369         vlog("[Error] %s not found\n", var);
 370     }
 371     else if (strcmp(var, "VT_LINUX_DISK_SDX_1ST_NONUSB") == 0)
 372     {
 373         for (i = 0; i < g_disk_num; i++)
 374         {
 375             node = g_disk_list + i;
 376             if (node->size > 0 && node->isSDX && node->isUSB == 0)
 377             {
 378                 vlog("%s=<%s>\n", var, node->name);
 379                 snprintf(value, len, "%s", node->name);
 380                 return 0;
 381             }
 382         }
 383
 384         vlog("[Error] %s not found\n", var);
 385     }
 386     else if (strcmp(var, "VT_LINUX_DISK_MAX_SIZE") == 0)
 387     {
 388         for (i = 0; i < g_disk_num; i++)
 389         {
 390             node = g_disk_list + i;
 391             if (node->size > 0 && node->size > maxsize)
 392             {
 393                 index = i;
 394                 maxsize = node->size;
 395             }
 396         }
 397
 398         if (index >= 0)
 399         {
 400             vlog("%s=<%s>\n", var, g_disk_list[index].name);
 401             snprintf(value, len, "%s", g_disk_list[index].name);
 402             return 0;
 403         }
 404         else
 405         {
 406             vlog("[Error] %s not found\n", var);
 407         }
 408     }
 409     else if (strncmp(var, "VT_LINUX_DISK_CLOSEST_", 22) == 0)
 410     {
 411         uiDst = strtoul(var + 22, NULL, 10);
 412         uiDst = uiDst * (1024ULL * 1024ULL * 1024ULL);
 413
 414         for (i = 0; i < g_disk_num; i++)
 415         {
 416             node = g_disk_list + i;
 417             if (node->size == 0)
 418             {
 419                 continue;
 420             }
 421
 422             if (node->size > uiDst)
 423             {
 424                 delta = node->size - uiDst;
 425             }
 426             else
 427             {
 428                 delta = uiDst - node->size;
 429             }
 430
 431             if (delta < maxdelta)
 432             {
 433                 index = i;
 434                 maxdelta = delta;
 435             }
 436         }
 437
 438         if (index >= 0)
 439         {
 440             vlog("%s=<%s>\n", var, g_disk_list[index].name);
 441             snprintf(value, len, "%s", g_disk_list[index].name);
 442             return 0;
 443         }
 444         else
 445         {
 446             vlog("[Error] %s not found\n", var);
 447         }
 448     }
 449     else
 450     {
 451         vlog("Invalid var name <%s>\n", var);
 452         snprintf(value, len, "$$%s$$", var);
 453     }
 454
 455     if (value[0] == 0)
 456     {
 457         snprintf(value, len, "$$%s$$", var);
 458     }
 459
 460     return 0;
 461 }
 462
 463 int vtoyexpand_main(int argc, char **argv)
 464 {
 465     FILE *fp = NULL;
 466     FILE *fout = NULL;
 467     char *start = NULL;
 468     char *end = NULL;
 469     char line[4096];
 470     char value[256];
 471
 472     vlog("========= vtoyexpand_main %d ========\n", argc);
 473
 474     if (argc != 3)
 475     {
 476         return 1;
 477     }
 478
 479     g_vtoy_disk_name = argv[2];
 480     if (strncmp(g_vtoy_disk_name, "/dev/", 5) == 0)
 481     {
 482         g_vtoy_disk_name += 5;
 483     }
 484     vlog("<%s> <%s> <%s>\n", argv[1], argv[2], g_vtoy_disk_name);
 485
 486     get_all_disk();
 487
 488     fp = fopen(argv[1], "r");
 489     if (!fp)
 490     {
 491         vlog("Failed to open file <%s>\n", argv[1]);
 492         return 1;
 493     }
 494
 495     fout = fopen(TMP_FILE, "w+");
 496     if (!fout)
 497     {
 498         vlog("Failed to create file <%s>\n", TMP_FILE);
 499         fclose(fp);
 500         return 1;
 501     }
 502
 503     memset(line, 0, sizeof(line));
 504     memset(value, 0, sizeof(value));
 505
 506     while (fgets(line, sizeof(line), fp))
 507     {
 508         start = strstr(line, "$$VT_");
 509         if (start)
 510         {
 511             end = strstr(start + 5, "$$");
 512         }
 513
 514         if (start && end)
 515         {
 516             *start = 0;
 517             fprintf(fout, "%s", line);
 518
 519             *end = 0;
 520             expand_var(start + 2, value, sizeof(value));
 521             fprintf(fout, "%s", value);
 522
 523             fprintf(fout, "%s", end + 2);
 524
 525             memset(value, 0, sizeof(value));
 526         }
 527         else
 528         {
 529             fprintf(fout, "%s", line);
 530         }
 531
 532         line[0] = line[4095] = 0;
 533     }
 534
 535     fclose(fp);
 536     fclose(fout);
 537
 538     vlog("delete file <%s>\n", argv[1]);
 539     remove(argv[1]);
 540
 541     vlog("Copy file <%s> --> <%s>\n", TMP_FILE, argv[1]);
 542     copy_file(TMP_FILE, argv[1]);
 543
 544     return 0;
 545 }
 546
 547 // wrapper main
 548 #ifndef BUILD_VTOY_TOOL
 549 int main(int argc, char **argv)
 550 {
 551     return vtoyexpand_main(argc, argv);
 552 }
 553 #endif
 554