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