From cd6541afd6f650915298f58a876dd1130b23876e Mon Sep 17 00:00:00 2001
From: Eric Biggers <ebiggers3@gmail.com>
Date: Mon, 28 Apr 2014 02:14:05 -0500
Subject: [PATCH] wimupdate:  Add --wimboot-config=FILE option

---
 doc/man1/imagex-update.1.in | 16 ++++++++++++++
 programs/imagex.c           | 44 ++++++++++++++++++++++++++++++++-----
 2 files changed, 55 insertions(+), 5 deletions(-)

diff --git a/doc/man1/imagex-update.1.in b/doc/man1/imagex-update.1.in
index f688e6b7..3e1e7417 100644
--- a/doc/man1/imagex-update.1.in
+++ b/doc/man1/imagex-update.1.in
@@ -128,6 +128,22 @@ quotes and single quotes for the inner quotes.  Example:
 .RE
 .RE
 .fi
+.TP
+\fB--wimboot-config\fR=\fIFILE\fR
+If this option is specified, no commands shall be read from standard input, and
+instead the following command shall be executed:
+.IP
+.nf
+.RS
+.RS
+\fBadd\fR \fIFILE\fR /Windows/System32/WimBootCompress.ini
+.RE
+.RE
+.fi
+.IP
+This sets \fIFILE\fR as the WIMBoot configuration file for the image.  See the
+documentation for the \fB--wimboot\fR option of \fB@IMAGEX_PROGNAME@ apply\fR
+(1) for more information.
 .SH NOTES
 \fB@IMAGEX_PROGNAME@ update\fR is partly redundant with \fB@IMAGEX_PROGNAME@
 mountrw\fR, since if a WIM image can be mounted read-write, then there
diff --git a/programs/imagex.c b/programs/imagex.c
index c3a08e8e..25a4471b 100644
--- a/programs/imagex.c
+++ b/programs/imagex.c
@@ -201,6 +201,7 @@ enum {
 	IMAGEX_UPDATE_OF_OPTION,
 	IMAGEX_VERBOSE_OPTION,
 	IMAGEX_WIMBOOT_OPTION,
+	IMAGEX_WIMBOOT_CONFIG_OPTION,
 	IMAGEX_XML_OPTION,
 };
 
@@ -387,6 +388,7 @@ static const struct option update_options[] = {
 	{T("check"),       no_argument,       NULL, IMAGEX_CHECK_OPTION},
 	{T("rebuild"),     no_argument,       NULL, IMAGEX_REBUILD_OPTION},
 	{T("command"),     required_argument, NULL, IMAGEX_COMMAND_OPTION},
+	{T("wimboot-config"), required_argument, NULL, IMAGEX_WIMBOOT_CONFIG_OPTION},
 
 	/* Default delete options */
 	{T("force"),       no_argument,       NULL, IMAGEX_FORCE_OPTION},
@@ -3720,6 +3722,7 @@ imagex_update(int argc, tchar **argv, int cmd)
 	size_t num_cmds;
 	tchar *command_str = NULL;
 	tchar *config_file = NULL;
+	tchar *wimboot_config = NULL;
 
 	for_opt(c, update_options) {
 		switch (c) {
@@ -3750,6 +3753,9 @@ imagex_update(int argc, tchar **argv, int cmd)
 				goto out_err;
 			}
 			break;
+		case IMAGEX_WIMBOOT_CONFIG_OPTION:
+			wimboot_config = optarg;
+			break;
 		/* Default delete options */
 		case IMAGEX_FORCE_OPTION:
 			default_delete_flags |= WIMLIB_DELETE_FLAG_FORCE;
@@ -3826,7 +3832,11 @@ imagex_update(int argc, tchar **argv, int cmd)
 		cmd_file_contents = NULL;
 		cmds = parse_update_command_file(&command_str, tstrlen(command_str),
 						 &num_cmds);
-	} else {
+		if (!cmds) {
+			ret = -1;
+			goto out_free_cmd_file_contents;
+		}
+	} else if (!wimboot_config) {
 		if (isatty(STDIN_FILENO)) {
 			tputs(T("Reading update commands from standard input..."));
 			recommend_man_page(CMD_UPDATE, stdout);
@@ -3840,10 +3850,14 @@ imagex_update(int argc, tchar **argv, int cmd)
 		/* Parse the update commands */
 		cmds = parse_update_command_file(&cmd_file_contents, cmd_file_nchars,
 						 &num_cmds);
-	}
-	if (!cmds) {
-		ret = -1;
-		goto out_free_cmd_file_contents;
+		if (!cmds) {
+			ret = -1;
+			goto out_free_cmd_file_contents;
+		}
+	} else {
+		cmd_file_contents = NULL;
+		cmds = NULL;
+		num_cmds = 0;
 	}
 
 	/* Set default flags and capture config on the update commands */
@@ -3867,6 +3881,26 @@ imagex_update(int argc, tchar **argv, int cmd)
 	if (ret)
 		goto out_free_cmds;
 
+	if (wimboot_config) {
+		/* --wimboot-config=FILE is short for an
+		 * "add FILE /Windows/System32/WimBootCompress.ini" command.
+		 */
+		struct wimlib_update_command cmd = {
+			.op = WIMLIB_UPDATE_OP_ADD,
+			.add = {
+				.fs_source_path = wimboot_config,
+				.wim_target_path =
+					T("/Windows/System32/WimBootCompress.ini"),
+				.config_file = NULL,
+				.add_flags = 0,
+			},
+		};
+		ret = wimlib_update_image(wim, image, &cmd, 1,
+					  update_flags, imagex_progress_func);
+		if (ret)
+			goto out_free_cmds;
+	}
+
 	/* Overwrite the updated WIM */
 	ret = wimlib_overwrite(wim, write_flags, num_threads,
 			       imagex_progress_func);
-- 
2.43.0