From fb31788504f035b9577ce230eee4e49ddc7b3baf Mon Sep 17 00:00:00 2001
From: Shivani Tipnis <shivani@espressif.com>
Date: Mon, 8 Oct 2018 11:11:56 +0530
Subject: [PATCH] nvs_part_gen: Add changes to partition size input

---
 .../nvs_partition_generator/README.rst        |  2 +-
 .../nvs_partition_gen.py                      | 23 +++++++++----------
 tools/mass_mfg/docs/README.rst                |  4 ++--
 tools/mass_mfg/mfg_gen.py                     |  5 ++--
 4 files changed, 17 insertions(+), 17 deletions(-)

diff --git a/components/nvs_flash/nvs_partition_generator/README.rst b/components/nvs_flash/nvs_partition_generator/README.rst
index aa7a945c2..f0c6d6623 100644
--- a/components/nvs_flash/nvs_partition_generator/README.rst
+++ b/components/nvs_flash/nvs_partition_generator/README.rst
@@ -123,7 +123,7 @@ A sample CSV file is provided with the utility::
 +------------------------+----------------------------------------------------------------------------------------------+
 | output                 | Path to output converted binary file. Will use stdout if omitted                             |
 +------------------------+----------------------------------------------------------------------------------------------+
-| size                   | Size of NVS Partition in KB. Eg. 12KB                                                        |
+| size                   | Size of NVS Partition in hex (must be multiple of 4096). Eg. 0x1000                          |
 +------------------------+----------------------------------------------------------------------------------------------+
 | --version {v1,v2}      |  Set version. Default: v2                                                                    |
 +-------------------------------+---------------------------------------------------------------------------------------+
diff --git a/components/nvs_flash/nvs_partition_generator/nvs_partition_gen.py b/components/nvs_flash/nvs_partition_generator/nvs_partition_gen.py
index a1cc81a65..2b11877ca 100755
--- a/components/nvs_flash/nvs_partition_generator/nvs_partition_gen.py
+++ b/components/nvs_flash/nvs_partition_generator/nvs_partition_gen.py
@@ -179,7 +179,7 @@ class Page(object):
 
             data_input = byte_arr
 
-        
+
         if sys.version_info[0] < 3:
             data_input = binascii.hexlify(bytearray(data_input))
         else:
@@ -214,7 +214,6 @@ class Page(object):
                 data_val = str(data_bytes) + (init_data_val * (data_len_needed - len(data_bytes)))
 
             encr_data_ret = self.encrypt_entry(data_val, tweak_val, encr_key_input)
-            #print("\n<<<\n")
             encr_data_to_write = encr_data_to_write + encr_data_ret
             # Update values for encrypting next set of data bytes
             start_idx = end_idx
@@ -233,7 +232,7 @@ class Page(object):
                 encr_data[0:len(encr_data_ret)] = encr_data_ret
             else:
                 encr_data[0:len(encr_data_ret)] = encr_data_ret
-            
+
             data = encr_data
 
         data_offset = Page.FIRST_ENTRY_OFFSET + (Page.SINGLE_ENTRY_SIZE * self.entry_num)
@@ -311,7 +310,7 @@ class Page(object):
             if not sys.version_info[0] < 3:
                 if type(data_chunk) == str:
                     data_chunk = codecs.encode(data_chunk, 'utf8')
-            
+
             crc = zlib.crc32(data_chunk, 0xFFFFFFFF)
             struct.pack_into('<I', entry_struct, 28, crc & 0xFFFFFFFF)
 
@@ -466,7 +465,7 @@ class Page(object):
         # write key
         key_array = bytearray(b'\x00')*16
         entry_struct[8:24] = key_array
-        
+
         if sys.version_info[0] < 3:
             entry_struct[8:8 + len(key)] = key
         else:
