imagex_update(): Acquire capture privileges on Windows

[wimlib] / programs / imagex.c

diff --git a/programs/imagex.c b/programs/imagex.c
index 591714e1684222a5b8f320025a3dfba7fbee7700..601ad01ca53374a06c63007c5174b4fee5b4f26b 100644 (file)
--- a/programs/imagex.c
+++ b/programs/imagex.c
@@ -22,7 +22,10 @@
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
 
-#include "config.h"
+#ifdef HAVE_CONFIG_H
+#  include "config.h" /* Need for PACKAGE_VERSION, etc. */
+#endif
+
 #include "wimlib.h"
 #include "wimlib_tchar.h"
 
@@ -40,7 +43,7 @@
 #include <locale.h>
 
 #ifdef HAVE_ALLOCA_H
-#include <alloca.h>
+#  include <alloca.h>
 #endif
 
 #ifdef __WIN32__
@@ -80,7 +83,7 @@ enum imagex_op_type {
 };
 
 static void usage(int cmd_type);
-static void usage_all();
+static void usage_all(void);
 
 static bool imagex_be_quiet = false;
 
@@ -165,11 +168,11 @@ 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(
-IMAGEX_PROGNAME" update WIMFILE IMAGE [--check] [--rebuild]\n"
+IMAGEX_PROGNAME" update WIMFILE [IMAGE_NUM | IMAGE_NAME] [--check] [--rebuild]\n"
 "                       [--threads=NUM_THREADS] [DEFAULT_ADD_OPTIONS]\n"
 "                       [DEFAULT_DELETE_OPTIONS] < CMDFILE\n"
 ),
@@ -179,14 +182,15 @@ IMAGEX_PROGNAME" update WIMFILE IMAGE [--check] [--rebuild]\n"
 static void
 recommend_man_page(const tchar *cmd_name)
 {
+       const tchar *format_str;
 #ifdef __WIN32__
-       tprintf(T("See "IMAGEX_PROGNAME"-%"TS".pdf in the "
-                 "doc directory for more details.\n"),
-               cmd_name);
+       format_str = T("See "IMAGEX_PROGNAME"%"TS"%"TS".pdf in the "
+                      "doc directory for more details.\n");
 #else
-       tprintf(T("Try `man "IMAGEX_PROGNAME"-%"TS"' "
-                 "for more details.\n"), cmd_name);
+       format_str = T("Try `man "IMAGEX_PROGNAME"%"TS"%"TS"' "
+                      "for more details.\n");
 #endif
+       tprintf(format_str, *cmd_name ? T("-") : T(""), cmd_name);
 }
 
 enum {
@@ -204,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,
@@ -331,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},
 };
 
@@ -1513,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;
@@ -1922,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;
@@ -2693,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);
@@ -2780,9 +2786,13 @@ imagex_optimize(int argc, tchar **argv)
 
        wimfile = argv[0];
 
+       ret = file_writable(wimfile);
+       if (ret)
+               return ret;
+
        ret = wimlib_open_wim(wimfile, open_flags, &w,
                              imagex_progress_func);
-       if (ret != 0)
+       if (ret)
                return ret;
 
        old_size = file_get_size(argv[0]);
@@ -2880,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;
@@ -2921,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;
@@ -2983,6 +2997,11 @@ imagex_update(int argc, tchar **argv)
        if (argc < 1 || argc > 2)
                goto out_usage;
        wimfile = argv[0];
+
+       ret = file_writable(wimfile);
+       if (ret)
+               goto out;
+
        ret = wimlib_open_wim(wimfile, open_flags, &wim, imagex_progress_func);
        if (ret)
                goto out;
@@ -2999,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;
@@ -3051,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:
@@ -3109,7 +3145,7 @@ static const struct imagex_command imagex_commands[] = {
 };
 
 static void
-version()
+version(void)
 {
        static const tchar *s =
        T(
@@ -3172,7 +3208,7 @@ usage(int cmd_type)
 }
 
 static void
-usage_all()
+usage_all(void)
 {
        tfputs(T("Usage:\n"), stdout);
        for (int i = 0; i < ARRAY_LEN(usage_strings); i++)
@@ -3184,9 +3220,9 @@ usage_all()
 "\n"
 "    The compression TYPE may be \"maximum\", \"fast\", or \"none\".\n"
 "\n"
-"    Try `man "IMAGEX_PROGNAME"' for more information.\n"
        );
        tfputs(extra, stdout);
+       recommend_man_page(T(""));
 }
 
 /* Entry point for wimlib's ImageX implementation.  On UNIX the command