1#!/usr/bin/env python
2#
3# Copyright 2018 Espressif Systems (Shanghai) PTE LTD
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#     http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16#
17
18from __future__ import print_function
19
20import argparse
21import distutils.dir_util
22import os
23import sys
24from io import open
25
26from future.moves.itertools import zip_longest
27
28try:
29    sys.path.insert(0, os.getenv('IDF_PATH') + '/components/nvs_flash/nvs_partition_generator/')
30    import nvs_partition_gen
31except Exception as e:
32    print(e)
33    sys.exit('Please check IDF_PATH')
34
35
36def verify_values_exist(input_values_file, values_file_data, key_count_in_values_file, line_no=1):
37    """ Verify all keys have corresponding values in values file
38    """
39    if len(values_file_data) != key_count_in_values_file:
40        raise SystemExit('\nError: Number of values is not equal to number of keys in file: %s at line No:%s\n'
41                         % (str(input_values_file), str(line_no)))
42
43
44def verify_keys_exist(values_file_keys, config_file_data):
45    """ Verify all keys from config file are present in values file
46    """
47    keys_missing = []
48
49    for line_no, config_data in enumerate(config_file_data,1):
50        if not isinstance(config_data, str):
51            config_data = config_data.encode('utf-8')
52        config_data_line = config_data.strip().split(',')
53        if 'namespace' not in config_data_line:
54            if values_file_keys:
55                if config_data_line[0] == values_file_keys[0]:
56                    del values_file_keys[0]
57                else:
58                    keys_missing.append([config_data_line[0], line_no])
59            else:
60                keys_missing.append([config_data_line[0], line_no])
61
62    if keys_missing:
63        for key, line_no in keys_missing:
64            print('Key:`', str(key), '` at line no:', str(line_no),
65                  ' in config file is not found in values file.')
66        raise SystemExit(1)
67
68
69def verify_datatype_encoding(input_config_file, config_file_data):
70    """ Verify datatype and encodings from config file is valid
71    """
72    valid_encodings = ['string', 'binary', 'hex2bin','u8', 'i8', 'u16', 'u32', 'i32','base64']
73    valid_datatypes = ['file','data','namespace']
74    line_no = 0
75
76    for data in config_file_data:
77        line_no += 1
78        if not isinstance(data, str):
79            data = data.encode('utf-8')
80        line = data.strip().split(',')
81        if line[1] not in valid_datatypes:
82            raise SystemExit('Error: config file: %s has invalid datatype at line no:%s\n'
83                             % (str(input_config_file), str(line_no)))
84        if 'namespace' not in line:
85            if line[2] not in valid_encodings:
86                raise SystemExit('Error: config file: %s has invalid encoding at line no:%s\n'
87                                 % (str(input_config_file), str(line_no)))
88
89
90def verify_file_data_count(cfg_file_data, keys_repeat):
91    """ Verify count of data on each line in config file is equal to 3
92    (as format must be: <key,type and encoding>)
93    """
94    line_no = 0
95
96    for data in cfg_file_data:
97        line_no += 1
98        if not isinstance(data, str):
99            data = data.encode('utf-8')
100        line = data.strip().split(',')
101        if len(line) != 3 and line[0] not in keys_repeat:
102            raise SystemExit('Error: data missing in config file at line no:%s <format needed:key,type,encoding>\n'
103                             % str(line_no))
104
105
106def verify_data_in_file(input_config_file, input_values_file, config_file_keys, keys_in_values_file, keys_repeat):
107    """ Verify count of data on each line in config file is equal to 3 \
108    (as format must be: <key,type and encoding>)
109    Verify datatype and encodings from config file is valid
110    Verify all keys from config file are present in values file and \
111    Verify each key has corresponding value in values file
112    """
113    try:
114        values_file_keys = []
115        values_file_line = None
116
117        # Get keys from values file present in config files
118        values_file_keys = get_keys(keys_in_values_file, config_file_keys)
119
120        with open(input_config_file, 'r', newline='\n') as cfg_file:
121            cfg_file_data = cfg_file.readlines()
122            verify_file_data_count(cfg_file_data, keys_repeat)
123            verify_datatype_encoding(input_config_file, cfg_file_data)
124            verify_keys_exist(values_file_keys, cfg_file_data)
125
126        with open(input_values_file, 'r', newline='\n') as values_file:
127            key_count_in_values_file = len(keys_in_values_file)
128            lineno = 0
129            # Read first keys(header) line
130            values_file_data = values_file.readline()
131            lineno += 1
132            while values_file_data:
133                # Read values line
134                values_file_line = values_file.readline()
135                if not isinstance(values_file_line, str):
136                    values_file_line = values_file_line.encode('utf-8')
137
138                values_file_data = values_file_line.strip().split(',')
139
140                lineno += 1
141                if len(values_file_data) == 1 and '' in values_file_data:
142                    break
143                verify_values_exist(input_values_file, values_file_data, key_count_in_values_file, line_no=lineno)
144
145    except Exception as err:
146        print(err)
147        exit(1)
148
149
150def get_keys(keys_in_values_file, config_file_keys):
151    """ Get keys from values file present in config file
152    """
153    values_file_keys = []
154    for key in keys_in_values_file:
155        if key in config_file_keys:
156            values_file_keys.append(key)
157
158    return values_file_keys
159
160
161def add_config_data_per_namespace(input_config_file):
162    """ Add config data per namespace to `config_data_to_write` list
163    """
164    config_data_to_write = []
165    config_data_per_namespace = []
166
167    with open(input_config_file, 'r', newline='\n') as cfg_file:
168        config_data = cfg_file.readlines()
169
170    # `config_data_per_namespace` is added to `config_data_to_write` list after reading next namespace
171    for data in config_data:
172        if not isinstance(data, str):
173            data = data.encode('utf-8')
174        cfg_data = data.strip().split(',')
175        if 'REPEAT' in cfg_data:
176            cfg_data.remove('REPEAT')
177        if 'namespace' in cfg_data:
178            if config_data_per_namespace:
179                config_data_to_write.append(config_data_per_namespace)
180                config_data_per_namespace = []
181                config_data_per_namespace.append(cfg_data)
182            else:
183                config_data_per_namespace.append(cfg_data)
184        else:
185            config_data_per_namespace.append(cfg_data)
186
187    # `config_data_per_namespace` is added to `config_data_to_write` list as EOF is reached
188    if (not config_data_to_write) or (config_data_to_write and config_data_per_namespace):
189        config_data_to_write.append(config_data_per_namespace)
190
191    return config_data_to_write
192
193
194def get_fileid_val(file_identifier, keys_in_config_file, keys_in_values_file,
195                   values_data_line, key_value_data, fileid_value):
196    """ Get file identifier value
197    """
198    file_id_found = False
199
200    for key in key_value_data:
201        if file_identifier and not file_id_found and file_identifier in key:
202            fileid_value = key[1]
203            file_id_found = True
204
205    if not file_id_found:
206        fileid_value = str(int(fileid_value) + 1)
207
208    return fileid_value
209
210
211def add_data_to_file(config_data_to_write, key_value_pair, output_csv_file):
212    """ Add data to csv target file
213    """
214    header = ['key', 'type', 'encoding', 'value']
215    data_to_write = []
216    newline = u'\n'
217
218    target_csv_file = open(output_csv_file, 'w', newline=None)
219
220    line_to_write = u','.join(header)
221    target_csv_file.write(line_to_write)
222    target_csv_file.write(newline)
223    for namespace_config_data in config_data_to_write:
224        for data in namespace_config_data:
225            data_to_write = data[:]
226            if 'namespace' in data:
227                data_to_write.append('')
228                line_to_write = u','.join(data_to_write)
229                target_csv_file.write(line_to_write)
230                target_csv_file.write(newline)
231            else:
232                key = data[0]
233                while key not in key_value_pair[0]:
234                    del key_value_pair[0]
235                if key in key_value_pair[0]:
236                    value = key_value_pair[0][1]
237                    data_to_write.append(value)
238                    del key_value_pair[0]
239                    line_to_write = u','.join(data_to_write)
240                    target_csv_file.write(line_to_write)
241                    target_csv_file.write(newline)
242
243    # Set index to start of file
244    target_csv_file.seek(0)
245    target_csv_file.close()
246
247
248def create_dir(filetype, output_dir_path):
249    """ Create new directory(if doesn't exist) to store file generated
250    """
251    output_target_dir = os.path.join(output_dir_path,filetype,'')
252    if not os.path.isdir(output_target_dir):
253        distutils.dir_util.mkpath(output_target_dir)
254
255    return output_target_dir
256
257
258def set_repeat_value(total_keys_repeat, keys, csv_file, target_filename):
259    key_val_pair = []
260    key_repeated = []
261    line = None
262    newline = u'\n'
263    with open(csv_file, 'r', newline=None) as read_from, open(target_filename,'w', newline=None) as write_to:
264        headers = read_from.readline()
265        values = read_from.readline()
266        write_to.write(headers)
267        write_to.write(values)
268        if not isinstance(values, str):
269            values = values.encode('utf-8')
270        values = values.strip().split(',')
271        total_keys_values = list(zip_longest(keys, values))
272
273        # read new data, add value if key has repeat tag, write to new file
274        line = read_from.readline()
275        if not isinstance(line, str):
276            line = line.encode('utf-8')
277        row = line.strip().split(',')
278        while row:
279            index = -1
280            key_val_new = list(zip_longest(keys, row))
281            key_val_pair = total_keys_values[:]
282            key_repeated = total_keys_repeat[:]
283            while key_val_new and key_repeated:
284                index = index + 1
285                #  if key has repeat tag, get its corresponding value, write to file
286                if key_val_new[0][0] == key_repeated[0]:
287                    val = key_val_pair[0][1]
288                    row[index] = val
289                    del key_repeated[0]
290                del key_val_new[0]
291                del key_val_pair[0]
292
293            line_to_write = u','.join(row)
294            write_to.write(line_to_write)
295            write_to.write(newline)
296
297            # Read next line
298            line = read_from.readline()
299            if not isinstance(line, str):
300                line = line.encode('utf-8')
301            row = line.strip().split(',')
302            if len(row) == 1 and '' in row:
303                break
304
305    return target_filename
306
307
308def create_intermediate_csv(args, keys_in_config_file, keys_in_values_file, keys_repeat, is_encr=False):
309    file_identifier_value = '0'
310    csv_str = 'csv'
311    bin_str = 'bin'
312    line = None
313    set_output_keyfile = False
314
315    # Add config data per namespace to `config_data_to_write` list
316    config_data_to_write = add_config_data_per_namespace(args.conf)
317
318    try:
319        with open(args.values, 'r', newline=None) as csv_values_file:
320            # first line must be keys in file
321            line = csv_values_file.readline()
322            if not isinstance(line, str):
323                line = line.encode('utf-8')
324            keys = line.strip().split(',')
325
326        filename, file_ext = os.path.splitext(args.values)
327        target_filename = filename + '_created' + file_ext
328        if keys_repeat:
329            target_values_file = set_repeat_value(keys_repeat, keys, args.values, target_filename)
330        else:
331            target_values_file = args.values
332
333        csv_values_file = open(target_values_file, 'r', newline=None)
334
335        # Read header line
336        csv_values_file.readline()
337
338        # Create new directory(if doesn't exist) to store csv file generated
339        output_csv_target_dir = create_dir(csv_str, args.outdir)
340        # Create new directory(if doesn't exist) to store bin file generated
341        output_bin_target_dir = create_dir(bin_str, args.outdir)
342        if args.keygen:
343            set_output_keyfile = True
344
345        line = csv_values_file.readline()
346        if not isinstance(line, str):
347            line = line.encode('utf-8')
348        values_data_line = line.strip().split(',')
349
350        while values_data_line:
351            key_value_data = list(zip_longest(keys_in_values_file, values_data_line))
352
353            # Get file identifier value from values file
354            file_identifier_value = get_fileid_val(args.fileid, keys_in_config_file,
355                                                   keys_in_values_file, values_data_line, key_value_data,
356                                                   file_identifier_value)
357
358            key_value_pair = key_value_data[:]
359
360            # Verify if output csv file does not exist
361            csv_filename = args.prefix + '-' + file_identifier_value + '.' + csv_str
362            output_csv_file = output_csv_target_dir + csv_filename
363            if os.path.isfile(output_csv_file):
364                raise SystemExit('Target csv file: %s already exists.`' % output_csv_file)
365
366            # Add values corresponding to each key to csv intermediate file
367            add_data_to_file(config_data_to_write, key_value_pair, output_csv_file)
368            print('\nCreated CSV file: ===>', output_csv_file)
369
370            # Verify if output bin file does not exist
371            bin_filename = args.prefix + '-' + file_identifier_value + '.' + bin_str
372            output_bin_file = output_bin_target_dir + bin_filename
373            if os.path.isfile(output_bin_file):
374                raise SystemExit('Target binary file: %s already exists.`' % output_bin_file)
375
376            args.input = output_csv_file
377            args.output = os.path.join(bin_str, bin_filename)
378            if set_output_keyfile:
379                args.keyfile = 'keys-' + args.prefix + '-' + file_identifier_value
380
381            if is_encr:
382                nvs_partition_gen.encrypt(args)
383            else:
384                nvs_partition_gen.generate(args)
385
386            # Read next line
387            line = csv_values_file.readline()
388            if not isinstance(line, str):
389                line = line.encode('utf-8')
390            values_data_line = line.strip().split(',')
391            if len(values_data_line) == 1 and '' in values_data_line:
392                break
393
394        print('\nFiles generated in %s ...' % args.outdir)
395
396    except Exception as e:
397        print(e)
398        exit(1)
399    finally:
400        csv_values_file.close()
401
402
403def verify_empty_lines_exist(file_name, input_file_data):
404    for data in input_file_data:
405        if not isinstance(data, str):
406            data = data.encode('utf-8')
407        cfg_data = data.strip().split(',')
408
409        if len(cfg_data) == 1 and '' in cfg_data:
410            raise SystemExit('Error: file: %s cannot have empty lines. ' % file_name)
411
412
413def verify_file_format(args):
414    keys_in_config_file = []
415    keys_in_values_file = []
416    keys_repeat = []
417    file_data_keys = None
418
419    # Verify config file is not empty
420    if os.stat(args.conf).st_size == 0:
421        raise SystemExit('Error: config file: %s is empty.' % args.conf)
422
423    # Verify values file is not empty
424    if os.stat(args.values).st_size == 0:
425        raise SystemExit('Error: values file: %s is empty.' % args.values)
426
427    # Verify config file does not have empty lines
428    with open(args.conf, 'r', newline='\n') as csv_config_file:
429        try:
430            file_data = csv_config_file.readlines()
431            verify_empty_lines_exist(args.conf, file_data)
432
433            csv_config_file.seek(0)
434            # Extract keys from config file
435            for data in file_data:
436                if not isinstance(data, str):
437                    data = data.encode('utf-8')
438                line_data = data.strip().split(',')
439                if 'namespace' not in line_data:
440                    keys_in_config_file.append(line_data[0])
441                if 'REPEAT' in line_data:
442                    keys_repeat.append(line_data[0])
443        except Exception as e:
444            print(e)
445
446    # Verify values file does not have empty lines
447    with open(args.values, 'r', newline='\n') as csv_values_file:
448        try:
449            # Extract keys from values file (first line of file)
450            file_data = [csv_values_file.readline()]
451
452            file_data_keys = file_data[0]
453            if not isinstance(file_data_keys, str):
454                file_data_keys = file_data_keys.encode('utf-8')
455
456            keys_in_values_file = file_data_keys.strip().split(',')
457
458            while file_data:
459                verify_empty_lines_exist(args.values, file_data)
460                file_data = [csv_values_file.readline()]
461                if '' in file_data:
462                    break
463
464        except Exception as e:
465            print(e)
466
467    # Verify file identifier exists in values file
468    if args.fileid:
469        if args.fileid not in keys_in_values_file:
470            raise SystemExit('Error: target_file_identifier: %s does not exist in values file.\n' % args.fileid)
471    else:
472        args.fileid = 1
473
474    return keys_in_config_file, keys_in_values_file, keys_repeat
475
476
477def generate(args):
478    keys_in_config_file = []
479    keys_in_values_file = []
480    keys_repeat = []
481    encryption_enabled = False
482
483    args.outdir = os.path.join(args.outdir, '')
484    # Verify input config and values file format
485    keys_in_config_file, keys_in_values_file, keys_repeat = verify_file_format(args)
486
487    # Verify data in the input_config_file and input_values_file
488    verify_data_in_file(args.conf, args.values, keys_in_config_file,
489                        keys_in_values_file, keys_repeat)
490
491    if (args.keygen or args.inputkey):
492        encryption_enabled = True
493        print('\nGenerating encrypted NVS binary images...')
494
495    # Create intermediate csv file
496    create_intermediate_csv(args, keys_in_config_file, keys_in_values_file,
497                            keys_repeat, is_encr=encryption_enabled)
498
499
500def generate_key(args):
501    nvs_partition_gen.generate_key(args)
502
503
504def main():
505    try:
506        parser = argparse.ArgumentParser(description='\nESP Manufacturing Utility', formatter_class=argparse.RawTextHelpFormatter)
507        subparser = parser.add_subparsers(title='Commands',
508                                          dest='command',
509                                          help='\nRun mfg_gen.py {command} -h for additional help\n\n')
510
511        parser_gen = subparser.add_parser('generate',
512                                          help='Generate NVS partition',
513                                          formatter_class=argparse.RawTextHelpFormatter)
514        parser_gen.set_defaults(func=generate)
515        parser_gen.add_argument('conf',
516                                default=None,
517                                help='Path to configuration csv file to parse')
518        parser_gen.add_argument('values',
519                                default=None,
520                                help='Path to values csv file to parse')
521        parser_gen.add_argument('prefix',
522                                default=None,
523                                help='Unique name for each output filename prefix')
524        parser_gen.add_argument('size',
525                                default=None,
526                                help='Size of NVS partition in bytes\
527                                    \n(must be multiple of 4096)')
528        parser_gen.add_argument('--fileid',
529                                default=None,
530                                help='''Unique file identifier(any key in values file) \
531                                    \nfor each filename suffix (Default: numeric value(1,2,3...)''')
532        parser_gen.add_argument('--version',
533                                choices=[1, 2],
534                                default=2,
535                                type=int,
536                                help='''Set multipage blob version.\
537                                    \nVersion 1 - Multipage blob support disabled.\
538                                    \nVersion 2 - Multipage blob support enabled.\
539                                    \nDefault: Version 2 ''')
540        parser_gen.add_argument('--keygen',
541                                action='store_true',
542                                default=False,
543                                help='Generates key for encrypting NVS partition')
544        parser_gen.add_argument('--keyfile',
545                                default=None,
546                                help=argparse.SUPPRESS)
547        parser_gen.add_argument('--inputkey',
548                                default=None,
549                                help='File having key for encrypting NVS partition')
550        parser_gen.add_argument('--outdir',
551                                default=os.getcwd(),
552                                help='Output directory to store files created\
553                                    \n(Default: current directory)')
554        parser_gen.add_argument('--input',
555                                default=None,
556                                help=argparse.SUPPRESS)
557        parser_gen.add_argument('--output',
558                                default=None,
559                                help=argparse.SUPPRESS)
560        parser_gen_key = subparser.add_parser('generate-key',
561                                              help='Generate keys for encryption',
562                                              formatter_class=argparse.RawTextHelpFormatter)
563        parser_gen_key.set_defaults(func=generate_key)
564        parser_gen_key.add_argument('--keyfile',
565                                    default=None,
566                                    help='Path to output encryption keys file')
567        parser_gen_key.add_argument('--outdir',
568                                    default=os.getcwd(),
569                                    help='Output directory to store files created.\
570                                        \n(Default: current directory)')
571
572        args = parser.parse_args()
573        args.func(args)
574
575    except ValueError as err:
576        print(err)
577    except Exception as e:
578        print(e)
579
580
581if __name__ == '__main__':
582    main()
583