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