env: Implement support for AES encryption into fw_* tools

Implement support for encrypting/decrypting the environment block
into the tools/env/fw_* tools. The cipher used is AES 128 CBC and
the implementation depends solely on components internal to U-Boot.

To allow building against the internal AES library, the library did
need minor adjustments to not include U-Boot's headers which are not
wanted to be included and define missing types.

Signed-off-by: Marek Vasut <marex@denx.de>
diff --git a/include/aes.h b/include/aes.h
index 4897e6f..ee0e6c2 100644
--- a/include/aes.h
+++ b/include/aes.h
@@ -8,6 +8,13 @@
 #ifndef _AES_REF_H_
 #define _AES_REF_H_
 
+#ifdef USE_HOSTCC
+/* Define compat stuff for use in fw_* tools. */
+typedef unsigned char u8;
+typedef unsigned int u32;
+#define debug(...) do {} while (0)
+#endif
+
 /*
  * AES encryption library, with small code size, supporting only 128-bit AES
  *
diff --git a/lib/aes.c b/lib/aes.c
index 9dadb22..05c97cd 100644
--- a/lib/aes.c
+++ b/lib/aes.c
@@ -22,7 +22,11 @@
  * REDISTRIBUTION OF THIS SOFTWARE.
 */
 
+#ifndef USE_HOSTCC
 #include <common.h>
+#else
+#include <string.h>
+#endif
 #include "aes.h"
 
 /* forward s-box */
diff --git a/tools/env/Makefile b/tools/env/Makefile
index fcb752d..f5368bc 100644
--- a/tools/env/Makefile
+++ b/tools/env/Makefile
@@ -25,7 +25,7 @@
 
 fw_printenv_unstripped-objs := fw_env.o fw_env_main.o \
 	crc32.o ctype.o linux_string.o \
-	env_attr.o env_flags.o
+	env_attr.o env_flags.o aes.o
 
 quiet_cmd_strip = STRIP   $@
       cmd_strip = $(STRIP) -o $@ $<
diff --git a/tools/env/fw_env.c b/tools/env/fw_env.c
index d228cc3..fba4c8c 100644
--- a/tools/env/fw_env.c
+++ b/tools/env/fw_env.c
@@ -31,6 +31,10 @@
 
 #include "fw_env.h"
 
+#include <aes.h>
+
+#define DIV_ROUND_UP(n, d)	(((n) + (d) - 1) / (d))
+
 #define WHITESPACE(c) ((c == '\t') || (c == ' '))
 
 #define min(x, y) ({				\
@@ -98,6 +102,11 @@
 	.flag_scheme = FLAG_NONE,
 };
 
+/* Is AES encryption used? */
+static int aes_flag;
+static uint8_t aes_key[AES_KEY_LENGTH] = { 0 };
+static int env_aes_cbc_crypt(char *data, const int enc);
+
 static int HaveRedundEnv = 0;
 
 static unsigned char active_flag = 1;
@@ -120,6 +129,10 @@
 
 	if (HaveRedundEnv)
 		rc -= sizeof (char);
+
+	if (aes_flag)
+		rc &= ~(AES_KEY_LENGTH - 1);
+
 	return rc;
 }
 
@@ -191,6 +204,36 @@
 	return NULL;
 }
 
+static int parse_aes_key(char *key)
+{
+	char tmp[5] = { '0', 'x', 0, 0, 0 };
+	unsigned long ul;
+	int i;
+
+	if (strnlen(key, 64) != 32) {
+		fprintf(stderr,
+			"## Error: '-a' option requires 16-byte AES key\n");
+		return -1;
+	}
+
+	for (i = 0; i < 16; i++) {
+		tmp[2] = key[0];
+		tmp[3] = key[1];
+		errno = 0;
+		ul = strtoul(tmp, NULL, 16);
+		if (errno) {
+			fprintf(stderr,
+				"## Error: '-a' option requires valid AES key\n");
+			return -1;
+		}
+		aes_key[i] = ul & 0xff;
+		key += 2;
+	}
+	aes_flag = 1;
+
+	return 0;
+}
+
 /*
  * Print the current definition of one, or more, or all
  * environment variables
@@ -201,6 +244,19 @@
 	int i, n_flag;
 	int rc = 0;
 
+	if (argc >= 2 && strcmp(argv[1], "-a") == 0) {
+		if (argc < 3) {
+			fprintf(stderr,
+				"## Error: '-a' option requires AES key\n");
+			return -1;
+		}
+		rc = parse_aes_key(argv[2]);
+		if (rc)
+			return rc;
+		argv += 2;
+		argc -= 2;
+	}
+
 	if (fw_env_open())
 		return -1;
 
@@ -266,6 +322,16 @@
 
 int fw_env_close(void)
 {
+	int ret;
+	if (aes_flag) {
+		ret = env_aes_cbc_crypt(environment.data, 1);
+		if (ret) {
+			fprintf(stderr,
+				"Error: can't encrypt env for flash\n");
+			return ret;
+		}
+	}
+
 	/*
 	 * Update CRC
 	 */
@@ -413,7 +479,7 @@
  */
 int fw_setenv(int argc, char *argv[])
 {
-	int i;
+	int i, rc;
 	size_t len;
 	char *name;
 	char *value = NULL;
@@ -423,6 +489,24 @@
 		return -1;
 	}
 
