wimlib.net Git - wimlib/blob - programs/imagex.c

   1 /*
   2  * imagex.c
   3  *
   4  * Use wimlib to create, modify, extract, mount, unmount, or display information
   5  * about a WIM file
   6  */
   7
   8 /*
   9  * Copyright (C) 2012, 2013 Eric Biggers
  10  *
  11  * This program is free software: you can redistribute it and/or modify
  12  * it under the terms of the GNU General Public License as published by
  13  * the Free Software Foundation, either version 3 of the License, or
  14  * (at your option) any later version.
  15  *
  16  * This program is distributed in the hope that it will be useful,
  17  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  18  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  19  * GNU General Public License for more details.
  20  *
  21  * You should have received a copy of the GNU General Public License
  22  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  23  */
  24
  25 #include "config.h"
  26 #include "wimlib.h"
  27 #include "wimlib_tchar.h"
  28
  29 #include <ctype.h>
  30 #include <errno.h>
  31
  32 #include <inttypes.h>
  33 #include <libgen.h>
  34 #include <limits.h>
  35 #include <stdarg.h>
  36 #include <stdlib.h>
  37 #include <string.h>
  38 #include <sys/stat.h>
  39 #include <unistd.h>
  40 #include <locale.h>
  41
  42 #ifdef HAVE_ALLOCA_H
  43 #include <alloca.h>
  44 #endif
  45
  46 #ifdef __WIN32__
  47 #  include "imagex-win32.h"
  48 #  define tbasename     win32_wbasename
  49 #  define tglob         win32_wglob
  50 #else /* __WIN32__ */
  51 #  include <glob.h>
  52 #  include <getopt.h>
  53 #  include <langinfo.h>
  54 #  define tbasename     basename
  55 #  define tglob         glob
  56 #endif /* !__WIN32 */
  57
  58
  59 #define ARRAY_LEN(array) (sizeof(array) / sizeof(array[0]))
  60
  61 #define for_opt(c, opts) while ((c = getopt_long_only(argc, (tchar**)argv, T(""), \
  62                                 opts, NULL)) != -1)
  63
  64 enum imagex_op_type {
  65         APPEND = 0,
  66         APPLY,
  67         CAPTURE,
  68         DELETE,
  69         DIR,
  70         EXPORT,
  71         INFO,
  72         JOIN,
  73         MOUNT,
  74         MOUNTRW,
  75         OPTIMIZE,
  76         SPLIT,
  77         UNMOUNT,
  78 };
  79
  80 static void usage(int cmd_type);
  81 static void usage_all();
  82
  83
  84 static const tchar *usage_strings[] = {
  85 [APPEND] =
  86 T(
  87 IMAGEX_PROGNAME" append (DIRECTORY | NTFS_VOLUME) WIMFILE [IMAGE_NAME]\n"
  88 "                     [DESCRIPTION] [--boot] [--check] [--flags EDITION_ID]\n"
  89 "                     [--verbose] [--dereference] [--config=FILE]\n"
  90 "                     [--threads=NUM_THREADS] [--rebuild] [--unix-data]\n"
  91 "                     [--source-list] [--no-acls] [--strict-acls]\n"
  92 "                     [--rpfix] [--norpfix]\n"
  93 ),
  94 [APPLY] =
  95 T(
  96 IMAGEX_PROGNAME" apply WIMFILE [IMAGE_NUM | IMAGE_NAME | all]\n"
  97 "                    (DIRECTORY | NTFS_VOLUME) [--check] [--hardlink]\n"
  98 "                    [--symlink] [--verbose] [--ref=\"GLOB\"] [--unix-data]\n"
  99 "                    [--no-acls] [--strict-acls] [--rpfix] [--norpfix]\n"
 100 ),
 101 [CAPTURE] =
 102 T(
 103 IMAGEX_PROGNAME" capture (DIRECTORY | NTFS_VOLUME) WIMFILE [IMAGE_NAME]\n"
 104 "                      [DESCRIPTION] [--boot] [--check] [--compress=TYPE]\n"
 105 "                      [--flags EDITION_ID] [--verbose] [--dereference]\n"
 106 "                      [--config=FILE] [--threads=NUM_THREADS] [--unix-data]\n"
 107 "                      [--source-list] [--no-acls] [--strict-acls]\n"
 108 "                      [--rpfix] [--norpfix]\n"
 109 ),
 110 [DELETE] =
 111 T(
 112 IMAGEX_PROGNAME" delete WIMFILE (IMAGE_NUM | IMAGE_NAME | all) [--check] [--soft]\n"
 113 ),
 114 [DIR] =
 115 T(
 116 IMAGEX_PROGNAME" dir WIMFILE (IMAGE_NUM | IMAGE_NAME | all)\n"
 117 ),
 118 [EXPORT] =
 119 T(
 120 IMAGEX_PROGNAME" export SRC_WIMFILE (SRC_IMAGE_NUM | SRC_IMAGE_NAME | all ) \n"
 121 "              DEST_WIMFILE [DEST_IMAGE_NAME] [DEST_IMAGE_DESCRIPTION]\n"
 122 "              [--boot] [--check] [--compress=TYPE] [--ref=\"GLOB\"]\n"
 123 "              [--threads=NUM_THREADS] [--rebuild]\n"
 124 ),
 125 [INFO] =
 126 T(
 127 IMAGEX_PROGNAME" info WIMFILE [IMAGE_NUM | IMAGE_NAME] [NEW_NAME]\n"
 128 "                   [NEW_DESC] [--boot] [--check] [--header] [--lookup-table]\n"
 129 "                   [--xml] [--extract-xml FILE] [--metadata]\n"
 130 ),
 131 [JOIN] =
 132 T(
 133 IMAGEX_PROGNAME" join [--check] WIMFILE SPLIT_WIM...\n"
 134 ),
 135 [MOUNT] =
 136 T(
 137 IMAGEX_PROGNAME" mount WIMFILE (IMAGE_NUM | IMAGE_NAME) DIRECTORY\n"
 138 "                    [--check] [--debug] [--streams-interface=INTERFACE]\n"
 139 "                    [--ref=\"GLOB\"] [--unix-data] [--allow-other]\n"
 140 ),
 141 [MOUNTRW] =
 142 T(
 143 IMAGEX_PROGNAME" mountrw WIMFILE [IMAGE_NUM | IMAGE_NAME] DIRECTORY\n"
 144 "                      [--check] [--debug] [--streams-interface=INTERFACE]\n"
 145 "                      [--staging-dir=DIR] [--unix-data] [--allow-other]\n"
 146 ),
 147 [OPTIMIZE] =
 148 T(
 149 IMAGEX_PROGNAME" optimize WIMFILE [--check] [--recompress]\n"
 150 "                      [--threads=NUM_THREADS]\n"
 151 ),
 152 [SPLIT] =
 153 T(
 154 IMAGEX_PROGNAME" split WIMFILE SPLIT_WIMFILE PART_SIZE_MB [--check]\n"
 155 ),
 156 [UNMOUNT] =
 157 T(
 158 IMAGEX_PROGNAME" unmount DIRECTORY [--commit] [--check] [--rebuild]\n"
 159 ),
 160 };
 161
 162 enum {
 163         IMAGEX_ALLOW_OTHER_OPTION,
 164         IMAGEX_BOOT_OPTION,
 165         IMAGEX_CHECK_OPTION,
 166         IMAGEX_COMMIT_OPTION,
 167         IMAGEX_COMPRESS_OPTION,
 168         IMAGEX_CONFIG_OPTION,
 169         IMAGEX_DEBUG_OPTION,
 170         IMAGEX_DEREFERENCE_OPTION,
 171         IMAGEX_EXTRACT_XML_OPTION,
 172         IMAGEX_FLAGS_OPTION,
 173         IMAGEX_HARDLINK_OPTION,
 174         IMAGEX_HEADER_OPTION,
 175         IMAGEX_LOOKUP_TABLE_OPTION,
 176         IMAGEX_METADATA_OPTION,
 177         IMAGEX_NO_ACLS_OPTION,
 178         IMAGEX_NORPFIX_OPTION,
 179         IMAGEX_REBULID_OPTION,
 180         IMAGEX_RECOMPRESS_OPTION,
 181         IMAGEX_REF_OPTION,
 182         IMAGEX_RPFIX_OPTION,
 183         IMAGEX_SOFT_OPTION,
 184         IMAGEX_SOURCE_LIST_OPTION,
 185         IMAGEX_STAGING_DIR_OPTION,
 186         IMAGEX_STREAMS_INTERFACE_OPTION,
 187         IMAGEX_STRICT_ACLS_OPTION,
 188         IMAGEX_SYMLINK_OPTION,
 189         IMAGEX_THREADS_OPTION,
 190         IMAGEX_UNIX_DATA_OPTION,
 191         IMAGEX_VERBOSE_OPTION,
 192         IMAGEX_XML_OPTION,
 193 };
 194
 195 static const struct option apply_options[] = {
 196         {T("check"),       no_argument,       NULL, IMAGEX_CHECK_OPTION},
 197         {T("hardlink"),    no_argument,       NULL, IMAGEX_HARDLINK_OPTION},
 198         {T("symlink"),     no_argument,       NULL, IMAGEX_SYMLINK_OPTION},
 199         {T("verbose"),     no_argument,       NULL, IMAGEX_VERBOSE_OPTION},
 200         {T("ref"),         required_argument, NULL, IMAGEX_REF_OPTION},
 201         {T("unix-data"),   no_argument,       NULL, IMAGEX_UNIX_DATA_OPTION},
 202         {T("noacls"),      no_argument,       NULL, IMAGEX_NO_ACLS_OPTION},
 203         {T("no-acls"),     no_argument,       NULL, IMAGEX_NO_ACLS_OPTION},
 204         {T("strict-acls"), no_argument,       NULL, IMAGEX_STRICT_ACLS_OPTION},
 205         {T("rpfix"),       no_argument,       NULL, IMAGEX_RPFIX_OPTION},
 206         {T("norpfix"),     no_argument,       NULL, IMAGEX_NORPFIX_OPTION},
 207         {NULL, 0, NULL, 0},
 208 };
 209 static const struct option capture_or_append_options[] = {
 210         {T("boot"),        no_argument,       NULL, IMAGEX_BOOT_OPTION},
 211         {T("check"),       no_argument,       NULL, IMAGEX_CHECK_OPTION},
 212         {T("compress"),    required_argument, NULL, IMAGEX_COMPRESS_OPTION},
 213         {T("config"),      required_argument, NULL, IMAGEX_CONFIG_OPTION},
 214         {T("dereference"), no_argument,       NULL, IMAGEX_DEREFERENCE_OPTION},
 215         {T("flags"),       required_argument, NULL, IMAGEX_FLAGS_OPTION},
 216         {T("verbose"),     no_argument,       NULL, IMAGEX_VERBOSE_OPTION},
 217         {T("threads"),     required_argument, NULL, IMAGEX_THREADS_OPTION},
 218         {T("rebuild"),     no_argument,       NULL, IMAGEX_REBULID_OPTION},
 219         {T("unix-data"),   no_argument,       NULL, IMAGEX_UNIX_DATA_OPTION},
 220         {T("source-list"), no_argument,       NULL, IMAGEX_SOURCE_LIST_OPTION},
 221         {T("noacls"),      no_argument,       NULL, IMAGEX_NO_ACLS_OPTION},
 222         {T("no-acls"),     no_argument,       NULL, IMAGEX_NO_ACLS_OPTION},
 223         {T("strict-acls"), no_argument,       NULL, IMAGEX_STRICT_ACLS_OPTION},
 224         {T("rpfix"),       no_argument,       NULL, IMAGEX_RPFIX_OPTION},
 225         {T("norpfix"),     no_argument,       NULL, IMAGEX_NORPFIX_OPTION},
 226         {NULL, 0, NULL, 0},
 227 };
 228 static const struct option delete_options[] = {
 229         {T("check"), no_argument, NULL, IMAGEX_CHECK_OPTION},
 230         {T("soft"),  no_argument, NULL, IMAGEX_SOFT_OPTION},
 231         {NULL, 0, NULL, 0},
 232 };
 233
 234 static const struct option export_options[] = {
 235         {T("boot"),       no_argument,       NULL, IMAGEX_BOOT_OPTION},
 236         {T("check"),      no_argument,       NULL, IMAGEX_CHECK_OPTION},
 237         {T("compress"),   required_argument, NULL, IMAGEX_COMPRESS_OPTION},
 238         {T("ref"),        required_argument, NULL, IMAGEX_REF_OPTION},
 239         {T("threads"),    required_argument, NULL, IMAGEX_THREADS_OPTION},
 240         {T("rebuild"),    no_argument,       NULL, IMAGEX_REBULID_OPTION},
 241         {NULL, 0, NULL, 0},
 242 };
 243
 244 static const struct option info_options[] = {
 245         {T("boot"),         no_argument,       NULL, IMAGEX_BOOT_OPTION},
 246         {T("check"),        no_argument,       NULL, IMAGEX_CHECK_OPTION},
 247         {T("extract-xml"),  required_argument, NULL, IMAGEX_EXTRACT_XML_OPTION},
 248         {T("header"),       no_argument,       NULL, IMAGEX_HEADER_OPTION},
 249         {T("lookup-table"), no_argument,       NULL, IMAGEX_LOOKUP_TABLE_OPTION},
 250         {T("metadata"),     no_argument,       NULL, IMAGEX_METADATA_OPTION},
 251         {T("xml"),          no_argument,       NULL, IMAGEX_XML_OPTION},
 252         {NULL, 0, NULL, 0},
 253 };
 254
 255 static const struct option join_options[] = {
 256         {T("check"), no_argument, NULL, IMAGEX_CHECK_OPTION},
 257         {NULL, 0, NULL, 0},
 258 };
 259
 260 static const struct option mount_options[] = {
 261         {T("check"),             no_argument,       NULL, IMAGEX_CHECK_OPTION},
 262         {T("debug"),             no_argument,       NULL, IMAGEX_DEBUG_OPTION},
 263         {T("streams-interface"), required_argument, NULL, IMAGEX_STREAMS_INTERFACE_OPTION},
 264         {T("ref"),               required_argument, NULL, IMAGEX_REF_OPTION},
 265         {T("staging-dir"),       required_argument, NULL, IMAGEX_STAGING_DIR_OPTION},
 266         {T("unix-data"),         no_argument,       NULL, IMAGEX_UNIX_DATA_OPTION},
 267         {T("allow-other"),       no_argument,       NULL, IMAGEX_ALLOW_OTHER_OPTION},
 268         {NULL, 0, NULL, 0},
 269 };
 270
 271 static const struct option optimize_options[] = {
 272         {T("check"),      no_argument, NULL, IMAGEX_CHECK_OPTION},
 273         {T("recompress"), no_argument, NULL, IMAGEX_RECOMPRESS_OPTION},
 274         {T("threads"),    required_argument, NULL, IMAGEX_THREADS_OPTION},
 275         {NULL, 0, NULL, 0},
 276 };
 277
 278 static const struct option split_options[] = {
 279         {T("check"), no_argument, NULL, IMAGEX_CHECK_OPTION},
 280         {NULL, 0, NULL, 0},
 281 };
 282
 283 static const struct option unmount_options[] = {
 284         {T("commit"),  no_argument, NULL, IMAGEX_COMMIT_OPTION},
 285         {T("check"),   no_argument, NULL, IMAGEX_CHECK_OPTION},
 286         {T("rebuild"), no_argument, NULL, IMAGEX_REBULID_OPTION},
 287         {NULL, 0, NULL, 0},
 288 };
 289
 290
 291
 292 /* Print formatted error message to stderr. */
 293 static void
 294 imagex_error(const tchar *format, ...)
 295 {
 296         va_list va;
 297         va_start(va, format);
 298         tfputs(T("ERROR: "), stderr);
 299         tvfprintf(stderr, format, va);
 300         tputc(T('\n'), stderr);
 301         va_end(va);
 302 }
 303
 304 /* Print formatted error message to stderr. */
 305 static void
 306 imagex_error_with_errno(const tchar *format, ...)
 307 {
 308         int errno_save = errno;
 309         va_list va;
 310         va_start(va, format);
 311         tfputs(T("ERROR: "), stderr);
 312         tvfprintf(stderr, format, va);
 313         tfprintf(stderr, T(": %"TS"\n"), tstrerror(errno_save));
 314         va_end(va);
 315 }
 316
 317 static int
 318 verify_image_exists(int image, const tchar *image_name, const tchar *wim_name)
 319 {
 320         if (image == WIMLIB_NO_IMAGE) {
 321                 imagex_error(T("\"%"TS"\" is not a valid image in \"%"TS"\"!\n"
 322                              "       Please specify a 1-based image index or "
 323                              "image name.\n"
 324                              "       You may use `"IMAGEX_PROGNAME" info' to list the images "
 325                              "contained in a WIM."),
 326                              image_name, wim_name);
 327                 return -1;
 328         }
 329         return 0;
 330 }
 331
 332 static int
 333 verify_image_is_single(int image)
 334 {
 335         if (image == WIMLIB_ALL_IMAGES) {
 336                 imagex_error(T("Cannot specify all images for this action!"));
 337                 return -1;
 338         }
 339         return 0;
 340 }
 341
 342 static int
 343 verify_image_exists_and_is_single(int image, const tchar *image_name,
 344                                   const tchar *wim_name)
 345 {
 346         int ret;
 347         ret = verify_image_exists(image, image_name, wim_name);
 348         if (ret == 0)
 349                 ret = verify_image_is_single(image);
 350         return ret;
 351 }
 352
 353 /* Parse the argument to --compress */
 354 static int
 355 get_compression_type(const tchar *optarg)
 356 {
 357         if (!tstrcasecmp(optarg, T("maximum")) || !tstrcasecmp(optarg, T("lzx")))
 358                 return WIMLIB_COMPRESSION_TYPE_LZX;
 359         else if (!tstrcasecmp(optarg, T("fast")) || !tstrcasecmp(optarg, T("xpress")))
 360                 return WIMLIB_COMPRESSION_TYPE_XPRESS;
 361         else if (!tstrcasecmp(optarg, T("none")))
 362                 return WIMLIB_COMPRESSION_TYPE_NONE;
 363         else {
 364                 imagex_error(T("Invalid compression type \"%"TS"\"! Must be "
 365                              "\"maximum\", \"fast\", or \"none\"."), optarg);
 366                 return WIMLIB_COMPRESSION_TYPE_INVALID;
 367         }
 368 }
 369
 370 /* Returns the size of a file given its name, or -1 if the file does not exist
 371  * or its size cannot be determined.  */
 372 static off_t
 373 file_get_size(const tchar *filename)
 374 {
 375         struct stat st;
 376         if (tstat(filename, &st) == 0)
 377                 return st.st_size;
 378         else
 379                 return (off_t)-1;
 380 }
 381
 382 tchar pat_ntfs_log[]                  = T("/$ntfs.log");
 383 tchar pat_hiberfil_sys[]              = T("/hiberfil.sys");
 384 tchar pat_pagefile_sys[]              = T("/pagefile.sys");
 385 tchar pat_system_volume_information[] = T("/System Volume Information");
 386 tchar pat_recycler[]                  = T("/RECYCLER");
 387 tchar pat_windows_csc[]               = T("/Windows/CSC");
 388
 389 tchar *default_pats[] = {
 390         pat_ntfs_log,
 391         pat_hiberfil_sys,
 392         pat_pagefile_sys,
 393         pat_system_volume_information,
 394         pat_recycler,
 395         pat_windows_csc,
 396 };
 397
 398 static struct wimlib_capture_config default_capture_config = {
 399         .exclusion_pats = {
 400                 .num_pats = sizeof(default_pats) / sizeof(*default_pats),
 401                 .pats = default_pats,
 402         },
 403 };
 404
 405 enum {
 406         PARSE_FILENAME_SUCCESS = 0,
 407         PARSE_FILENAME_FAILURE = 1,
 408         PARSE_FILENAME_NONE = 2,
 409 };
 410
 411 /*
 412  * Parses a filename in the source list file format.  (See the man page for
 413  * 'wimlib-imagex capture' for details on this format and the meaning.)
 414  * Accepted formats for filenames are an unquoted string (whitespace-delimited),
 415  * or a double or single-quoted string.
 416  *
 417  * @line_p:  Pointer to the pointer to the line of data.  Will be updated
 418  *           to point past the filename iff the return value is
 419  *           PARSE_FILENAME_SUCCESS.  If *len_p > 0, (*line_p)[*len_p - 1] must
 420  *           be '\0'.
 421  *
 422  * @len_p:   @len_p initially stores the length of the line of data, which may
 423  *           be 0, and it will be updated to the number of bytes remaining in
 424  *           the line iff the return value is PARSE_FILENAME_SUCCESS.
 425  *
 426  * @fn_ret:  Iff the return value is PARSE_FILENAME_SUCCESS, a pointer to the
 427  *           parsed filename will be returned here.
 428  *
 429  * Returns: PARSE_FILENAME_SUCCESS if a filename was successfully parsed; or
 430  *          PARSE_FILENAME_FAILURE if the data was invalid due to a missing
 431  *          closing quote; or PARSE_FILENAME_NONE if the line ended before the
 432  *          beginning of a filename was found.
 433  */
 434 static int
 435 parse_filename(tchar **line_p, size_t *len_p, tchar **fn_ret)
 436 {
 437         size_t len = *len_p;
 438         tchar *line = *line_p;
 439         tchar *fn;
 440         tchar quote_char;
 441
 442         /* Skip leading whitespace */
 443         for (;;) {
 444                 if (len == 0)
 445                         return PARSE_FILENAME_NONE;
 446                 if (!istspace(*line) && *line != T('\0'))
 447                         break;
 448                 line++;
 449                 len--;
 450         }
 451         quote_char = *line;
 452         if (quote_char == T('"') || quote_char == T('\'')) {
 453                 /* Quoted filename */
 454                 line++;
 455                 len--;
 456                 fn = line;
 457                 line = tmemchr(line, quote_char, len);
 458                 if (!line) {
 459                         imagex_error(T("Missing closing quote: %"TS), fn - 1);
 460                         return PARSE_FILENAME_FAILURE;
 461                 }
 462         } else {
 463                 /* Unquoted filename.  Go until whitespace.  Line is terminated
 464                  * by '\0', so no need to check 'len'. */
 465                 fn = line;
 466                 do {
 467                         line++;
 468                 } while (!istspace(*line) && *line != T('\0'));
 469         }
 470         *line = T('\0');
 471         len -= line - fn;
 472         *len_p = len;
 473         *line_p = line;
 474         *fn_ret = fn;
 475         return PARSE_FILENAME_SUCCESS;
 476 }
 477
 478 /* Parses a line of data (not an empty line or comment) in the source list file
 479  * format.  (See the man page for 'wimlib-imagex capture' for details on this
 480  * format and the meaning.)
 481  *
 482  * @line:  Line of data to be parsed.  line[len - 1] must be '\0', unless
 483  *         len == 0.  The data in @line will be modified by this function call.
 484  *
 485  * @len:   Length of the line of data.
 486  *
 487  * @source:  On success, the capture source and target described by the line is
 488  *           written into this destination.  Note that it will contain pointers
 489  *           to data in the @line array.
 490  *
 491  * Returns true if the line was valid; false otherwise.  */
 492 static bool
 493 parse_source_list_line(tchar *line, size_t len,
 494                        struct wimlib_capture_source *source)
 495 {
 496         /* SOURCE [DEST] */
 497         int ret;
 498         ret = parse_filename(&line, &len, &source->fs_source_path);
 499         if (ret != PARSE_FILENAME_SUCCESS)
 500                 return false;
 501         ret = parse_filename(&line, &len, &source->wim_target_path);
 502         if (ret == PARSE_FILENAME_NONE)
 503                 source->wim_target_path = source->fs_source_path;
 504         return ret != PARSE_FILENAME_FAILURE;
 505 }
 506
 507 /* Returns %true if the given line of length @len > 0 is a comment or empty line
 508  * in the source list file format. */
 509 static bool
 510 is_comment_line(const tchar *line, size_t len)
 511 {
 512         for (;;) {
 513                 if (*line == T('#'))
 514                         return true;
 515                 if (!istspace(*line) && *line != T('\0'))
 516                         return false;
 517                 ++line;
 518                 --len;
 519                 if (len == 0)
 520                         return true;
 521         }
 522 }
 523
 524 static ssize_t
 525 text_file_count_lines(tchar **contents_p, size_t *nchars_p)
 526 {
 527         ssize_t nlines = 0;
 528         tchar *contents = *contents_p;
 529         size_t nchars = *nchars_p;
 530         size_t i;
 531
 532         for (i = 0; i < nchars; i++)
 533                 if (contents[i] == T('\n'))
 534                         nlines++;
 535
 536         /* Handle last line not terminated by a newline */
 537         if (nchars != 0 && contents[nchars - 1] != T('\n')) {
 538                 contents = realloc(contents, (nchars + 1) * sizeof(tchar));
 539                 if (!contents) {
 540                         imagex_error(T("Out of memory!"));
 541                         return -1;
 542                 }
 543                 contents[nchars] = T('\n');
 544                 *contents_p = contents;
 545                 nchars++;
 546                 nlines++;
 547         }
 548         *nchars_p = nchars;
 549         return nlines;
 550 }
 551
 552 /* Parses a file in the source list format.  (See the man page for
 553  * 'wimlib-imagex capture' for details on this format and the meaning.)
 554  *
 555  * @source_list_contents:  Contents of the source list file.  Note that this
 556  *                         buffer will be modified to save memory allocations,
 557  *                         and cannot be freed until the returned array of
 558  *                         wimlib_capture_source's has also been freed.
 559  *
 560  * @source_list_nbytes:    Number of bytes of data in the @source_list_contents
 561  *                         buffer.
 562  *
 563  * @nsources_ret:          On success, the length of the returned array is
 564  *                         returned here.
 565  *
 566  * Returns:   An array of `struct wimlib_capture_source's that can be passed to
 567  * the wimlib_add_image_multisource() function to specify how a WIM image is to
 568  * be created.  */
 569 static struct wimlib_capture_source *
 570 parse_source_list(tchar **source_list_contents_p, size_t source_list_nchars,
 571                   size_t *nsources_ret)
 572 {
 573         ssize_t nlines;
 574         tchar *p;
 575         struct wimlib_capture_source *sources;
 576         size_t i, j;
 577
 578         nlines = text_file_count_lines(source_list_contents_p,
 579                                        &source_list_nchars);
 580         if (nlines < 0)
 581                 return NULL;
 582
 583         /* Always allocate at least 1 slot, just in case the implementation of
 584          * calloc() returns NULL if 0 bytes are requested. */
 585         sources = calloc(nlines ?: 1, sizeof(*sources));
 586         if (!sources) {
 587                 imagex_error(T("out of memory"));
 588                 return NULL;
 589         }
 590         p = *source_list_contents_p;
 591         j = 0;
 592         for (i = 0; i < nlines; i++) {
 593                 /* XXX: Could use rawmemchr() here instead, but it may not be
 594                  * available on all platforms. */
 595                 tchar *endp = tmemchr(p, T('\n'), source_list_nchars);
 596                 size_t len = endp - p + 1;
 597                 *endp = T('\0');
 598                 if (!is_comment_line(p, len)) {
 599                         if (!parse_source_list_line(p, len, &sources[j++])) {
 600                                 free(sources);
 601                                 return NULL;
 602                         }
 603                 }
 604                 p = endp + 1;
 605
 606         }
 607         *nsources_ret = j;
 608         return sources;
 609 }
 610
 611
 612 enum capture_config_section {
 613         CAPTURE_CONFIG_NO_SECTION,
 614         CAPTURE_CONFIG_EXCLUSION_SECTION,
 615         CAPTURE_CONFIG_EXCLUSION_EXCEPTION_SECTION,
 616         CAPTURE_CONFIG_IGNORE_SECTION,
 617 };
 618
 619 enum {
 620         CAPTURE_CONFIG_INVALID_SECTION,
 621         CAPTURE_CONFIG_CHANGED_SECTION,
 622         CAPTURE_CONFIG_SAME_SECTION,
 623 };
 624
 625 static int
 626 check_config_section(tchar *line, size_t len,
 627                      enum capture_config_section *cur_section)
 628 {
 629         while (istspace(*line))
 630                 line++;
 631
 632         if (*line != T('['))
 633                 return CAPTURE_CONFIG_SAME_SECTION;
 634
 635         line++;
 636         tchar *endbrace = tstrrchr(line, T(']'));
 637         if (!endbrace)
 638                 return CAPTURE_CONFIG_SAME_SECTION;
 639
 640         if (!tmemcmp(line, T("ExclusionList"), endbrace - line)) {
 641                 *cur_section = CAPTURE_CONFIG_EXCLUSION_SECTION;
 642         } else if (!tmemcmp(line, T("ExclusionException"), endbrace - line)) {
 643                 *cur_section = CAPTURE_CONFIG_EXCLUSION_EXCEPTION_SECTION;
 644         } else if (!tmemcmp(line, T("CompressionExclusionList"), endbrace - line)) {
 645                 *cur_section = CAPTURE_CONFIG_IGNORE_SECTION;
 646                 tfputs(T("WARNING: Ignoring [CompressionExclusionList] section "
 647                          "of capture config file\n"),
 648                        stderr);
 649         } else if (!tmemcmp(line, T("AlignmentList"), endbrace - line)) {
 650                 *cur_section = CAPTURE_CONFIG_IGNORE_SECTION;
 651                 tfputs(T("WARNING: Ignoring [AlignmentList] section "
 652                          "of capture config file\n"),
 653                        stderr);
 654         } else {
 655                 imagex_error(T("Invalid capture config file section \"%"TS"\""),
 656                              line - 1);
 657                 return CAPTURE_CONFIG_INVALID_SECTION;
 658         }
 659         return CAPTURE_CONFIG_CHANGED_SECTION;
 660 }
 661
 662
 663 static bool
 664 pattern_list_add_pattern(struct wimlib_pattern_list *pat_list,
 665                          tchar *pat)
 666 {
 667         if (pat_list->num_pats == pat_list->num_allocated_pats) {
 668                 tchar **pats;
 669                 size_t num_allocated_pats = pat_list->num_pats + 8;
 670
 671                 pats = realloc(pat_list->pats,
 672                                num_allocated_pats * sizeof(pat_list->pats[0]));
 673                 if (!pats) {
 674                         imagex_error(T("Out of memory!"));
 675                         return false;
 676                 }
 677                 pat_list->pats = pats;
 678                 pat_list->num_allocated_pats = num_allocated_pats;
 679         }
 680         pat_list->pats[pat_list->num_pats++] = pat;
 681         return true;
 682 }
 683
 684 static bool
 685 parse_capture_config_line(tchar *line, size_t len,
 686                           enum capture_config_section *cur_section,
 687                           struct wimlib_capture_config *config)
 688 {
 689         tchar *filename;
 690         int ret;
 691
 692         ret = check_config_section(line, len, cur_section);
 693         if (ret == CAPTURE_CONFIG_INVALID_SECTION)
 694                 return false;
 695         if (ret == CAPTURE_CONFIG_CHANGED_SECTION)
 696                 return true;
 697
 698         switch (*cur_section) {
 699         case CAPTURE_CONFIG_NO_SECTION:
 700                 imagex_error(T("Line \"%"TS"\" is not in a section "
 701                                "(such as [ExclusionList]"), line);
 702                 return false;
 703         case CAPTURE_CONFIG_EXCLUSION_SECTION:
 704                 if (parse_filename(&line, &len, &filename) != PARSE_FILENAME_SUCCESS)
 705                         return false;
 706                 return pattern_list_add_pattern(&config->exclusion_pats,
 707                                                 filename);
 708         case CAPTURE_CONFIG_EXCLUSION_EXCEPTION_SECTION:
 709                 if (parse_filename(&line, &len, &filename) != PARSE_FILENAME_SUCCESS)
 710                         return false;
 711                 return pattern_list_add_pattern(&config->exclusion_exception_pats,
 712                                                 filename);
 713         case CAPTURE_CONFIG_IGNORE_SECTION:
 714                 return true;
 715         }
 716         return false;
 717 }
 718
 719 static int
 720 parse_capture_config(tchar **contents_p, size_t nchars,
 721                      struct wimlib_capture_config *config)
 722 {
 723         ssize_t nlines;
 724         tchar *p;
 725         size_t i;
 726         enum capture_config_section cur_section;
 727
 728         memset(config, 0, sizeof(*config));
 729
 730         nlines = text_file_count_lines(contents_p, &nchars);
 731         if (nlines < 0)
 732                 return -1;
 733
 734         cur_section = CAPTURE_CONFIG_NO_SECTION;
 735         p = *contents_p;
 736         for (i = 0; i < nlines; i++) {
 737                 tchar *endp = tmemchr(p, T('\n'), nchars);
 738                 size_t len = endp - p + 1;
 739                 *endp = T('\0');
 740                 if (!is_comment_line(p, len))
 741                         if (!parse_capture_config_line(p, len, &cur_section, config))
 742                                 return -1;
 743                 p = endp + 1;
 744
 745         }
 746         return 0;
 747 }
 748
 749 /* Reads the contents of a file into memory. */
 750 static char *
 751 file_get_contents(const tchar *filename, size_t *len_ret)
 752 {
 753         struct stat stbuf;
 754         void *buf = NULL;
 755         size_t len;
 756         FILE *fp;
 757
 758         if (tstat(filename, &stbuf) != 0) {
 759                 imagex_error_with_errno(T("Failed to stat the file \"%"TS"\""), filename);
 760                 goto out;
 761         }
 762         len = stbuf.st_size;
 763
 764         fp = tfopen(filename, T("rb"));
 765         if (!fp) {
 766                 imagex_error_with_errno(T("Failed to open the file \"%"TS"\""), filename);
 767                 goto out;
 768         }
 769
 770         buf = malloc(len);
 771         if (!buf) {
 772                 imagex_error(T("Failed to allocate buffer of %zu bytes to hold "
 773                                "contents of file \"%"TS"\""), len, filename);
 774                 goto out_fclose;
 775         }
 776         if (fread(buf, 1, len, fp) != len) {
 777                 imagex_error_with_errno(T("Failed to read %zu bytes from the "
 778                                           "file \"%"TS"\""), len, filename);
 779                 goto out_free_buf;
 780         }
 781         *len_ret = len;
 782         goto out_fclose;
 783 out_free_buf:
 784         free(buf);
 785         buf = NULL;
 786 out_fclose:
 787         fclose(fp);
 788 out:
 789         return buf;
 790 }
 791
 792 /* Read standard input until EOF and return the full contents in a malloc()ed
 793  * buffer and the number of bytes of data in @len_ret.  Returns NULL on read
 794  * error. */
 795 static char *
 796 stdin_get_contents(size_t *len_ret)
 797 {
 798         /* stdin can, of course, be a pipe or other non-seekable file, so the
 799          * total length of the data cannot be pre-determined */
 800         char *buf = NULL;
 801         size_t newlen = 1024;
 802         size_t pos = 0;
 803         size_t inc = 1024;
 804         for (;;) {
 805                 char *p = realloc(buf, newlen);
 806                 size_t bytes_read, bytes_to_read;
 807                 if (!p) {
 808                         imagex_error(T("out of memory while reading stdin"));
 809                         break;
 810                 }
 811                 buf = p;
 812                 bytes_to_read = newlen - pos;
 813                 bytes_read = fread(&buf[pos], 1, bytes_to_read, stdin);
 814                 pos += bytes_read;
 815                 if (bytes_read != bytes_to_read) {
 816                         if (feof(stdin)) {
 817                                 *len_ret = pos;
 818                                 return buf;
 819                         } else {
 820                                 imagex_error_with_errno(T("error reading stdin"));
 821                                 break;
 822                         }
 823                 }
 824                 newlen += inc;
 825                 inc *= 3;
 826                 inc /= 2;
 827         }
 828         free(buf);
 829         return NULL;
 830 }
 831
 832
 833 static tchar *
 834 translate_text_to_tstr(char *text, size_t num_bytes,
 835                        size_t *num_tchars_ret)
 836 {
 837 #ifndef __WIN32__
 838         /* On non-Windows, assume an ASCII-compatible encoding, such as UTF-8.
 839          * */
 840         *num_tchars_ret = num_bytes;
 841         return text;
 842 #else /* !__WIN32__ */
 843         /* On Windows, translate the text to UTF-16LE */
 844         wchar_t *text_wstr;
 845         size_t num_wchars;
 846
 847         if (num_bytes >= 2 &&
 848             ((text[0] == 0xff && text[1] == 0xfe) ||
 849              (text[0] <= 0x7f && text[1] == 0x00)))
 850         {
 851                 /* File begins with 0xfeff, the BOM for UTF-16LE, or it begins
 852                  * with something that looks like an ASCII character encoded as
 853                  * a UTF-16LE code unit.  Assume the file is encoded as
 854                  * UTF-16LE.  This is not a 100% reliable check. */
 855                 num_wchars = num_bytes / 2;
 856                 text_wstr = (wchar_t*)text;
 857         } else {
 858                 /* File does not look like UTF-16LE.  Assume it is encoded in
 859                  * the current Windows code page.  I think these are always
 860                  * ASCII-compatible, so any so-called "plain-text" (ASCII) files
 861                  * should work as expected. */
 862                 text_wstr = win32_mbs_to_wcs(text,
 863                                              num_bytes,
 864                                              &num_wchars);
 865                 free(text);
 866         }
 867         *num_tchars_ret = num_wchars;
 868         return text_wstr;
 869 #endif /* __WIN32__ */
 870 }
 871
 872 static tchar *
 873 file_get_text_contents(const tchar *filename, size_t *num_tchars_ret)
 874 {
 875         char *contents;
 876         size_t num_bytes;
 877
 878         contents = file_get_contents(filename, &num_bytes);
 879         if (!contents)
 880                 return NULL;
 881         return translate_text_to_tstr(contents, num_bytes, num_tchars_ret);
 882 }
 883
 884 static tchar *
 885 stdin_get_text_contents(size_t *num_tchars_ret)
 886 {
 887         char *contents;
 888         size_t num_bytes;
 889
 890         contents = stdin_get_contents(&num_bytes);
 891         if (!contents)
 892                 return NULL;
 893         return translate_text_to_tstr(contents, num_bytes, num_tchars_ret);
 894 }
 895
 896 /* Return 0 if a path names a file to which the current user has write access;
 897  * -1 otherwise (and print an error message). */
 898 static int
 899 file_writable(const tchar *path)
 900 {
 901         int ret;
 902         ret = taccess(path, W_OK);
 903         if (ret != 0)
 904                 imagex_error_with_errno(T("Can't modify \"%"TS"\""), path);
 905         return ret;
 906 }
 907
 908 #define TO_PERCENT(numerator, denominator) \
 909         (((denominator) == 0) ? 0 : ((numerator) * 100 / (denominator)))
 910
 911 /* Given an enumerated value for WIM compression type, return a descriptive
 912  * string. */
 913 static const tchar *
 914 get_data_type(int ctype)
 915 {
 916         switch (ctype) {
 917         case WIMLIB_COMPRESSION_TYPE_NONE:
 918                 return T("uncompressed");
 919         case WIMLIB_COMPRESSION_TYPE_LZX:
 920                 return T("LZX-compressed");
 921         case WIMLIB_COMPRESSION_TYPE_XPRESS:
 922                 return T("XPRESS-compressed");
 923         }
 924         return NULL;
 925 }
 926
 927 /* Progress callback function passed to various wimlib functions. */
 928 static int
 929 imagex_progress_func(enum wimlib_progress_msg msg,
 930                      const union wimlib_progress_info *info)
 931 {
 932         unsigned percent_done;
 933         switch (msg) {
 934         case WIMLIB_PROGRESS_MSG_WRITE_STREAMS:
 935                 percent_done = TO_PERCENT(info->write_streams.completed_bytes,
 936                                           info->write_streams.total_bytes);
 937                 if (info->write_streams.completed_streams == 0) {
 938                         const tchar *data_type;
 939
 940                         data_type = get_data_type(info->write_streams.compression_type);
 941                         tprintf(T("Writing %"TS" data using %u thread%"TS"\n"),
 942                                 data_type, info->write_streams.num_threads,
 943                                 (info->write_streams.num_threads == 1) ? T("") : T("s"));
 944                 }
 945                 tprintf(T("\r%"PRIu64" MiB of %"PRIu64" MiB (uncompressed) "
 946                         "written (%u%% done)"),
 947                         info->write_streams.completed_bytes >> 20,
 948                         info->write_streams.total_bytes >> 20,
 949                         percent_done);
 950                 if (info->write_streams.completed_bytes >= info->write_streams.total_bytes)
 951                         tputchar(T('\n'));
 952                 break;
 953         case WIMLIB_PROGRESS_MSG_SCAN_BEGIN:
 954                 tprintf(T("Scanning \"%"TS"\""), info->scan.source);
 955                 if (*info->scan.wim_target_path) {
 956                         tprintf(T(" (loading as WIM path: \"/%"TS"\")...\n"),
 957                                info->scan.wim_target_path);
 958                 } else {
 959                         tprintf(T(" (loading as root of WIM image)...\n"));
 960                 }
 961                 break;
 962         case WIMLIB_PROGRESS_MSG_SCAN_DENTRY:
 963                 if (info->scan.excluded)
 964                         tprintf(T("Excluding \"%"TS"\" from capture\n"), info->scan.cur_path);
 965                 else
 966                         tprintf(T("Scanning \"%"TS"\"\n"), info->scan.cur_path);
 967                 break;
 968         /*case WIMLIB_PROGRESS_MSG_SCAN_END:*/
 969                 /*break;*/
 970         case WIMLIB_PROGRESS_MSG_VERIFY_INTEGRITY:
 971                 percent_done = TO_PERCENT(info->integrity.completed_bytes,
 972                                           info->integrity.total_bytes);
 973                 tprintf(T("\rVerifying integrity of \"%"TS"\": %"PRIu64" MiB "
 974                         "of %"PRIu64" MiB (%u%%) done"),
 975                         info->integrity.filename,
 976                         info->integrity.completed_bytes >> 20,
 977                         info->integrity.total_bytes >> 20,
 978                         percent_done);
 979                 if (info->integrity.completed_bytes == info->integrity.total_bytes)
 980                         tputchar(T('\n'));
 981                 break;
 982         case WIMLIB_PROGRESS_MSG_CALC_INTEGRITY:
 983                 percent_done = TO_PERCENT(info->integrity.completed_bytes,
 984                                           info->integrity.total_bytes);
 985                 tprintf(T("\rCalculating integrity table for WIM: %"PRIu64" MiB "
 986                           "of %"PRIu64" MiB (%u%%) done"),
 987                         info->integrity.completed_bytes >> 20,
 988                         info->integrity.total_bytes >> 20,
 989                         percent_done);
 990                 if (info->integrity.completed_bytes == info->integrity.total_bytes)
 991                         tputchar(T('\n'));
 992                 break;
 993         case WIMLIB_PROGRESS_MSG_EXTRACT_IMAGE_BEGIN:
 994                 tprintf(T("Applying image %d (%"TS") from \"%"TS"\" "
 995                           "to %"TS" \"%"TS"\"\n"),
 996                         info->extract.image,
 997                         info->extract.image_name,
 998                         info->extract.wimfile_name,
 999                         ((info->extract.extract_flags & WIMLIB_EXTRACT_FLAG_NTFS) ?
1000                          T("NTFS volume") : T("directory")),
1001                         info->extract.target);
1002                 break;
1003         /*case WIMLIB_PROGRESS_MSG_EXTRACT_DIR_STRUCTURE_BEGIN:*/
1004                 /*tprintf(T("Applying directory structure to %"TS"\n"),*/
1005                         /*info->extract.target);*/
1006                 /*break;*/
1007         case WIMLIB_PROGRESS_MSG_EXTRACT_STREAMS:
1008                 percent_done = TO_PERCENT(info->extract.completed_bytes,
1009                                           info->extract.total_bytes);
1010                 tprintf(T("\rExtracting files: "
1011                           "%"PRIu64" MiB of %"PRIu64" MiB (%u%%) done"),
1012                         info->extract.completed_bytes >> 20,
1013                         info->extract.total_bytes >> 20,
1014                         percent_done);
1015                 if (info->extract.completed_bytes >= info->extract.total_bytes)
1016                         tputchar(T('\n'));
1017                 break;
1018         case WIMLIB_PROGRESS_MSG_EXTRACT_DENTRY:
1019                 tprintf(T("%"TS"\n"), info->extract.cur_path);
1020                 break;
1021         case WIMLIB_PROGRESS_MSG_APPLY_TIMESTAMPS:
1022                 tprintf(T("Setting timestamps on all extracted files...\n"));
1023                 break;
1024         case WIMLIB_PROGRESS_MSG_EXTRACT_IMAGE_END:
1025                 if (info->extract.extract_flags & WIMLIB_EXTRACT_FLAG_NTFS) {
1026                         tprintf(T("Unmounting NTFS volume \"%"TS"\"...\n"),
1027                                 info->extract.target);
1028                 }
1029                 break;
1030         case WIMLIB_PROGRESS_MSG_JOIN_STREAMS:
1031                 percent_done = TO_PERCENT(info->join.completed_bytes,
1032                                           info->join.total_bytes);
1033                 tprintf(T("Writing resources from part %u of %u: "
1034                           "%"PRIu64 " MiB of %"PRIu64" MiB (%u%%) written\n"),
1035                         (info->join.completed_parts == info->join.total_parts) ?
1036                         info->join.completed_parts : info->join.completed_parts + 1,
1037                         info->join.total_parts,
1038                         info->join.completed_bytes >> 20,
1039                         info->join.total_bytes >> 20,
1040                         percent_done);
1041                 break;
1042         case WIMLIB_PROGRESS_MSG_SPLIT_BEGIN_PART:
1043                 percent_done = TO_PERCENT(info->split.completed_bytes,
1044                                           info->split.total_bytes);
1045                 tprintf(T("Writing \"%"TS"\": %"PRIu64" MiB of "
1046                           "%"PRIu64" MiB (%u%%) written\n"),
1047                         info->split.part_name,
1048                         info->split.completed_bytes >> 20,
1049                         info->split.total_bytes >> 20,
1050                         percent_done);
1051                 break;
1052         case WIMLIB_PROGRESS_MSG_SPLIT_END_PART:
1053                 if (info->split.completed_bytes == info->split.total_bytes) {
1054                         tprintf(T("Finished writing %u split WIM parts\n"),
1055                                 info->split.cur_part_number);
1056                 }
1057                 break;
1058         default:
1059                 break;
1060         }
1061         fflush(stdout);
1062         return 0;
1063 }
1064
1065 /* Open all the split WIM parts that correspond to a file glob.
1066  *
1067  * @first_part specifies the first part of the split WIM and it may be either
1068  * included or omitted from the glob. */
1069 static int
1070 open_swms_from_glob(const tchar *swm_glob,
1071                     const tchar *first_part,
1072                     int open_flags,
1073                     WIMStruct ***additional_swms_ret,
1074                     unsigned *num_additional_swms_ret)
1075 {
1076         unsigned num_additional_swms = 0;
1077         WIMStruct **additional_swms = NULL;
1078         glob_t globbuf;
1079         int ret;
1080
1081         /* Warning: glob() is replaced in Windows native builds */
1082         ret = tglob(swm_glob, GLOB_ERR | GLOB_NOSORT, NULL, &globbuf);
1083         if (ret != 0) {
1084                 if (ret == GLOB_NOMATCH) {
1085                         imagex_error(T("Found no files for glob \"%"TS"\""),
1086                                      swm_glob);
1087                 } else {
1088                         imagex_error_with_errno(T("Failed to process glob \"%"TS"\""),
1089                                                 swm_glob);
1090                 }
1091                 ret = -1;
1092                 goto out;
1093         }
1094         num_additional_swms = globbuf.gl_pathc;
1095         additional_swms = calloc(num_additional_swms, sizeof(additional_swms[0]));
1096         if (!additional_swms) {
1097                 imagex_error(T("Out of memory"));
1098                 ret = -1;
1099                 goto out_globfree;
1100         }
1101         unsigned offset = 0;
1102         for (unsigned i = 0; i < num_additional_swms; i++) {
1103                 if (tstrcmp(globbuf.gl_pathv[i], first_part) == 0) {
1104                         offset++;
1105                         continue;
1106                 }
1107                 ret = wimlib_open_wim(globbuf.gl_pathv[i],
1108                                       open_flags | WIMLIB_OPEN_FLAG_SPLIT_OK,
1109                                       &additional_swms[i - offset],
1110                                       imagex_progress_func);
1111                 if (ret != 0)
1112                         goto out_close_swms;
1113         }
1114         *additional_swms_ret = additional_swms;
1115         *num_additional_swms_ret = num_additional_swms - offset;
1116         ret = 0;
1117         goto out_globfree;
1118 out_close_swms:
1119         for (unsigned i = 0; i < num_additional_swms; i++)
1120                 wimlib_free(additional_swms[i]);
1121         free(additional_swms);
1122 out_globfree:
1123         globfree(&globbuf);
1124 out:
1125         return ret;
1126 }
1127
1128
1129 static unsigned
1130 parse_num_threads(const tchar *optarg)
1131 {
1132         tchar *tmp;
1133         unsigned long ul_nthreads = tstrtoul(optarg, &tmp, 10);
1134         if (ul_nthreads >= UINT_MAX || *tmp || tmp == optarg) {
1135                 imagex_error(T("Number of threads must be a non-negative integer!"));
1136                 return UINT_MAX;
1137         } else {
1138                 return ul_nthreads;
1139         }
1140 }
1141
1142
1143 /* Apply one image, or all images, from a WIM file into a directory, OR apply
1144  * one image from a WIM file to a NTFS volume. */
1145 static int
1146 imagex_apply(int argc, tchar **argv)
1147 {
1148         int c;
1149         int open_flags = WIMLIB_OPEN_FLAG_SPLIT_OK;
1150         int image;
1151         int num_images;
1152         WIMStruct *w;
1153         int ret;
1154         const tchar *wimfile;
1155         const tchar *target;
1156         const tchar *image_num_or_name;
1157         int extract_flags = WIMLIB_EXTRACT_FLAG_SEQUENTIAL;
1158
1159         const tchar *swm_glob = NULL;
1160         WIMStruct **additional_swms = NULL;
1161         unsigned num_additional_swms = 0;
1162
1163         for_opt(c, apply_options) {
1164                 switch (c) {
1165                 case IMAGEX_CHECK_OPTION:
1166                         open_flags |= WIMLIB_OPEN_FLAG_CHECK_INTEGRITY;
1167                         break;
1168                 case IMAGEX_HARDLINK_OPTION:
1169                         extract_flags |= WIMLIB_EXTRACT_FLAG_HARDLINK;
1170                         break;
1171                 case IMAGEX_SYMLINK_OPTION:
1172                         extract_flags |= WIMLIB_EXTRACT_FLAG_SYMLINK;
1173                         break;
1174                 case IMAGEX_VERBOSE_OPTION:
1175                         extract_flags |= WIMLIB_EXTRACT_FLAG_VERBOSE;
1176                         break;
1177                 case IMAGEX_REF_OPTION:
1178                         swm_glob = optarg;
1179                         break;
1180                 case IMAGEX_UNIX_DATA_OPTION:
1181                         extract_flags |= WIMLIB_EXTRACT_FLAG_UNIX_DATA;
1182                         break;
1183                 case IMAGEX_NO_ACLS_OPTION:
1184                         extract_flags |= WIMLIB_EXTRACT_FLAG_NO_ACLS;
1185                         break;
1186                 case IMAGEX_STRICT_ACLS_OPTION:
1187                         extract_flags |= WIMLIB_EXTRACT_FLAG_STRICT_ACLS;
1188                         break;
1189                 case IMAGEX_NORPFIX_OPTION:
1190                         extract_flags |= WIMLIB_EXTRACT_FLAG_NORPFIX;
1191                         break;
1192                 case IMAGEX_RPFIX_OPTION:
1193                         extract_flags |= WIMLIB_EXTRACT_FLAG_RPFIX;
1194                         break;
1195                 default:
1196                         usage(APPLY);
1197                         return -1;
1198                 }
1199         }
1200         argc -= optind;
1201         argv += optind;
1202         if (argc != 2 && argc != 3) {
1203                 usage(APPLY);
1204                 return -1;
1205         }
1206
1207         wimfile = argv[0];
1208         if (argc == 2) {
1209                 image_num_or_name = T("1");
1210                 target = argv[1];
1211         } else {
1212                 image_num_or_name = argv[1];
1213                 target = argv[2];
1214         }
1215
1216         ret = wimlib_open_wim(wimfile, open_flags, &w, imagex_progress_func);
1217         if (ret != 0)
1218                 return ret;
1219
1220         image = wimlib_resolve_image(w, image_num_or_name);
1221         ret = verify_image_exists(image, image_num_or_name, wimfile);
1222         if (ret != 0)
1223                 goto out;
1224
1225         num_images = wimlib_get_num_images(w);
1226         if (argc == 2 && num_images != 1) {
1227                 imagex_error(T("\"%"TS"\" contains %d images; Please select one "
1228                                "(or all)"), wimfile, num_images);
1229                 usage(APPLY);
1230                 ret = -1;
1231                 goto out;
1232         }
1233
1234         if (swm_glob) {
1235                 ret = open_swms_from_glob(swm_glob, wimfile, open_flags,
1236                                           &additional_swms,
1237                                           &num_additional_swms);
1238                 if (ret != 0)
1239                         goto out;
1240         }
1241
1242         struct stat stbuf;
1243
1244         ret = tstat(target, &stbuf);
1245         if (ret == 0) {
1246                 if (S_ISBLK(stbuf.st_mode) || S_ISREG(stbuf.st_mode))
1247                         extract_flags |= WIMLIB_EXTRACT_FLAG_NTFS;
1248         } else {
1249                 if (errno != ENOENT) {
1250                         imagex_error_with_errno(T("Failed to stat \"%"TS"\""),
1251                                                 target);
1252                         ret = -1;
1253                         goto out;
1254                 }
1255         }
1256
1257 #ifdef __WIN32__
1258         win32_acquire_restore_privileges();
1259 #endif
1260         ret = wimlib_extract_image(w, image, target, extract_flags,
1261                                    additional_swms, num_additional_swms,
1262                                    imagex_progress_func);
1263         if (ret == 0)
1264                 tprintf(T("Done applying WIM image.\n"));
1265 #ifdef __WIN32__
1266         win32_release_restore_privileges();
1267 #endif
1268 out:
1269         wimlib_free(w);
1270         if (additional_swms) {
1271                 for (unsigned i = 0; i < num_additional_swms; i++)
1272                         wimlib_free(additional_swms[i]);
1273                 free(additional_swms);
1274         }
1275         return ret;
1276 }
1277
1278 /* Create a WIM image from a directory tree, NTFS volume, or multiple files or
1279  * directory trees.  'wimlib-imagex capture': create a new WIM file containing
1280  * the desired image.  'wimlib-imagex append': add a new image to an existing
1281  * WIM file. */
1282 static int
1283 imagex_capture_or_append(int argc, tchar **argv)
1284 {
1285         int c;
1286         int open_flags = 0;
1287         int add_image_flags = WIMLIB_ADD_IMAGE_FLAG_EXCLUDE_VERBOSE;
1288         int write_flags = 0;
1289         int compression_type = WIMLIB_COMPRESSION_TYPE_XPRESS;
1290         const tchar *wimfile;
1291         const tchar *name;
1292         const tchar *desc;
1293         const tchar *flags_element = NULL;
1294         WIMStruct *w;
1295         int ret;
1296         int cur_image;
1297         int cmd = tstrcmp(argv[0], T("append")) ? CAPTURE : APPEND;
1298         unsigned num_threads = 0;
1299
1300         tchar *source;
1301         size_t source_name_len;
1302         tchar *source_copy;
1303
1304         const tchar *config_file = NULL;
1305         tchar *config_str;
1306         struct wimlib_capture_config *config = NULL;
1307
1308         bool source_list = false;
1309         size_t source_list_nchars;
1310         tchar *source_list_contents;
1311         bool capture_sources_malloced;
1312         struct wimlib_capture_source *capture_sources;
1313         size_t num_sources;
1314
1315         for_opt(c, capture_or_append_options) {
1316                 switch (c) {
1317                 case IMAGEX_BOOT_OPTION:
1318                         add_image_flags |= WIMLIB_ADD_IMAGE_FLAG_BOOT;
1319                         break;
1320                 case IMAGEX_CHECK_OPTION:
1321                         open_flags |= WIMLIB_OPEN_FLAG_CHECK_INTEGRITY;
1322                         write_flags |= WIMLIB_WRITE_FLAG_CHECK_INTEGRITY;
1323                         break;
1324                 case IMAGEX_CONFIG_OPTION:
1325                         config_file = optarg;
1326                         break;
1327                 case IMAGEX_COMPRESS_OPTION:
1328                         compression_type = get_compression_type(optarg);
1329                         if (compression_type == WIMLIB_COMPRESSION_TYPE_INVALID)
1330                                 return -1;
1331                         break;
1332                 case IMAGEX_FLAGS_OPTION:
1333                         flags_element = optarg;
1334                         break;
1335                 case IMAGEX_DEREFERENCE_OPTION:
1336                         add_image_flags |= WIMLIB_ADD_IMAGE_FLAG_DEREFERENCE;
1337                         break;
1338                 case IMAGEX_VERBOSE_OPTION:
1339                         add_image_flags |= WIMLIB_ADD_IMAGE_FLAG_VERBOSE;
1340                         break;
1341                 case IMAGEX_THREADS_OPTION:
1342                         num_threads = parse_num_threads(optarg);
1343                         if (num_threads == UINT_MAX)
1344                                 return -1;
1345                         break;
1346                 case IMAGEX_REBULID_OPTION:
1347                         write_flags |= WIMLIB_WRITE_FLAG_REBUILD;
1348                         break;
1349                 case IMAGEX_UNIX_DATA_OPTION:
1350                         add_image_flags |= WIMLIB_ADD_IMAGE_FLAG_UNIX_DATA;
1351                         break;
1352                 case IMAGEX_SOURCE_LIST_OPTION:
1353                         source_list = true;
1354                         break;
1355                 case IMAGEX_NO_ACLS_OPTION:
1356                         add_image_flags |= WIMLIB_ADD_IMAGE_FLAG_NO_ACLS;
1357                         break;
1358                 case IMAGEX_STRICT_ACLS_OPTION:
1359                         add_image_flags |= WIMLIB_ADD_IMAGE_FLAG_STRICT_ACLS;
1360                         break;
1361                 case IMAGEX_RPFIX_OPTION:
1362                         add_image_flags |= WIMLIB_ADD_IMAGE_FLAG_RPFIX;
1363                         break;
1364                 case IMAGEX_NORPFIX_OPTION:
1365                         add_image_flags |= WIMLIB_ADD_IMAGE_FLAG_NORPFIX;
1366                         break;
1367                 default:
1368                         usage(cmd);
1369                         return -1;
1370                 }
1371         }
1372         argc -= optind;
1373         argv += optind;
1374
1375         if (argc < 2 || argc > 4) {
1376                 usage(cmd);
1377                 return -1;
1378         }
1379
1380         source = argv[0];
1381         wimfile = argv[1];
1382
1383         if (argc >= 3) {
1384                 name = argv[2];
1385         } else {
1386                 /* Set default name to SOURCE argument, omitting any directory
1387                  * prefixes and trailing slashes.  This requires making a copy
1388                  * of @source. */
1389                 source_name_len = tstrlen(source);
1390                 source_copy = alloca((source_name_len + 1) * sizeof(tchar));
1391                 name = tbasename(tstrcpy(source_copy, source));
1392         }
1393         /* Image description defaults to NULL if not given. */
1394         desc = (argc >= 4) ? argv[3] : NULL;
1395
1396         if (source_list) {
1397                 /* Set up capture sources in source list mode */
1398                 if (source[0] == T('-') && source[1] == T('\0')) {
1399                         source_list_contents = stdin_get_text_contents(&source_list_nchars);
1400                 } else {
1401                         source_list_contents = file_get_text_contents(source,
1402                                                                       &source_list_nchars);
1403                 }
1404                 if (!source_list_contents)
1405                         return -1;
1406
1407                 capture_sources = parse_source_list(&source_list_contents,
1408                                                     source_list_nchars,
1409                                                     &num_sources);
1410                 if (!capture_sources) {
1411                         ret = -1;
1412                         goto out_free_source_list_contents;
1413                 }
1414                 capture_sources_malloced = true;
1415         } else {
1416                 /* Set up capture source in non-source-list mode (could be
1417                  * either "normal" mode or "NTFS mode"--- see the man page). */
1418                 capture_sources = alloca(sizeof(struct wimlib_capture_source));
1419                 capture_sources[0].fs_source_path = source;
1420                 capture_sources[0].wim_target_path = NULL;
1421                 capture_sources[0].reserved = 0;
1422                 num_sources = 1;
1423                 capture_sources_malloced = false;
1424                 source_list_contents = NULL;
1425         }
1426
1427         if (config_file) {
1428                 size_t config_len;
1429
1430                 config_str = file_get_text_contents(config_file, &config_len);
1431                 if (!config_str) {
1432                         ret = -1;
1433                         goto out_free_capture_sources;
1434                 }
1435
1436                 config = alloca(sizeof(*config));
1437                 ret = parse_capture_config(&config_str, config_len, config);
1438                 if (ret)
1439                         goto out_free_config;
1440         } else {
1441                 config = &default_capture_config;
1442         }
1443
1444         if (cmd == APPEND)
1445                 ret = wimlib_open_wim(wimfile, open_flags, &w,
1446                                       imagex_progress_func);
1447         else
1448                 ret = wimlib_create_new_wim(compression_type, &w);
1449         if (ret)
1450                 goto out_free_config;
1451
1452         if (!source_list) {
1453                 struct stat stbuf;
1454                 ret = tstat(source, &stbuf);
1455                 if (ret == 0) {
1456                         if (S_ISBLK(stbuf.st_mode) || S_ISREG(stbuf.st_mode)) {
1457                                 tprintf(T("Capturing WIM image from NTFS "
1458                                           "filesystem on \"%"TS"\"\n"), source);
1459                                 add_image_flags |= WIMLIB_ADD_IMAGE_FLAG_NTFS;
1460                         }
1461                 } else {
1462                         if (errno != ENOENT) {
1463                                 imagex_error_with_errno(T("Failed to stat "
1464                                                           "\"%"TS"\""), source);
1465                                 ret = -1;
1466                                 goto out_wimlib_free;
1467                         }
1468                 }
1469         }
1470 #ifdef __WIN32__
1471         win32_acquire_capture_privileges();
1472 #endif
1473
1474         ret = wimlib_add_image_multisource(w,
1475                                            capture_sources,
1476                                            num_sources,
1477                                            name,
1478                                            config,
1479                                            add_image_flags,
1480                                            imagex_progress_func);
1481         if (ret != 0)
1482                 goto out_release_privs;
1483         cur_image = wimlib_get_num_images(w);
1484         if (desc) {
1485                 ret = wimlib_set_image_descripton(w, cur_image, desc);
1486                 if (ret != 0)
1487                         goto out_release_privs;
1488         }
1489         if (flags_element) {
1490                 ret = wimlib_set_image_flags(w, cur_image, flags_element);
1491                 if (ret != 0)
1492                         goto out_release_privs;
1493         }
1494         if (cmd == APPEND) {
1495                 ret = wimlib_overwrite(w, write_flags, num_threads,
1496                                        imagex_progress_func);
1497         } else {
1498                 ret = wimlib_write(w, wimfile, WIMLIB_ALL_IMAGES, write_flags,
1499                                    num_threads, imagex_progress_func);
1500         }
1501         if (ret == WIMLIB_ERR_REOPEN)
1502                 ret = 0;
1503         if (ret != 0)
1504                 imagex_error(T("Failed to write the WIM file \"%"TS"\""),
1505                              wimfile);
1506 out_release_privs:
1507 #ifdef __WIN32__
1508         win32_release_capture_privileges();
1509 #endif
1510 out_wimlib_free:
1511         wimlib_free(w);
1512 out_free_config:
1513         if (config != NULL && config != &default_capture_config) {
1514                 free(config->exclusion_pats.pats);
1515                 free(config->exclusion_exception_pats.pats);
1516                 free(config_str);
1517         }
1518 out_free_capture_sources:
1519         if (capture_sources_malloced)
1520                 free(capture_sources);
1521 out_free_source_list_contents:
1522         free(source_list_contents);
1523         return ret;
1524 }
1525
1526 /* Remove image(s) from a WIM. */
1527 static int
1528 imagex_delete(int argc, tchar **argv)
1529 {
1530         int c;
1531         int open_flags = 0;
1532         int write_flags = 0;
1533         const tchar *wimfile;
1534         const tchar *image_num_or_name;
1535         WIMStruct *w;
1536         int image;
1537         int ret;
1538
1539         for_opt(c, delete_options) {
1540                 switch (c) {
1541                 case IMAGEX_CHECK_OPTION:
1542                         open_flags |= WIMLIB_OPEN_FLAG_CHECK_INTEGRITY;
1543                         write_flags |= WIMLIB_WRITE_FLAG_CHECK_INTEGRITY;
1544                         break;
1545                 case IMAGEX_SOFT_OPTION:
1546                         write_flags |= WIMLIB_WRITE_FLAG_SOFT_DELETE;
1547                         break;
1548                 default:
1549                         usage(DELETE);
1550                         return -1;
1551                 }
1552         }
1553         argc -= optind;
1554         argv += optind;
1555
1556         if (argc != 2) {
1557                 if (argc < 1)
1558                         imagex_error(T("Must specify a WIM file"));
1559                 if (argc < 2)
1560                         imagex_error(T("Must specify an image"));
1561                 usage(DELETE);
1562                 return -1;
1563         }
1564         wimfile = argv[0];
1565         image_num_or_name = argv[1];
1566
1567         ret = file_writable(wimfile);
1568         if (ret != 0)
1569                 return ret;
1570
1571         ret = wimlib_open_wim(wimfile, open_flags, &w,
1572                               imagex_progress_func);
1573         if (ret != 0)
1574                 return ret;
1575
1576         image = wimlib_resolve_image(w, image_num_or_name);
1577
1578         ret = verify_image_exists(image, image_num_or_name, wimfile);
1579         if (ret != 0)
1580                 goto out;
1581
1582         ret = wimlib_delete_image(w, image);
1583         if (ret != 0) {
1584                 imagex_error(T("Failed to delete image from \"%"TS"\""), wimfile);
1585                 goto out;
1586         }
1587
1588         ret = wimlib_overwrite(w, write_flags, 0, imagex_progress_func);
1589         if (ret == WIMLIB_ERR_REOPEN)
1590                 ret = 0;
1591         if (ret != 0) {
1592                 imagex_error(T("Failed to write the file \"%"TS"\" with image "
1593                                "deleted"), wimfile);
1594         }
1595 out:
1596         wimlib_free(w);
1597         return ret;
1598 }
1599
1600 /* Print the files contained in an image(s) in a WIM file. */
1601 static int
1602 imagex_dir(int argc, tchar **argv)
1603 {
1604         const tchar *wimfile;
1605         WIMStruct *w;
1606         int image;
1607         int ret;
1608         int num_images;
1609
1610         if (argc < 2) {
1611                 imagex_error(T("Must specify a WIM file"));
1612                 usage(DIR);
1613                 return -1;
1614         }
1615         if (argc > 3) {
1616                 imagex_error(T("Too many arguments"));
1617                 usage(DIR);
1618                 return -1;
1619         }
1620
1621         wimfile = argv[1];
1622         ret = wimlib_open_wim(wimfile, WIMLIB_OPEN_FLAG_SPLIT_OK, &w,
1623                               imagex_progress_func);
1624         if (ret != 0)
1625                 return ret;
1626
1627         if (argc == 3) {
1628                 image = wimlib_resolve_image(w, argv[2]);
1629                 ret = verify_image_exists(image, argv[2], wimfile);
1630                 if (ret != 0)
1631                         goto out;
1632         } else {
1633                 /* Image was not specified.  If the WIM only contains one image,
1634                  * choose that one; otherwise, print an error. */
1635                 num_images = wimlib_get_num_images(w);
1636                 if (num_images != 1) {
1637                         imagex_error(T("The file \"%"TS"\" contains %d images; Please "
1638                                        "select one."), wimfile, num_images);
1639                         usage(DIR);
1640                         ret = -1;
1641                         goto out;
1642                 }
1643                 image = 1;
1644         }
1645
1646         ret = wimlib_print_files(w, image);
1647 out:
1648         wimlib_free(w);
1649         return ret;
1650 }
1651
1652 /* Exports one, or all, images from a WIM file to a new WIM file or an existing
1653  * WIM file. */
1654 static int
1655 imagex_export(int argc, tchar **argv)
1656 {
1657         int c;
1658         int open_flags = 0;
1659         int export_flags = 0;
1660         int write_flags = 0;
1661         int compression_type = WIMLIB_COMPRESSION_TYPE_NONE;
1662         bool compression_type_specified = false;
1663         const tchar *src_wimfile;
1664         const tchar *src_image_num_or_name;
1665         const tchar *dest_wimfile;
1666         const tchar *dest_name;
1667         const tchar *dest_desc;
1668         WIMStruct *src_w = NULL;
1669         WIMStruct *dest_w = NULL;
1670         int ret;
1671         int image;
1672         struct stat stbuf;
1673         bool wim_is_new;
1674         const tchar *swm_glob = NULL;
1675         WIMStruct **additional_swms = NULL;
1676         unsigned num_additional_swms = 0;
1677         unsigned num_threads = 0;
1678
1679         for_opt(c, export_options) {
1680                 switch (c) {
1681                 case IMAGEX_BOOT_OPTION:
1682                         export_flags |= WIMLIB_EXPORT_FLAG_BOOT;
1683                         break;
1684                 case IMAGEX_CHECK_OPTION:
1685                         open_flags |= WIMLIB_OPEN_FLAG_CHECK_INTEGRITY;
1686                         write_flags |= WIMLIB_WRITE_FLAG_CHECK_INTEGRITY;
1687                         break;
1688                 case IMAGEX_COMPRESS_OPTION:
1689                         compression_type = get_compression_type(optarg);
1690                         if (compression_type == WIMLIB_COMPRESSION_TYPE_INVALID)
1691                                 return -1;
1692                         compression_type_specified = true;
1693                         break;
1694                 case IMAGEX_REF_OPTION:
1695                         swm_glob = optarg;
1696                         break;
1697                 case IMAGEX_THREADS_OPTION:
1698                         num_threads = parse_num_threads(optarg);
1699                         if (num_threads == UINT_MAX)
1700                                 return -1;
1701                         break;
1702                 case IMAGEX_REBULID_OPTION:
1703                         write_flags |= WIMLIB_WRITE_FLAG_REBUILD;
1704                         break;
1705                 default:
1706                         usage(EXPORT);
1707                         return -1;
1708                 }
1709         }
1710         argc -= optind;
1711         argv += optind;
1712         if (argc < 3 || argc > 5) {
1713                 usage(EXPORT);
1714                 return -1;
1715         }
1716         src_wimfile           = argv[0];
1717         src_image_num_or_name = argv[1];
1718         dest_wimfile          = argv[2];
1719         dest_name             = (argc >= 4) ? argv[3] : NULL;
1720         dest_desc             = (argc >= 5) ? argv[4] : NULL;
1721         ret = wimlib_open_wim(src_wimfile,
1722                               open_flags | WIMLIB_OPEN_FLAG_SPLIT_OK, &src_w,
1723                               imagex_progress_func);
1724         if (ret != 0)
1725                 return ret;
1726
1727         /* Determine if the destination is an existing file or not.
1728          * If so, we try to append the exported image(s) to it; otherwise, we
1729          * create a new WIM containing the exported image(s). */
1730         if (tstat(dest_wimfile, &stbuf) == 0) {
1731                 int dest_ctype;
1732
1733                 wim_is_new = false;
1734                 /* Destination file exists. */
1735
1736                 if (!S_ISREG(stbuf.st_mode)) {
1737                         imagex_error(T("\"%"TS"\" is not a regular file"),
1738                                      dest_wimfile);
1739                         ret = -1;
1740                         goto out;
1741                 }
1742                 ret = wimlib_open_wim(dest_wimfile, open_flags, &dest_w,
1743                                       imagex_progress_func);
1744                 if (ret != 0)
1745                         goto out;
1746
1747                 ret = file_writable(dest_wimfile);
1748                 if (ret != 0)
1749                         goto out;
1750
1751                 dest_ctype = wimlib_get_compression_type(dest_w);
1752                 if (compression_type_specified
1753                     && compression_type != dest_ctype)
1754                 {
1755                         imagex_error(T("Cannot specify a compression type that is "
1756                                        "not the same as that used in the "
1757                                        "destination WIM"));
1758                         ret = -1;
1759                         goto out;
1760                 }
1761         } else {
1762                 wim_is_new = true;
1763                 /* dest_wimfile is not an existing file, so create a new WIM. */
1764                 if (!compression_type_specified)
1765                         compression_type = wimlib_get_compression_type(src_w);
1766                 if (errno == ENOENT) {
1767                         ret = wimlib_create_new_wim(compression_type, &dest_w);
1768                         if (ret != 0)
1769                                 goto out;
1770                 } else {
1771                         imagex_error_with_errno(T("Cannot stat file \"%"TS"\""),
1772                                                 dest_wimfile);
1773                         ret = -1;
1774                         goto out;
1775                 }
1776         }
1777
1778         image = wimlib_resolve_image(src_w, src_image_num_or_name);
1779         ret = verify_image_exists(image, src_image_num_or_name, src_wimfile);
1780         if (ret != 0)
1781                 goto out;
1782
1783         if (swm_glob) {
1784                 ret = open_swms_from_glob(swm_glob, src_wimfile, open_flags,
1785                                           &additional_swms,
1786                                           &num_additional_swms);
1787                 if (ret != 0)
1788                         goto out;
1789         }
1790
1791         ret = wimlib_export_image(src_w, image, dest_w, dest_name, dest_desc,
1792                                   export_flags, additional_swms,
1793                                   num_additional_swms, imagex_progress_func);
1794         if (ret != 0)
1795                 goto out;
1796
1797
1798         if (wim_is_new)
1799                 ret = wimlib_write(dest_w, dest_wimfile, WIMLIB_ALL_IMAGES,
1800                                    write_flags, num_threads,
1801                                    imagex_progress_func);
1802         else
1803                 ret = wimlib_overwrite(dest_w, write_flags, num_threads,
1804                                        imagex_progress_func);
1805 out:
1806         if (ret == WIMLIB_ERR_REOPEN)
1807                 ret = 0;
1808         wimlib_free(src_w);
1809         wimlib_free(dest_w);
1810         if (additional_swms) {
1811                 for (unsigned i = 0; i < num_additional_swms; i++)
1812                         wimlib_free(additional_swms[i]);
1813                 free(additional_swms);
1814         }
1815         return ret;
1816 }
1817
1818 /* Prints information about a WIM file; also can mark an image as bootable,
1819  * change the name of an image, or change the description of an image. */
1820 static int
1821 imagex_info(int argc, tchar **argv)
1822 {
1823         int c;
1824         bool boot         = false;
1825         bool check        = false;
1826         bool header       = false;
1827         bool lookup_table = false;
1828         bool xml          = false;
1829         bool metadata     = false;
1830         bool short_header = true;
1831         const tchar *xml_out_file = NULL;
1832         const tchar *wimfile;
1833         const tchar *image_num_or_name = T("all");
1834         const tchar *new_name = NULL;
1835         const tchar *new_desc = NULL;
1836         WIMStruct *w;
1837         FILE *fp;
1838         int image;
1839         int ret;
1840         int open_flags = WIMLIB_OPEN_FLAG_SPLIT_OK;
1841         int part_number;
1842         int total_parts;
1843         int num_images;
1844
1845         for_opt(c, info_options) {
1846                 switch (c) {
1847                 case IMAGEX_BOOT_OPTION:
1848                         boot = true;
1849                         break;
1850                 case IMAGEX_CHECK_OPTION:
1851                         check = true;
1852                         break;
1853                 case IMAGEX_HEADER_OPTION:
1854                         header = true;
1855                         short_header = false;
1856                         break;
1857                 case IMAGEX_LOOKUP_TABLE_OPTION:
1858                         lookup_table = true;
1859                         short_header = false;
1860                         break;
1861                 case IMAGEX_XML_OPTION:
1862                         xml = true;
1863                         short_header = false;
1864                         break;
1865                 case IMAGEX_EXTRACT_XML_OPTION:
1866                         xml_out_file = optarg;
1867                         short_header = false;
1868                         break;
1869                 case IMAGEX_METADATA_OPTION:
1870                         metadata = true;
1871                         short_header = false;
1872                         break;
1873                 default:
1874                         usage(INFO);
1875                         return -1;
1876                 }
1877         }
1878
1879         argc -= optind;
1880         argv += optind;
1881         if (argc == 0 || argc > 4) {
1882                 usage(INFO);
1883                 return -1;
1884         }
1885         wimfile = argv[0];
1886         if (argc > 1) {
1887                 image_num_or_name = argv[1];
1888                 if (argc > 2) {
1889                         new_name = argv[2];
1890                         if (argc > 3) {
1891                                 new_desc = argv[3];
1892                         }
1893                 }
1894         }
1895
1896         if (check)
1897                 open_flags |= WIMLIB_OPEN_FLAG_CHECK_INTEGRITY;
1898
1899         ret = wimlib_open_wim(wimfile, open_flags, &w,
1900                               imagex_progress_func);
1901         if (ret != 0)
1902                 return ret;
1903
1904         part_number = wimlib_get_part_number(w, &total_parts);
1905
1906         image = wimlib_resolve_image(w, image_num_or_name);
1907         if (image == WIMLIB_NO_IMAGE && tstrcmp(image_num_or_name, T("0"))) {
1908                 imagex_error(T("The image \"%"TS"\" does not exist"),
1909                              image_num_or_name);
1910                 if (boot) {
1911                         imagex_error(T("If you would like to set the boot "
1912                                        "index to 0, specify image \"0\" with "
1913                                        "the --boot flag."));
1914                 }
1915                 ret = WIMLIB_ERR_INVALID_IMAGE;
1916                 goto out;
1917         }
1918
1919         num_images = wimlib_get_num_images(w);
1920
1921         if (num_images == 0) {
1922                 if (boot) {
1923                         imagex_error(T("--boot is meaningless on a WIM with no "
1924                                        "images"));
1925                         ret = WIMLIB_ERR_INVALID_IMAGE;
1926                         goto out;
1927                 }
1928         }
1929
1930         if (image == WIMLIB_ALL_IMAGES && num_images > 1) {
1931                 if (boot) {
1932                         imagex_error(T("Cannot specify the --boot flag "
1933                                        "without specifying a specific "
1934                                        "image in a multi-image WIM"));
1935                         ret = WIMLIB_ERR_INVALID_IMAGE;
1936                         goto out;
1937                 }
1938                 if (new_name) {
1939                         imagex_error(T("Cannot specify the NEW_NAME "
1940                                        "without specifying a specific "
1941                                        "image in a multi-image WIM"));
1942                         ret = WIMLIB_ERR_INVALID_IMAGE;
1943                         goto out;
1944                 }
1945         }
1946
1947         /* Operations that print information are separated from operations that
1948          * recreate the WIM file. */
1949         if (!new_name && !boot) {
1950
1951                 /* Read-only operations */
1952
1953                 if (image == WIMLIB_NO_IMAGE) {
1954                         imagex_error(T("\"%"TS"\" is not a valid image"),
1955                                      image_num_or_name);
1956                         ret = WIMLIB_ERR_INVALID_IMAGE;
1957                         goto out;
1958                 }
1959
1960                 if (image == WIMLIB_ALL_IMAGES && short_header)
1961                         wimlib_print_wim_information(w);
1962
1963                 if (header)
1964                         wimlib_print_header(w);
1965
1966                 if (lookup_table) {
1967                         if (total_parts != 1) {
1968                                 tprintf(T("Warning: Only showing the lookup table "
1969                                           "for part %d of a %d-part WIM.\n"),
1970                                         part_number, total_parts);
1971                         }
1972                         wimlib_print_lookup_table(w);
1973                 }
1974
1975                 if (xml) {
1976                         ret = wimlib_extract_xml_data(w, stdout);
1977                         if (ret != 0)
1978                                 goto out;
1979                 }
1980
1981                 if (xml_out_file) {
1982                         fp = tfopen(xml_out_file, T("wb"));
1983                         if (!fp) {
1984                                 imagex_error_with_errno(T("Failed to open the "
1985                                                           "file \"%"TS"\" for "
1986                                                           "writing "),
1987                                                         xml_out_file);
1988                                 ret = -1;
1989                                 goto out;
1990                         }
1991                         ret = wimlib_extract_xml_data(w, fp);
1992                         if (fclose(fp) != 0) {
1993                                 imagex_error(T("Failed to close the file "
1994                                                "\"%"TS"\""),
1995                                              xml_out_file);
1996                                 ret = -1;
1997                         }
1998
1999                         if (ret != 0)
2000                                 goto out;
2001                 }
2002
2003                 if (short_header)
2004                         wimlib_print_available_images(w, image);
2005
2006                 if (metadata) {
2007                         ret = wimlib_print_metadata(w, image);
2008                         if (ret != 0)
2009                                 goto out;
2010                 }
2011         } else {
2012
2013                 /* Modification operations */
2014                 if (total_parts != 1) {
2015                         imagex_error(T("Modifying a split WIM is not supported."));
2016                         ret = -1;
2017                         goto out;
2018                 }
2019                 if (image == WIMLIB_ALL_IMAGES)
2020                         image = 1;
2021
2022                 if (image == WIMLIB_NO_IMAGE && new_name) {
2023                         imagex_error(T("Cannot specify new_name (\"%"TS"\") "
2024                                        "when using image 0"), new_name);
2025                         ret = -1;
2026                         goto out;
2027                 }
2028
2029                 if (boot) {
2030                         if (image == wimlib_get_boot_idx(w)) {
2031                                 tprintf(T("Image %d is already marked as "
2032                                           "bootable.\n"), image);
2033                                 boot = false;
2034                         } else {
2035                                 tprintf(T("Marking image %d as bootable.\n"),
2036                                         image);
2037                                 wimlib_set_boot_idx(w, image);
2038                         }
2039                 }
2040                 if (new_name) {
2041                         if (!tstrcmp(wimlib_get_image_name(w, image), new_name))
2042                         {
2043                                 tprintf(T("Image %d is already named \"%"TS"\".\n"),
2044                                         image, new_name);
2045                                 new_name = NULL;
2046                         } else {
2047                                 tprintf(T("Changing the name of image %d to "
2048                                           "\"%"TS"\".\n"), image, new_name);
2049                                 ret = wimlib_set_image_name(w, image, new_name);
2050                                 if (ret != 0)
2051                                         goto out;
2052                         }
2053                 }
2054                 if (new_desc) {
2055                         const tchar *old_desc;
2056                         old_desc = wimlib_get_image_description(w, image);
2057                         if (old_desc && !tstrcmp(old_desc, new_desc)) {
2058                                 tprintf(T("The description of image %d is already "
2059                                           "\"%"TS"\".\n"), image, new_desc);
2060                                 new_desc = NULL;
2061                         } else {
2062                                 tprintf(T("Changing the description of image %d "
2063                                           "to \"%"TS"\".\n"), image, new_desc);
2064                                 ret = wimlib_set_image_descripton(w, image,
2065                                                                   new_desc);
2066                                 if (ret != 0)
2067                                         goto out;
2068                         }
2069                 }
2070
2071                 /* Only call wimlib_overwrite() if something actually needs to
2072                  * be changed. */
2073                 if (boot || new_name || new_desc ||
2074                     (check && !wimlib_has_integrity_table(w)))
2075                 {
2076                         int write_flags;
2077
2078                         ret = file_writable(wimfile);
2079                         if (ret != 0)
2080                                 return ret;
2081
2082                         if (check)
2083                                 write_flags = WIMLIB_WRITE_FLAG_CHECK_INTEGRITY;
2084                         else
2085                                 write_flags = 0;
2086
2087                         ret = wimlib_overwrite(w, write_flags, 1,
2088                                                imagex_progress_func);
2089                         if (ret == WIMLIB_ERR_REOPEN)
2090                                 ret = 0;
2091                 } else {
2092                         tprintf(T("The file \"%"TS"\" was not modified because nothing "
2093                                   "needed to be done.\n"), wimfile);
2094                         ret = 0;
2095                 }
2096         }
2097 out:
2098         wimlib_free(w);
2099         return ret;
2100 }
2101
2102 /* Join split WIMs into one part WIM */
2103 static int
2104 imagex_join(int argc, tchar **argv)
2105 {
2106         int c;
2107         int swm_open_flags = WIMLIB_OPEN_FLAG_SPLIT_OK;
2108         int wim_write_flags = 0;
2109         const tchar *output_path;
2110
2111         for_opt(c, join_options) {
2112                 switch (c) {
2113                 case IMAGEX_CHECK_OPTION:
2114                         swm_open_flags |= WIMLIB_OPEN_FLAG_CHECK_INTEGRITY;
2115                         wim_write_flags |= WIMLIB_WRITE_FLAG_CHECK_INTEGRITY;
2116                         break;
2117                 default:
2118                         goto err;
2119                 }
2120         }
2121         argc -= optind;
2122         argv += optind;
2123
2124         if (argc < 2) {
2125                 imagex_error(T("Must specify one or more split WIM (.swm) "
2126                                "parts to join"));
2127                 goto err;
2128         }
2129         output_path = argv[0];
2130         return wimlib_join((const tchar * const *)++argv,
2131                            --argc,
2132                            output_path,
2133                            swm_open_flags,
2134                            wim_write_flags,
2135                            imagex_progress_func);
2136 err:
2137         usage(JOIN);
2138         return -1;
2139 }
2140
2141 /* Mounts an image using a FUSE mount. */
2142 static int
2143 imagex_mount_rw_or_ro(int argc, tchar **argv)
2144 {
2145         int c;
2146         int mount_flags = 0;
2147         int open_flags = WIMLIB_OPEN_FLAG_SPLIT_OK;
2148         const tchar *wimfile;
2149         const tchar *dir;
2150         WIMStruct *w;
2151         int image;
2152         int num_images;
2153         int ret;
2154         const tchar *swm_glob = NULL;
2155         WIMStruct **additional_swms = NULL;
2156         unsigned num_additional_swms = 0;
2157         const tchar *staging_dir = NULL;
2158
2159         if (!tstrcmp(argv[0], T("mountrw")))
2160                 mount_flags |= WIMLIB_MOUNT_FLAG_READWRITE;
2161
2162         for_opt(c, mount_options) {
2163                 switch (c) {
2164                 case IMAGEX_ALLOW_OTHER_OPTION:
2165                         mount_flags |= WIMLIB_MOUNT_FLAG_ALLOW_OTHER;
2166                         break;
2167                 case IMAGEX_CHECK_OPTION:
2168                         open_flags |= WIMLIB_OPEN_FLAG_CHECK_INTEGRITY;
2169                         break;
2170                 case IMAGEX_DEBUG_OPTION:
2171                         mount_flags |= WIMLIB_MOUNT_FLAG_DEBUG;
2172                         break;
2173                 case IMAGEX_STREAMS_INTERFACE_OPTION:
2174                         if (!tstrcasecmp(optarg, T("none")))
2175                                 mount_flags |= WIMLIB_MOUNT_FLAG_STREAM_INTERFACE_NONE;
2176                         else if (!tstrcasecmp(optarg, T("xattr")))
2177                                 mount_flags |= WIMLIB_MOUNT_FLAG_STREAM_INTERFACE_XATTR;
2178                         else if (!tstrcasecmp(optarg, T("windows")))
2179                                 mount_flags |= WIMLIB_MOUNT_FLAG_STREAM_INTERFACE_WINDOWS;
2180                         else {
2181                                 imagex_error(T("Unknown stream interface \"%"TS"\""),
2182                                              optarg);
2183                                 goto mount_usage;
2184                         }
2185                         break;
2186                 case IMAGEX_REF_OPTION:
2187                         swm_glob = optarg;
2188                         break;
2189                 case IMAGEX_STAGING_DIR_OPTION:
2190                         staging_dir = optarg;
2191                         break;
2192                 case IMAGEX_UNIX_DATA_OPTION:
2193                         mount_flags |= WIMLIB_MOUNT_FLAG_UNIX_DATA;
2194                         break;
2195                 default:
2196                         goto mount_usage;
2197                 }
2198         }
2199         argc -= optind;
2200         argv += optind;
2201         if (argc != 2 && argc != 3)
2202                 goto mount_usage;
2203
2204         wimfile = argv[0];
2205
2206         ret = wimlib_open_wim(wimfile, open_flags, &w,
2207                               imagex_progress_func);
2208         if (ret != 0)
2209                 return ret;
2210
2211         if (swm_glob) {
2212                 ret = open_swms_from_glob(swm_glob, wimfile, open_flags,
2213                                           &additional_swms,
2214                                           &num_additional_swms);
2215                 if (ret != 0)
2216                         goto out;
2217         }
2218
2219         if (argc == 2) {
2220                 image = 1;
2221                 num_images = wimlib_get_num_images(w);
2222                 if (num_images != 1) {
2223                         imagex_error(T("The file \"%"TS"\" contains %d images; Please "
2224                                        "select one."), wimfile, num_images);
2225                         usage((mount_flags & WIMLIB_MOUNT_FLAG_READWRITE)
2226                                         ? MOUNTRW : MOUNT);
2227                         ret = -1;
2228                         goto out;
2229                 }
2230                 dir = argv[1];
2231         } else {
2232                 image = wimlib_resolve_image(w, argv[1]);
2233                 dir = argv[2];
2234                 ret = verify_image_exists_and_is_single(image, argv[1], wimfile);
2235                 if (ret != 0)
2236                         goto out;
2237         }
2238
2239         if (mount_flags & WIMLIB_MOUNT_FLAG_READWRITE) {
2240                 ret = file_writable(wimfile);
2241                 if (ret != 0)
2242                         goto out;
2243         }
2244
2245         ret = wimlib_mount_image(w, image, dir, mount_flags, additional_swms,
2246                                  num_additional_swms, staging_dir);
2247         if (ret != 0) {
2248                 imagex_error(T("Failed to mount image %d from \"%"TS"\" "
2249                                "on \"%"TS"\""),
2250                              image, wimfile, dir);
2251
2252         }
2253 out:
2254         wimlib_free(w);
2255         if (additional_swms) {
2256                 for (unsigned i = 0; i < num_additional_swms; i++)
2257                         wimlib_free(additional_swms[i]);
2258                 free(additional_swms);
2259         }
2260         return ret;
2261 mount_usage:
2262         usage((mount_flags & WIMLIB_MOUNT_FLAG_READWRITE)
2263                         ? MOUNTRW : MOUNT);
2264         return -1;
2265 }
2266
2267 /* Rebuild a WIM file */
2268 static int
2269 imagex_optimize(int argc, tchar **argv)
2270 {
2271         int c;
2272         int open_flags = 0;
2273         int write_flags = WIMLIB_WRITE_FLAG_REBUILD;
2274         int ret;
2275         WIMStruct *w;
2276         const tchar *wimfile;
2277         off_t old_size;
2278         off_t new_size;
2279         unsigned num_threads = 0;
2280
2281         for_opt(c, optimize_options) {
2282                 switch (c) {
2283                 case IMAGEX_CHECK_OPTION:
2284                         open_flags |= WIMLIB_OPEN_FLAG_CHECK_INTEGRITY;
2285                         write_flags |= WIMLIB_WRITE_FLAG_CHECK_INTEGRITY;
2286                         break;
2287                 case IMAGEX_RECOMPRESS_OPTION:
2288                         write_flags |= WIMLIB_WRITE_FLAG_RECOMPRESS;
2289                         break;
2290                 case IMAGEX_THREADS_OPTION:
2291                         num_threads = parse_num_threads(optarg);
2292                         if (num_threads == UINT_MAX)
2293                                 return -1;
2294                         break;
2295                 default:
2296                         usage(OPTIMIZE);
2297                         return -1;
2298                 }
2299         }
2300         argc -= optind;
2301         argv += optind;
2302
2303         if (argc != 1) {
2304                 usage(OPTIMIZE);
2305                 return -1;
2306         }
2307
2308         wimfile = argv[0];
2309
2310         ret = wimlib_open_wim(wimfile, open_flags, &w,
2311                               imagex_progress_func);
2312         if (ret != 0)
2313                 return ret;
2314
2315         old_size = file_get_size(argv[0]);
2316         tprintf(T("\"%"TS"\" original size: "), wimfile);
2317         if (old_size == -1)
2318                 tfputs(T("Unknown\n"), stdout);
2319         else
2320                 tprintf(T("%"PRIu64" KiB\n"), old_size >> 10);
2321
2322         ret = wimlib_overwrite(w, write_flags, num_threads,
2323                                imagex_progress_func);
2324
2325         if (ret == 0) {
2326                 new_size = file_get_size(argv[0]);
2327                 tprintf(T("\"%"TS"\" optimized size: "), wimfile);
2328                 if (new_size == -1)
2329                         tfputs(T("Unknown\n"), stdout);
2330                 else
2331                         tprintf(T("%"PRIu64" KiB\n"), new_size >> 10);
2332
2333                 tfputs(T("Space saved: "), stdout);
2334                 if (new_size != -1 && old_size != -1) {
2335                         tprintf(T("%lld KiB\n"),
2336                                ((long long)old_size - (long long)new_size) >> 10);
2337                 } else {
2338                         tfputs(T("Unknown\n"), stdout);
2339                 }
2340         }
2341
2342         wimlib_free(w);
2343         return ret;
2344 }
2345
2346 /* Split a WIM into a spanned set */
2347 static int
2348 imagex_split(int argc, tchar **argv)
2349 {
2350         int c;
2351         int open_flags = WIMLIB_OPEN_FLAG_SPLIT_OK;
2352         int write_flags = 0;
2353         unsigned long part_size;
2354         tchar *tmp;
2355         int ret;
2356         WIMStruct *w;
2357
2358         for_opt(c, split_options) {
2359                 switch (c) {
2360                 case IMAGEX_CHECK_OPTION:
2361                         open_flags |= WIMLIB_OPEN_FLAG_CHECK_INTEGRITY;
2362                         write_flags |= WIMLIB_WRITE_FLAG_CHECK_INTEGRITY;
2363                         break;
2364                 default:
2365                         usage(SPLIT);
2366                         return -1;
2367                 }
2368         }
2369         argc -= optind;
2370         argv += optind;
2371
2372         if (argc != 3) {
2373                 usage(SPLIT);
2374                 return -1;
2375         }
2376         part_size = tstrtod(argv[2], &tmp) * (1 << 20);
2377         if (tmp == argv[2] || *tmp) {
2378                 imagex_error(T("Invalid part size \"%"TS"\""), argv[2]);
2379                 imagex_error(T("The part size must be an integer or "
2380                                "floating-point number of megabytes."));
2381                 return -1;
2382         }
2383         ret = wimlib_open_wim(argv[0], open_flags, &w, imagex_progress_func);
2384         if (ret != 0)
2385                 return ret;
2386         ret = wimlib_split(w, argv[1], part_size, write_flags, imagex_progress_func);
2387         wimlib_free(w);
2388         return ret;
2389 }
2390
2391 /* Unmounts a mounted WIM image. */
2392 static int
2393 imagex_unmount(int argc, tchar **argv)
2394 {
2395         int c;
2396         int unmount_flags = 0;
2397         int ret;
2398
2399         for_opt(c, unmount_options) {
2400                 switch (c) {
2401                 case IMAGEX_COMMIT_OPTION:
2402                         unmount_flags |= WIMLIB_UNMOUNT_FLAG_COMMIT;
2403                         break;
2404                 case IMAGEX_CHECK_OPTION:
2405                         unmount_flags |= WIMLIB_UNMOUNT_FLAG_CHECK_INTEGRITY;
2406                         break;
2407                 case IMAGEX_REBULID_OPTION:
2408                         unmount_flags |= WIMLIB_UNMOUNT_FLAG_REBUILD;
2409                         break;
2410                 default:
2411                         usage(UNMOUNT);
2412                         return -1;
2413                 }
2414         }
2415         argc -= optind;
2416         argv += optind;
2417         if (argc != 1) {
2418                 usage(UNMOUNT);
2419                 return -1;
2420         }
2421
2422         ret = wimlib_unmount_image(argv[0], unmount_flags,
2423                                    imagex_progress_func);
2424         if (ret != 0)
2425                 imagex_error(T("Failed to unmount \"%"TS"\""), argv[0]);
2426         return ret;
2427 }
2428
2429 struct imagex_command {
2430         const tchar *name;
2431         int (*func)(int , tchar **);
2432         int cmd;
2433 };
2434
2435
2436 #define for_imagex_command(p) for (p = &imagex_commands[0]; \
2437                 p != &imagex_commands[ARRAY_LEN(imagex_commands)]; p++)
2438
2439 static const struct imagex_command imagex_commands[] = {
2440         {T("append"),  imagex_capture_or_append, APPEND},
2441         {T("apply"),   imagex_apply,             APPLY},
2442         {T("capture"), imagex_capture_or_append, CAPTURE},
2443         {T("delete"),  imagex_delete,            DELETE},
2444         {T("dir"),     imagex_dir,               DIR},
2445         {T("export"),  imagex_export,            EXPORT},
2446         {T("info"),    imagex_info,              INFO},
2447         {T("join"),    imagex_join,              JOIN},
2448         {T("mount"),   imagex_mount_rw_or_ro,    MOUNT},
2449         {T("mountrw"), imagex_mount_rw_or_ro,    MOUNTRW},
2450         {T("optimize"),imagex_optimize,          OPTIMIZE},
2451         {T("split"),   imagex_split,             SPLIT},
2452         {T("unmount"), imagex_unmount,           UNMOUNT},
2453 };
2454
2455 static void
2456 version()
2457 {
2458         static const tchar *s =
2459         T(
2460 IMAGEX_PROGNAME " (" PACKAGE ") " PACKAGE_VERSION "\n"
2461 "Copyright (C) 2012, 2013 Eric Biggers\n"
2462 "License GPLv3+; GNU GPL version 3 or later <http://gnu.org/licenses/gpl.html>.\n"
2463 "This is free software: you are free to change and redistribute it.\n"
2464 "There is NO WARRANTY, to the extent permitted by law.\n"
2465 "\n"
2466 "Report bugs to "PACKAGE_BUGREPORT".\n"
2467         );
2468         tfputs(s, stdout);
2469 }
2470
2471
2472 static void
2473 help_or_version(int argc, tchar **argv)
2474 {
2475         int i;
2476         const tchar *p;
2477         const struct imagex_command *cmd;
2478
2479         for (i = 1; i < argc; i++) {
2480                 p = argv[i];
2481                 if (*p == T('-'))
2482                         p++;
2483                 else
2484                         continue;
2485                 if (*p == T('-'))
2486                         p++;
2487                 if (!tstrcmp(p, T("help"))) {
2488                         for_imagex_command(cmd) {
2489                                 if (!tstrcmp(cmd->name, argv[1])) {
2490                                         usage(cmd->cmd);
2491                                         exit(0);
2492                                 }
2493                         }
2494                         usage_all();
2495                         exit(0);
2496                 }
2497                 if (!tstrcmp(p, T("version"))) {
2498                         version();
2499                         exit(0);
2500                 }
2501         }
2502 }
2503
2504
2505 static void
2506 usage(int cmd_type)
2507 {
2508         const struct imagex_command *cmd;
2509         tprintf(T("Usage:\n%"TS), usage_strings[cmd_type]);
2510         for_imagex_command(cmd) {
2511                 if (cmd->cmd == cmd_type) {
2512                         tprintf(T("\nTry `man "IMAGEX_PROGNAME"-%"TS"' "
2513                                   "for more details.\n"), cmd->name);
2514                 }
2515         }
2516 }
2517
2518 static void
2519 usage_all()
2520 {
2521         tfputs(T("Usage:\n"), stdout);
2522         for (int i = 0; i < ARRAY_LEN(usage_strings); i++)
2523                 tprintf(T("    %"TS), usage_strings[i]);
2524         static const tchar *extra =
2525         T(
2526 "    "IMAGEX_PROGNAME" --help\n"
2527 "    "IMAGEX_PROGNAME" --version\n"
2528 "\n"
2529 "    The compression TYPE may be \"maximum\", \"fast\", or \"none\".\n"
2530 "\n"
2531 "    Try `man "IMAGEX_PROGNAME"' for more information.\n"
2532         );
2533         tfputs(extra, stdout);
2534 }
2535
2536 /* Entry point for wimlib's ImageX implementation.  On UNIX the command
2537  * arguments will just be 'char' strings (ideally UTF-8 encoded, but could be
2538  * something else), while an Windows the command arguments will be UTF-16LE
2539  * encoded 'wchar_t' strings. */
2540 int
2541 #ifdef __WIN32__
2542 wmain(int argc, wchar_t **argv, wchar_t **envp)
2543 #else
2544 main(int argc, char **argv)
2545 #endif
2546 {
2547         const struct imagex_command *cmd;
2548         int ret;
2549         int init_flags = 0;
2550
2551 #ifndef __WIN32__
2552         if (getenv("WIMLIB_IMAGEX_USE_UTF8")) {
2553                 init_flags |= WIMLIB_INIT_FLAG_ASSUME_UTF8;
2554         } else {
2555                 char *codeset;
2556
2557                 setlocale(LC_ALL, "");
2558                 codeset = nl_langinfo(CODESET);
2559                 if (!strstr(codeset, "UTF-8") &&
2560                     !strstr(codeset, "UTF8") &&
2561                     !strstr(codeset, "utf-8") &&
2562                     !strstr(codeset, "utf8"))
2563                 {
2564                         fputs(
2565 "WARNING: Running "IMAGEX_PROGNAME" in a UTF-8 locale is recommended!\n"
2566 "         Maybe try: `export LANG=en_US.UTF-8'?\n"
2567 "         Alternatively, set the environmental variable WIMLIB_IMAGEX_USE_UTF8\n"
2568 "         to any value to force wimlib to use UTF-8.\n",
2569                         stderr);
2570
2571                 }
2572         }
2573 #endif /* !__WIN32__ */
2574
2575         if (argc < 2) {
2576                 imagex_error(T("No command specified"));
2577                 usage_all();
2578                 ret = 2;
2579                 goto out;
2580         }
2581
2582         /* Handle --help and --version for all commands.  Note that this will
2583          * not return if either of these arguments are present. */
2584         help_or_version(argc, argv);
2585         argc--;
2586         argv++;
2587
2588         /* The user may like to see more informative error messages. */
2589         wimlib_set_print_errors(true);
2590
2591         /* Do any initializations that the library needs */
2592         ret = wimlib_global_init(init_flags);
2593         if (ret)
2594                 goto out_check_status;
2595
2596         /* Search for the function to handle the ImageX subcommand. */
2597         for_imagex_command(cmd) {
2598                 if (!tstrcmp(cmd->name, *argv)) {
2599                         ret = cmd->func(argc, argv);
2600                         goto out_check_write_error;
2601                 }
2602         }
2603
2604         imagex_error(T("Unrecognized command: `%"TS"'"), argv[0]);
2605         usage_all();
2606         ret = 2;
2607         goto out_cleanup;
2608 out_check_write_error:
2609         /* For 'wimlib-imagex info' and 'wimlib-imagex dir', data printed to
2610          * standard output is part of the program's actual behavior and not just
2611          * for informational purposes, so we should set a failure exit status if
2612          * there was a write error. */
2613         if (cmd == &imagex_commands[INFO] || cmd == &imagex_commands[DIR]) {
2614                 if (ferror(stdout) || fclose(stdout)) {
2615                         imagex_error_with_errno(T("error writing to standard output"));
2616                         if (ret == 0)
2617                                 ret = -1;
2618                 }
2619         }
2620 out_check_status:
2621         /* Exit status (ret):  -1 indicates an error found by 'wimlib-imagex'
2622          * outside of the wimlib library code.  0 indicates success.  > 0
2623          * indicates a wimlib error code from which an error message can be
2624          * printed. */
2625         if (ret > 0) {
2626                 imagex_error(T("Exiting with error code %d:\n"
2627                                "       %"TS"."), ret,
2628                              wimlib_get_error_string(ret));
2629                 if (ret == WIMLIB_ERR_NTFS_3G && errno != 0)
2630                         imagex_error_with_errno(T("errno"));
2631         }
2632 out_cleanup:
2633         /* Make the library free any resources it's holding (not strictly
2634          * necessary because the process is ending anyway). */
2635         wimlib_global_cleanup();
2636 out:
2637         return ret;
2638 }