Ventoy/VtoyTool/vtoyexpand.c
2022-06-19 13:56:49 +08:00

555 lines
12 KiB
C

/******************************************************************************
* vtoyexpand.c ---- ventoy auto install script variable expansion
*
* Copyright (c) 2022, longpanda <admin@ventoy.net>
*
* This program is free software; you can redistribute it and/or
* modify it under the terms of the GNU General Public License as
* published by the Free Software Foundation; either version 3 of the
* License, or (at your option) any later version.
*
* This program is distributed in the hope that it will be useful, but
* WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
* General Public License for more details.
*
* You should have received a copy of the GNU General Public License
* along with this program; if not, see <http://www.gnu.org/licenses/>.
*
*/
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdarg.h>
#include <errno.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/types.h>
#include <sys/mman.h>
#include <sys/ioctl.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <linux/fs.h>
#include <dirent.h>
#include "vtoytool.h"
#ifndef O_BINARY
#define O_BINARY 0
#endif
#define TMP_FILE "/ventoy/tmp_var_expansion"
#define SIZE_1MB (1024 * 1024)
#define ulong unsigned long
#define ulonglong unsigned long long
typedef struct disk_info
{
char name[128];
ulonglong size;
int isUSB;
int isSDX;
}disk_info;
static disk_info *g_disk_list = NULL;
static int g_disk_num = 0;
static const char *g_vtoy_disk_name = NULL;
static void vlog(const char *fmt, ...)
{
int n = 0;
va_list arg;
FILE *fp = NULL;
char log[1024];
fp = fopen("/ventoy/autoinstall.log", "a+");
if (fp)
{
va_start(arg, fmt);
n += vsnprintf(log, sizeof(log) - 1, fmt, arg);
va_end(arg);
fwrite(log, 1, n, fp);
fclose(fp);
}
}
static int copy_file(const char *file1, const char *file2)
{
int n;
int size;
int ret = 1;
FILE *fp1 = NULL;
FILE *fp2 = NULL;
char *buf = NULL;
fp1 = fopen(file1, "rb");
if (!fp1)
{
vlog("Failed to read file <%s>\n", file1);
goto end;
}
fp2 = fopen(file2, "wb+");
if (!fp2)
{
vlog("Failed to create file <%s>\n", file2);
goto end;
}
fseek(fp1, 0, SEEK_END);
size = (int)ftell(fp1);
fseek(fp1, 0, SEEK_SET);
buf = malloc(size);
if (!buf)
{
vlog("Failed to malloc buf\n");
goto end;
}
n = fread(buf, 1, size, fp1);
if (n != size)
{
vlog("Failed to read <%s> %d %d\n", file1, n, size);
goto end;
}
n = fwrite(buf, 1, size, fp2);
if (n != size)
{
vlog("Failed to write <%s> %d %d\n", file2, n, size);
goto end;
}
ret = 0;
end:
if (fp1)
fclose(fp1);
if (fp2)
fclose(fp2);
if (buf)
free(buf);
return ret;
}
static int vtoy_is_possible_blkdev(const char *name)
{
if (name[0] == '.')
{
return 0;
}
/* /dev/ramX */
if (name[0] == 'r' && name[1] == 'a' && name[2] == 'm')
{
return 0;
}
/* /dev/loopX */
if (name[0] == 'l' && name[1] == 'o' && name[2] == 'o' && name[3] == 'p')
{
return 0;
}
/* /dev/dm-X */
if (name[0] == 'd' && name[1] == 'm' && name[2] == '-' && IS_DIGIT(name[3]))
{
return 0;
}
/* /dev/srX */
if (name[0] == 's' && name[1] == 'r' && (name[2] >= '0' && name[2] <= '9'))
{
return 0;
}
return 1;
}
static ulonglong vtoy_get_disk_size_in_byte(const char *disk)
{
int fd;
int rc;
unsigned long long size = 0;
char diskpath[256] = {0};
char sizebuf[64] = {0};
// Try 1: get size from sysfs
snprintf(diskpath, sizeof(diskpath) - 1, "/sys/block/%s/size", disk);
if (access(diskpath, F_OK) >= 0)
{
vlog("get disk size from sysfs for %s\n", disk);
fd = open(diskpath, O_RDONLY | O_BINARY);
if (fd >= 0)
{
read(fd, sizebuf, sizeof(sizebuf));
size = strtoull(sizebuf, NULL, 10);
close(fd);
return (size * 512);
}
}
else
{
vlog("%s not exist \n", diskpath);
}
// Try 2: get size from ioctl
snprintf(diskpath, sizeof(diskpath) - 1, "/dev/%s", disk);
fd = open(diskpath, O_RDONLY);
if (fd >= 0)
{
vlog("get disk size from ioctl for %s\n", disk);
rc = ioctl(fd, BLKGETSIZE64, &size);
if (rc == -1)
{
size = 0;
vlog("failed to ioctl %d\n", rc);
}
close(fd);
}
else
{
vlog("failed to open %s %d\n", diskpath, errno);
}
vlog("disk %s size %llu bytes\n", disk, (ulonglong)size);
return size;
}
static int get_disk_num(void)
{
int n = 0;
DIR* dir = NULL;
struct dirent* p = NULL;
dir = opendir("/sys/block");
if (!dir)
{
return 0;
}
while ((p = readdir(dir)) != NULL)
{
n++;
}
closedir(dir);
return n;
}
static int is_usb_disk(const char *diskname)
{
int rc;
char dstpath[1024] = { 0 };
char syspath[1024] = { 0 };
snprintf(syspath, sizeof(syspath), "/sys/block/%s", diskname);
rc = readlink(syspath, dstpath, sizeof(dstpath) - 1);
if (rc > 0 && strstr(dstpath, "/usb"))
{
return 1;
}
return 0;
}
static int get_all_disk(void)
{
int i = 0;
int j = 0;
int num = 0;
ulonglong cursize = 0;
DIR* dir = NULL;
struct dirent* p = NULL;
disk_info *node = NULL;
disk_info tmpnode;
num = get_disk_num();
if (num <= 0)
{
return 1;
}
g_disk_list = malloc(num * sizeof(disk_info));
if (!g_disk_list)
{
return 1;
}
memset(g_disk_list, 0, num * sizeof(disk_info));
dir = opendir("/sys/block");
if (!dir)
{
return 0;
}
while (((p = readdir(dir)) != NULL) && g_disk_num < num)
{
if (!vtoy_is_possible_blkdev(p->d_name))
{
vlog("disk %s is filted by name\n", p->d_name);
continue;
}
cursize = vtoy_get_disk_size_in_byte(p->d_name);
node = g_disk_list + g_disk_num;
g_disk_num++;
snprintf(node->name, sizeof(node->name), p->d_name);
node->size = cursize;
node->isUSB = is_usb_disk(p->d_name);
if (strncmp(node->name, "sd", 2) == 0)
{
node->isSDX = 1;
}
}
closedir(dir);
/* sort disks */
for (i = 0; i < g_disk_num; i++)
{
for (j = i + 1; j < g_disk_num; j++)
{
if (g_disk_list[i].isSDX && g_disk_list[j].isSDX)
{
if (strcmp(g_disk_list[i].name, g_disk_list[j].name) > 0)
{
memcpy(&tmpnode, g_disk_list + i, sizeof(tmpnode));
memcpy(g_disk_list + i, g_disk_list + j, sizeof(tmpnode));
memcpy(g_disk_list + j, &tmpnode, sizeof(tmpnode));
}
}
}
}
vlog("============ DISK DUMP BEGIN ===========\n");
for (i = 0; i < g_disk_num; i++)
{
node = g_disk_list + i;
vlog("[%d] %s %dGB(%llu) USB:%d\n", i, node->name,
node->size / 1024 / 1024 / 1024, node->size, node->isUSB);
}
vlog("============ DISK DUMP END ===========\n");
return 0;
}
static int expand_var(const char *var, char *value, int len)
{
int i;
int index = -1;
ulonglong uiDst = 0;
ulonglong delta = 0;
ulonglong maxsize = 0;
ulonglong maxdelta = 0xFFFFFFFFFFFFFFFFULL;
disk_info *node = NULL;
value[0] = 0;
if (strcmp(var, "VT_LINUX_DISK_SDX_1ST_NONVTOY") == 0)
{
for (i = 0; i < g_disk_num; i++)
{
node = g_disk_list + i;
if (node->size > 0 && node->isSDX && strcmp(node->name, g_vtoy_disk_name) != 0)
{
vlog("%s=<%s>\n", var, node->name);
snprintf(value, len, "%s", node->name);
return 0;
}
}
vlog("[Error] %s not found\n", var);
}
else if (strcmp(var, "VT_LINUX_DISK_SDX_1ST_NONUSB") == 0)
{
for (i = 0; i < g_disk_num; i++)
{
node = g_disk_list + i;
if (node->size > 0 && node->isSDX && node->isUSB == 0)
{
vlog("%s=<%s>\n", var, node->name);
snprintf(value, len, "%s", node->name);
return 0;
}
}
vlog("[Error] %s not found\n", var);
}
else if (strcmp(var, "VT_LINUX_DISK_MAX_SIZE") == 0)
{
for (i = 0; i < g_disk_num; i++)
{
node = g_disk_list + i;
if (node->size > 0 && node->size > maxsize)
{
index = i;
maxsize = node->size;
}
}
if (index >= 0)
{
vlog("%s=<%s>\n", var, g_disk_list[index].name);
snprintf(value, len, "%s", g_disk_list[index].name);
return 0;
}
else
{
vlog("[Error] %s not found\n", var);
}
}
else if (strncmp(var, "VT_LINUX_DISK_CLOSEST_", 22) == 0)
{
uiDst = strtoul(var + 22, NULL, 10);
uiDst = uiDst * (1024ULL * 1024ULL * 1024ULL);
for (i = 0; i < g_disk_num; i++)
{
node = g_disk_list + i;
if (node->size == 0)
{
continue;
}
if (node->size > uiDst)
{
delta = node->size - uiDst;
}
else
{
delta = uiDst - node->size;
}
if (delta < maxdelta)
{
index = i;
maxdelta = delta;
}
}
if (index >= 0)
{
vlog("%s=<%s>\n", var, g_disk_list[index].name);
snprintf(value, len, "%s", g_disk_list[index].name);
return 0;
}
else
{
vlog("[Error] %s not found\n", var);
}
}
else
{
vlog("Invalid var name <%s>\n", var);
snprintf(value, len, "$$%s$$", var);
}
if (value[0] == 0)
{
snprintf(value, len, "$$%s$$", var);
}
return 0;
}
int vtoyexpand_main(int argc, char **argv)
{
FILE *fp = NULL;
FILE *fout = NULL;
char *start = NULL;
char *end = NULL;
char line[4096];
char value[256];
vlog("========= vtoyexpand_main %d ========\n", argc);
if (argc != 3)
{
return 1;
}
g_vtoy_disk_name = argv[2];
if (strncmp(g_vtoy_disk_name, "/dev/", 5) == 0)
{
g_vtoy_disk_name += 5;
}
vlog("<%s> <%s> <%s>\n", argv[1], argv[2], g_vtoy_disk_name);
get_all_disk();
fp = fopen(argv[1], "r");
if (!fp)
{
vlog("Failed to open file <%s>\n", argv[1]);
return 1;
}
fout = fopen(TMP_FILE, "w+");
if (!fout)
{
vlog("Failed to create file <%s>\n", TMP_FILE);
fclose(fp);
return 1;
}
memset(line, 0, sizeof(line));
memset(value, 0, sizeof(value));
while (fgets(line, sizeof(line), fp))
{
start = strstr(line, "$$VT_");
if (start)
{
end = strstr(start + 5, "$$");
}
if (start && end)
{
*start = 0;
fprintf(fout, "%s", line);
*end = 0;
expand_var(start + 2, value, sizeof(value));
fprintf(fout, "%s", value);
fprintf(fout, "%s", end + 2);
memset(value, 0, sizeof(value));
}
else
{
fprintf(fout, "%s", line);
}
line[0] = line[4095] = 0;
}
fclose(fp);
fclose(fout);
vlog("delete file <%s>\n", argv[1]);
remove(argv[1]);
vlog("Copy file <%s> --> <%s>\n", TMP_FILE, argv[1]);
copy_file(TMP_FILE, argv[1]);
return 0;
}
// wrapper main
#ifndef BUILD_VTOY_TOOL
int main(int argc, char **argv)
{
return vtoyexpand_main(argc, argv);
}
#endif