+	if (strcmp(argv[1], "-a") == 0) {
+		if (argc < 3) {
+			fprintf(stderr,
+				"## Error: '-a' option requires AES key\n");
+			return -1;
+		}
+		rc = parse_aes_key(argv[2]);
+		if (rc)
+			return rc;
+		argv += 2;
+		argc -= 2;
+	}
+
+	if (argc < 2) {
+		errno = EINVAL;
+		return -1;
+	}
+
 	if (fw_env_open()) {
 		fprintf(stderr, "Error: environment not initialized\n");
 		return -1;
@@ -900,6 +984,28 @@
 	return rc;
 }
 
+/* Encrypt or decrypt the environment before writing or reading it. */
+static int env_aes_cbc_crypt(char *payload, const int enc)
+{
+	uint8_t *data = (uint8_t *)payload;
+	const int len = getenvsize();
+	uint8_t key_exp[AES_EXPAND_KEY_LENGTH];
+	uint32_t aes_blocks;
+
+	/* First we expand the key. */
+	aes_expand_key(aes_key, key_exp);
+
+	/* Calculate the number of AES blocks to encrypt. */
+	aes_blocks = DIV_ROUND_UP(len, AES_KEY_LENGTH);
+
+	if (enc)
+		aes_cbc_encrypt_blocks(key_exp, data, data, aes_blocks);
+	else
+		aes_cbc_decrypt_blocks(key_exp, data, data, aes_blocks);
+
+	return 0;
+}
+
 static int flash_write (int fd_current, int fd_target, int dev_target)
 {
 	int rc;
@@ -923,6 +1029,7 @@
 	fprintf(stderr, "Writing new environment at 0x%lx on %s\n",
 		DEVOFFSET (dev_target), DEVNAME (dev_target));
 #endif
+
 	rc = flash_write_buf(dev_target, fd_target, environment.image,
 			      CUR_ENVSIZE, DEVOFFSET(dev_target),
 			      DEVTYPE(dev_target));
@@ -981,8 +1088,10 @@
 
 	rc = flash_read_buf(dev_current, fd, environment.image, CUR_ENVSIZE,
 			     DEVOFFSET (dev_current), mtdinfo.type);
+	if (rc != CUR_ENVSIZE)
+		return -1;
 
-	return (rc != CUR_ENVSIZE) ? -1 : 0;
+	return 0;
 }
 
 static int flash_io (int mode)
@@ -1075,6 +1184,8 @@
 	unsigned char flag1;
 	void *addr1;
 
+	int ret;
+
 	struct env_image_single *single;
 	struct env_image_redundant *redundant;
 
@@ -1109,6 +1220,13 @@
 		return -1;
 
 	crc0 = crc32 (0, (uint8_t *) environment.data, ENV_SIZE);
+
+	if (aes_flag) {
+		ret = env_aes_cbc_crypt(environment.data, 0);
+		if (ret)
+			return ret;
+	}
+
 	crc0_ok = (crc0 == *environment.crc);
 	if (!HaveRedundEnv) {
 		if (!crc0_ok) {
@@ -1159,6 +1277,13 @@
 		}
 
 		crc1 = crc32 (0, (uint8_t *) redundant->data, ENV_SIZE);
+
+		if (aes_flag) {
+			ret = env_aes_cbc_crypt(redundant->data, 0);
+			if (ret)
+				return ret;
+		}
+
 		crc1_ok = (crc1 == redundant->crc);
 		flag1 = redundant->flags;
 
diff --git a/tools/env/fw_env_main.c b/tools/env/fw_env_main.c
index 2b85d78..ce50d58 100644
--- a/tools/env/fw_env_main.c
+++ b/tools/env/fw_env_main.c
@@ -9,18 +9,22 @@
  * Command line user interface to firmware (=U-Boot) environment.
  *
  * Implements:
- *	fw_printenv [[ -n name ] | [ name ... ]]
+ *	fw_printenv [ -a key ] [[ -n name ] | [ name ... ]]
  *              - prints the value of a single environment variable
  *                "name", the ``name=value'' pairs of one or more
  *                environment variables "name", or the whole
  *                environment if no names are specified.
- *	fw_setenv name [ value ... ]
+ *	fw_setenv [ -a key ] name [ value ... ]
  *		- If a name without any values is given, the variable
  *		  with this name is deleted from the environment;
  *		  otherwise, all "value" arguments are concatenated,
  *		  separated by single blank characters, and the
  *		  resulting string is assigned to the environment
  *		  variable "name"
+ *
+ * If '-a key' is specified, the env block is encrypted with AES 128 CBC.
+ * The 'key' argument is in the format of 32 hexadecimal numbers (16 bytes
+ * of AES key), eg. '-a aabbccddeeff00112233445566778899'.
  */
 
 #include <fcntl.h>
@@ -46,8 +50,8 @@
 
 	fprintf(stderr, "fw_printenv/fw_setenv, "
 		"a command line interface to U-Boot environment\n\n"
-		"usage:\tfw_printenv [-n] [variable name]\n"
-		"\tfw_setenv [variable name] [variable value]\n"
+		"usage:\tfw_printenv [-a key] [-n] [variable name]\n"
+		"\tfw_setenv [-a key] [variable name] [variable value]\n"
 		"\tfw_setenv -s [ file ]\n"
 		"\tfw_setenv -s - < [ file ]\n\n"
 		"The file passed as argument contains only pairs "
@@ -94,9 +98,12 @@
 		cmdname = p + 1;
 	}
 
-	while ((c = getopt_long (argc, argv, "ns:h",
+	while ((c = getopt_long (argc, argv, "a:ns:h",
 		long_options, NULL)) != EOF) {
 		switch (c) {
+		case 'a':
+			/* AES key, handled later */
+			break;
 		case 'n':
 			/* handled in fw_printenv */
 			break;