From 3ff05e03364e27c7132d1e015d46042dbf50de95 Mon Sep 17 00:00:00 2001
From: Martin Olejar <martin.olejar@nxp.com>
Date: Thu, 30 May 2019 11:31:28 +0200
Subject: [PATCH] Add srkgen command into imxim tool

---
 doc/imxim.md            | 26 ++++++++++++++++++++++++++
 imx/img/__main__.py     | 37 ++++++++++++++++++++++++++++++++++++-
 tests/test_cli_imxim.py | 14 ++++++++++++++
 3 files changed, 76 insertions(+), 1 deletion(-)

diff --git a/doc/imxim.md b/doc/imxim.md
index cdf9754..60f5a3d 100644
--- a/doc/imxim.md
+++ b/doc/imxim.md
@@ -25,6 +25,7 @@ Commands:
   create2b  Create new i.MX8M boot image from attached files
   create3a  Create new i.MX8QXP boot image from attached files
   create3b  Create new i.MX8QM boot image from attached files
+  srkgen    SRK Table and Fuses Generator
   dcdfc     DCD file converter (*.bin, *.txt)
   extract   Extract i.MX boot image content
   info      List i.MX boot image content
@@ -246,6 +247,31 @@ Create new i.MX8QM boot image from attached files:
 
 <br>
 
+#### $ imxim srkgen [OPTIONS] [INFILES]
+
+SRK Table and Fuses Generator.
+
+**INFILES** - Input certificates with *.pem extension
+
+##### options:
+* **-t, --table** - Output file name of SRK table (default: srk_table.bin)
+* **-f, --fuses** - Output file name of SRK fuses (default: srk_fuses.bin)
+* **-v, --version** - HAB version (default: 0x40)
+* **-?, --help** - Show help message and exit
+
+##### Example:
+
+```sh
+ $ imxim srkgen -t srk_table.bin -f srk_fuses.bin SRK1_sha256_4096_65537_v3_ca_crt.pem 
+ SRK2_sha256_4096_65537_v3_ca_crt.pem SRK3_sha256_4096_65537_v3_ca_crt.pem SRK4_sha256_4096_65537_v3_ca_crt.pem
+
+ Generated successfully !
+ SRK Table: srk_table.bin
+ SRK Fuses: srk_fuses.bin
+```
+
+<br>
+
 #### $ imxim dcdfc [OPTIONS] OUTFILE [INFILES]
 
 Convert DCD binary blob (*.bin) into readable text file (*.txt) and vice versa.
diff --git a/imx/img/__main__.py b/imx/img/__main__.py
index 732be96..6cc55a0 100644
--- a/imx/img/__main__.py
+++ b/imx/img/__main__.py
@@ -11,7 +11,10 @@
 import yaml
 import click
 
-from imx.img import parse, SegDCD, BootImg2, BootImg3a, BootImg3b, BootImg4, EnumAppType
+from cryptography import x509
+from cryptography.hazmat.backends import default_backend
+
+from imx.img import parse, SegDCD, BootImg2, BootImg3a, BootImg3b, BootImg4, EnumAppType, SrkItem, SrkTable
 from imx import __version__
 
 
@@ -612,6 +615,38 @@ def extract(file, type, offset, step, embedded):
     click.secho(" Image successfully extracted\n Path: %s\n" % out_path)
 
 
+@cli.command(short_help="SRK Table and Fuses Generator")
+@click.argument('infiles', nargs=-1, type=click.Path(exists=True))
+@click.option('-t', '--table', default='srk_table.bin', show_default=True, help="Output file name")
+@click.option('-f', '--fuses', default='srk_fuses.bin', show_default=True, help="Output file name")
+@click.option('-v', '--version', type=UINT, default=0x40, show_default=True, help="HAB version")
+def srkgen(infiles, table, fuses, version):
+    """ SRK table generator """
+    try:
+        srk_table = SrkTable(version)
+
+        for infile in infiles:
+            with open(infile, 'rb') as f:
+                cert = x509.load_pem_x509_certificate(f.read(), default_backend())
+                srk_table.append(SrkItem.from_certificate(cert))
+
+        if table:
+            # Save SRK table
+            with open(table, 'wb') as f:
+                f.write(srk_table.export())
+
+        if fuses:
+            # Save SRK fuses
+            with open(fuses, 'wb') as f:
+                f.write(srk_table.export_fuses())
+
+    except Exception as e:
+        click.echo(str(e) if str(e) else "Unknown Error !")
+        sys.exit(ERROR_CODE)
+
+    click.secho(" Generated successfully\n SRK Table: %s\n SRK Fuses: %s\n" % (table, fuses))
+
+
 @cli.command(short_help="DCD file converter (*.bin, *.txt)")
 @click.argument('outfile', nargs=1, type=click.Path(readable=False))
 @click.argument('infile', nargs=1, type=click.Path(exists=True))
diff --git a/tests/test_cli_imxim.py b/tests/test_cli_imxim.py
index 184afd3..146eecb 100644
--- a/tests/test_cli_imxim.py
+++ b/tests/test_cli_imxim.py
@@ -70,6 +70,20 @@ def test_imxim_info(script_runner):
     assert ret.success
 
 
+@pytest.mark.script_launch_mode('subprocess')
+def test_imxim_srktable(script_runner):
+    # generate SRK table and fuses
+    ret = script_runner.run('imxim',
+                            'srkgen',
+                            '-t', os.path.join(TEMP_DIR, 'srk_table.bin'),
+                            '-f', os.path.join(TEMP_DIR, 'srk_fuses.bin'),
+                            os.path.join(DATA_DIR, 'SRK1_sha256_4096_65537_v3_ca_crt.pem'),
+                            os.path.join(DATA_DIR, 'SRK2_sha256_4096_65537_v3_ca_crt.pem'),
+                            os.path.join(DATA_DIR, 'SRK3_sha256_4096_65537_v3_ca_crt.pem'),
+                            os.path.join(DATA_DIR, 'SRK4_sha256_4096_65537_v3_ca_crt.pem'))
+    assert ret.success
+
+
 @pytest.mark.script_launch_mode('subprocess')
 def test_imxim_dcdfc(script_runner):
     # convert DCD in TXT format to Binary format (default conversion)