imagex_update(): Acquire capture privileges on Windows

[wimlib] / programs / imagex.c

diff --git a/programs/imagex.c b/programs/imagex.c
index 7dbbc8a4330addd60c8a1f668d62853996ff4e37..601ad01ca53374a06c63007c5174b4fee5b4f26b 100644 (file)
--- a/programs/imagex.c
+++ b/programs/imagex.c
@@ -168,7 +168,7 @@ IMAGEX_PROGNAME" split WIMFILE SPLIT_WIMFILE PART_SIZE_MB [--check]\n"
 ),
 [UNMOUNT] =
 T(
-IMAGEX_PROGNAME" unmount DIRECTORY [--commit] [--check] [--rebuild]\n"
+IMAGEX_PROGNAME" unmount DIRECTORY [--commit] [--check] [--rebuild] [--lazy]\n"
 ),
 [UPDATE] =
 T(
@@ -208,6 +208,7 @@ enum {
        IMAGEX_FORCE_OPTION,
        IMAGEX_HARDLINK_OPTION,
        IMAGEX_HEADER_OPTION,
+       IMAGEX_LAZY_OPTION,
        IMAGEX_LOOKUP_TABLE_OPTION,
        IMAGEX_METADATA_OPTION,
        IMAGEX_NORPFIX_OPTION,
@@ -335,6 +336,7 @@ static const struct option unmount_options[] = {
        {T("commit"),  no_argument, NULL, IMAGEX_COMMIT_OPTION},
        {T("check"),   no_argument, NULL, IMAGEX_CHECK_OPTION},
        {T("rebuild"), no_argument, NULL, IMAGEX_REBUILD_OPTION},
+       {T("lazy"),    no_argument, NULL, IMAGEX_LAZY_OPTION},
        {NULL, 0, NULL, 0},
 };
 
@@ -1517,7 +1519,7 @@ imagex_apply(int argc, tchar **argv)
        num_images = wimlib_get_num_images(w);
        if (argc == 2 && num_images != 1) {
                imagex_error(T("\"%"TS"\" contains %d images; Please select one "
-                              "(or all)"), wimfile, num_images);
+                              "(or all)."), wimfile, num_images);
                usage(APPLY);
                ret = -1;
                goto out;
@@ -1926,7 +1928,7 @@ imagex_dir(int argc, tchar **argv)
                 * choose that one; otherwise, print an error. */
                num_images = wimlib_get_num_images(w);
                if (num_images != 1) {
-                       imagex_error(T("The file \"%"TS"\" contains %d images; Please "
+                       imagex_error(T("\"%"TS"\" contains %d images; Please "
                                       "select one."), wimfile, num_images);
                        usage(DIR);
                        ret = -1;
@@ -2697,7 +2699,7 @@ imagex_mount_rw_or_ro(int argc, tchar **argv)
                image = 1;
                num_images = wimlib_get_num_images(w);
                if (num_images != 1) {
-                       imagex_error(T("The file \"%"TS"\" contains %d images; Please "
+                       imagex_error(T("\"%"TS"\" contains %d images; Please "
                                       "select one."), wimfile, num_images);
                        usage((mount_flags & WIMLIB_MOUNT_FLAG_READWRITE)
                                        ? MOUNTRW : MOUNT);
@@ -2888,6 +2890,9 @@ imagex_unmount(int argc, tchar **argv)
                case IMAGEX_REBUILD_OPTION:
                        unmount_flags |= WIMLIB_UNMOUNT_FLAG_REBUILD;
                        break;
+               case IMAGEX_LAZY_OPTION:
+                       unmount_flags |= WIMLIB_UNMOUNT_FLAG_LAZY;
+                       break;
                default:
                        usage(UNMOUNT);
                        return -1;
@@ -2929,6 +2934,7 @@ imagex_update(int argc, tchar **argv)
        size_t cmd_file_nchars;
        struct wimlib_update_command *cmds;
        size_t num_cmds;
+       int num_images;
 
        const tchar *config_file = NULL;
        tchar *config_str;
@@ -3012,6 +3018,15 @@ imagex_update(int argc, tchar **argv)
        if (ret)
                goto out_wimlib_free;
 
+       num_images = wimlib_get_num_images(wim);
+       if (argc == 1 && num_images != 1) {
+               imagex_error(T("\"%"TS"\" contains %d images; Please select one."),
+                            wimfile, num_images);
+               usage(UPDATE);
+               ret = -1;
+               goto out_wimlib_free;
+       }
+
        /* Parse capture configuration file if specified */
        if (config_file) {
                size_t config_len;
@@ -3064,15 +3079,23 @@ imagex_update(int argc, tchar **argv)
                }
        }
 
+#ifdef __WIN32__
+       win32_acquire_capture_privileges();
+#endif
+
        /* Execute the update commands */
        ret = wimlib_update_image(wim, image, cmds, num_cmds, update_flags,
                                  imagex_progress_func);
        if (ret)
-               goto out_free_cmds;
+               goto out_release_privs;
 
        /* Overwrite the updated WIM */
        ret = wimlib_overwrite(wim, write_flags, num_threads,
                               imagex_progress_func);
+out_release_privs:
+#ifdef __WIN32__
+       win32_release_capture_privileges();
+#endif
 out_free_cmds:
        free(cmds);
 out_free_cmd_file_contents: