1 files changed, 71 insertions, 2 deletions

diff --git a/ch-flash.c b/ch-flash.c
index 548487d..e90a995 100644
--- a/ch-flash.c
+++ b/ch-flash.c
@@ -18,6 +18,7 @@
 
 #include "ch-flash.h"
 
+#include <assert.h>
 #include <err.h>
 #include <errno.h>
 #include <fcntl.h>
@@ -134,7 +135,7 @@ static void usage(void)
 {
   printf("ISP programmer for some WinChipHead MCUs\n");
   printf("Options:\n");
-  printf("  --code-flash, -f    firmware to flash\n");
+  printf("  --code-flash, -f    firmware to flash.\n");
   printf("  --code-verify, -c   verify existing firwmare\n");
   printf("  --data-flash, -k    data to flash\n");
   printf("  --data-verify, -l   verify existing data\n");
@@ -157,6 +158,18 @@ static void hexdump(const char *name, const void *data, int len)
   printf("\n");
 }
 
+static uint32_t chksum(const void *data, int len)
+{
+  const uint8_t *data_u8 = data;
+  uint32_t ret = 0;
+  while (len) {
+    ret += *data_u8;
+    data_u8++;
+    len--;
+  }
+  return ret;
+}
+
 /* Open and claim the USB device */
 static void open_usb_device(struct device *dev)
 {
@@ -318,6 +331,55 @@ static void erase_code_flash(struct device *dev)
     errx(EXIT_FAILURE, "The device refused to erase the code flash");
 }
 
+static void load_shell_command(struct device *dev, struct content *info)
+{
+  assert(strncmp(info->filename, "shell:", 6) == 0);
+  const char *shell_cmd = info->filename + 6;
+
+  size_t allocd_size = 4096;
+  size_t total_size = 0;
+  size_t block_size = 0;
+
+  uint8_t *buffer = malloc(allocd_size + 16);
+  size_t cursor = 0;
+
+  FILE *fp;
+
+  // Open the command for reading
+  printf("Running the shell command: %s\n", shell_cmd);
+  fp = popen(shell_cmd, "r");
+  if (fp == NULL) {
+    err(EXIT_FAILURE, "Unable to run the command");
+  }
+
+  int chunk_size = 0;
+  while (
+      (chunk_size = fread(buffer + cursor, 1, allocd_size - total_size, fp))) {
+    total_size += chunk_size;
+    cursor += chunk_size;
+    if (total_size == allocd_size) {
+      buffer = realloc(buffer, allocd_size * 2 + 16);
+      allocd_size *= 2;
+    }
+  }
+
+  // Close the pipe
+  if (pclose(fp) == -1) {
+    err(EXIT_FAILURE, "Pclose failed.");
+  }
+
+  info->buf = buffer;
+  info->len = total_size;
+
+  /* Round up to 8 bytes boundary as upload protocol requires
+   * it. Extra bytes are zeroes. */
+  info->len = (info->len + 7) & ~7;
+  memset(info->buf + total_size, 0xff, info->len - total_size);
+
+  if (dev->debug) hexdump("shell cmd output", info->buf, info->len);
+  printf("Command contents Checksum: %08x\n", chksum(info->buf, info->len));
+}
+
 static void load_file(struct device *dev, struct content *info)
 {
   struct stat statbuf;
@@ -348,6 +410,9 @@ static void load_file(struct device *dev, struct content *info)
   if (ret != statbuf.st_size) err(EXIT_FAILURE, "Can't read firmware file");
 
   close(fd);
+
+  if (dev->debug) hexdump("code file", info->buf, info->len);
+  printf("File contents Checksum: %08x\n", chksum(info->buf, info->len));
 }
 
 /* Encrypt (or decrypt) some data */
@@ -712,7 +777,11 @@ int main(int argc, char *argv[])
   create_key(&dev);
 
   if (do_code_flash || do_code_verify) {
-    load_file(&dev, &dev.fw);
+    if (strncmp(dev.fw.filename, "shell:", 6) == 0) {
+      load_shell_command(&dev, &dev.fw);
+    } else {
+      load_file(&dev, &dev.fw);
+    }
     encrypt(&dev, &dev.fw);
   }