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