@@ -748,7 +747,7 @@ def nvs_part_gen(input_filename=None, output_filename=None, input_size=None, key
         input_file = open(input_filename, 'rb')
     else:
         input_file = open(input_filename, 'r', newline='')
-    
+
     output_file = open(output_filename, 'wb')
 
     with nvs_open(output_file, input_size) as nvs_obj:
@@ -768,24 +767,24 @@ def nvs_part_gen(input_filename=None, output_filename=None, input_size=None, key
     if is_encrypt_data:
         output_keys_file = open("encryption_keys.bin",'wb')
         keys_page_buf = bytearray(b'\xff')*Page.PAGE_PARAMS["max_size"]
-        
+
         if sys.version_info[0] < 3:
             key_bytes = key_input.decode('hex')
         else:
             key_bytes = bytearray()
-            key_bytes = codecs.decode(key_input, 'hex') 
+            key_bytes = codecs.decode(key_input, 'hex')
+
 
-        
         key_len = len(key_bytes)
         keys_page_buf[0:key_len] = key_bytes
-        
+
         crc_data = keys_page_buf[0:key_len]
-        
+
         if sys.version_info[0] < 3:
             crc = zlib.crc32(buffer(crc_data), 0xFFFFFFFF)
         else:
             crc = zlib.crc32(memoryview(crc_data), 0xFFFFFFFF)
-        
+
         struct.pack_into('<I', keys_page_buf, key_len,  crc & 0xFFFFFFFF)
         output_keys_file.write(keys_page_buf)
 
diff --git a/tools/mass_mfg/docs/README.rst b/tools/mass_mfg/docs/README.rst
index f71ec7eac..8d50aedf8 100644
--- a/tools/mass_mfg/docs/README.rst
+++ b/tools/mass_mfg/docs/README.rst
@@ -120,7 +120,7 @@ The mfg\_gen.py utility is using the generated CSV Configuration file and Master
 +------------------------+----------------------------------------------------------------------------------------------+
 |   Arguments            |                                     Description                                              |                                   
 +========================+==============================================================================================+
-| --size PART_SIZE       |  the size of NVS Partition in KB. Eg. 12KB                                                   |
+| --size PART_SIZE       |  Size of NVS Partition in hex (must be multiple of 4096). Eg. 0x1000                         |
 +------------------------+----------------------------------------------------------------------------------------------+
 | --conf CONFIG_FILE     |  the input configuration csv file                                                            |
 +------------------------+----------------------------------------------------------------------------------------------+
@@ -136,7 +136,7 @@ The mfg\_gen.py utility is using the generated CSV Configuration file and Master
 
 **You can use the below command to run this utility with the sample files provided**::
    
-   $ ./mfg_gen.py --size 12KB --conf samples/sample_config.csv --values samples/sample_values.csv --prefix Fan
+   $ ./mfg_gen.py --size 0x3000 --conf samples/sample_config.csv --values samples/sample_values.csv --prefix Fan
 
 
 .. note:: The default numeric value: 1,2,3... of ``fileid`` argument, corresponds to each row having device instance values in master csv values file.
diff --git a/tools/mass_mfg/mfg_gen.py b/tools/mass_mfg/mfg_gen.py
index a9813ae22..ca9c94cb0 100755
--- a/tools/mass_mfg/mfg_gen.py
+++ b/tools/mass_mfg/mfg_gen.py
@@ -22,9 +22,10 @@ import csv
 import argparse
 import shutil
 import distutils.dir_util
-from itertools import zip_longest
 sys.path.insert(0, os.getenv('IDF_PATH') + "/components/nvs_flash/nvs_partition_generator/")
 import nvs_partition_gen
+if not sys.version_info[0] < 3:
+    from itertools import zip_longest
 
 def verify_values_exist(input_values_file, keys_in_values_file):
     """ Verify all keys have corresponding values in values file
@@ -306,7 +307,7 @@ file_identifier=None,output_dir_path=None):
             parser.add_argument("--size",
                                 dest='part_size',
                                 required=True,
-                                help='Size of NVS Partition in KB. Eg. 12KB')
+                                help='Size of NVS Partition in hex (must be multiple of 4096). Eg. 0x1000')
 
             parser.add_argument('--conf',
                                 dest='config_file',