ampk: add pc encrypt tool [1/1]

PD#SWPL-177486

Problem:
Need a tool to encrypt secrets for AMPK

Solution:
Add python script to do it

Verify:
./gen_keypair.py -k vendor_priv.pem -u vendor_pub.pem -v
./encrypt_secret.py -k vendor_priv.pem -d $(cat device-pub.txt) -i
dgpk1.bin  -o edgpk1.bin -x edgpk1.txt -v

Change-Id: Ib973da46a3bee0c3aa51ece6c1d59083fe953967
Signed-off-by: Lawrence Mok <lawrence.mok@amlogic.com>
diff --git a/tools/ampk/encrypt_secret.py b/tools/ampk/encrypt_secret.py
new file mode 100755
index 0000000..23dadff
--- /dev/null
+++ b/tools/ampk/encrypt_secret.py
@@ -0,0 +1,138 @@
+#!/usr/bin/env python3
+
+# Copyright Amlogic
+
+import argparse
+import logging
+import os
+import sys
+import random
+
+from Cryptodome.PublicKey import ECC
+from Cryptodome.Cipher import AES
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        description='\n\
+    Encrypts OTP secret for programming.\n\
+    encrypt_secret.py -k vendor_priv.pem -d device_pub.pem -i dgpk1.bin -o edgpk1.bin -x edpgk1.txt')
+
+    parser.add_argument('-k', '--private-key',
+                        metavar='vendor_priv.pem',
+                        required=True,
+                        help='Input vendor private key file',
+                        )
+    parser.add_argument('-d', '--device-public-key',
+                        metavar='device_public_key',
+                        required=True,
+                        help='Input device public key in hex string',
+                        )
+    parser.add_argument('-i', '--input',
+                        metavar='input',
+                        required=True,
+                        help='Input secret file',
+                        )
+    parser.add_argument('-o', '--output',
+                        metavar='output',
+                        required=False,
+                        help='Output encrypted secret file',
+                        )
+    parser.add_argument('-x', '--output-hex',
+                        metavar='output',
+                        required=False,
+                        help='Output encrypted secret file in hex',
+                        )
+
+    parser.add_argument('-v', '--verbose',
+                        action='store_true',
+                        help='Enable verbose log')
+
+    args = parser.parse_args()
+
+    if args.verbose:
+        logging.basicConfig(level=logging.DEBUG)
+    else:
+        logging.basicConfig(level=logging.INFO)
+
+    logging.info('Running {}'.format(sys.argv))
+
+    return args
+
+def key_import(key):
+    """ Check keys are valid, on the curve.
+        returns status, EccKey
+    """
+    if not key:
+        logging.error('key arg missing')
+        return False, None
+    if len(key) < 128:
+        logging.error('key arg should be 64 bytes in hex')
+        logging.error(' {}'.format(len(key)))
+        return False, None
+    kb = bytes.fromhex(key)
+    if len(kb) != 64:
+        logging.error('key arg should be 64 bytes in hex')
+        logging.error(' {}'.format(len(kb)))
+        return False, None
+    xb = kb[:32]
+    yb = kb[32:]
+    x = int.from_bytes(xb, byteorder='big', signed=False)
+    y = int.from_bytes(yb, byteorder='big', signed=False)
+    ecc_key = ECC.construct(curve='p256', point_x=x, point_y=y)
+
+    return True, ecc_key
+
+if __name__ == '__main__':
+    args = parse_args()
+
+    if not os.path.exists(args.private_key):
+        logging.error('Private key {} does not exist'.format(args.private_key))
+        sys.exit(1)
+    device_key_r, device_key = key_import(args.device_public_key)
+    if not device_key_r:
+        logging.error('Device public key is invalid')
+        sys.exit(1)
+    if not os.path.exists(args.input):
+        logging.error('Input {} does not exist'.format(args.input))
+        sys.exit(1)
+    if not args.output and not args.output_hex:
+        logging.error('Missing output arg')
+        sys.exit(1)
+    if os.path.exists(args.output):
+        logging.error('Output {} already exists'.format(args.output))
+        sys.exit(1)
+    if os.path.exists(args.output_hex):
+        logging.error('Output {} already exists'.format(args.output_hex))
+        sys.exit(1)
+
+    # Load keys
+    with open(args.private_key, 'rt') as f:
+        vendor_key = ECC.import_key(f.read())
+
+    # Derive shared secret with ECDH
+    pointZ = device_key.pointQ * vendor_key.d
+    shared_secret = int.to_bytes(int(pointZ.x), length=32, byteorder='big')
+
+    # AES-GCM encrypt
+    # randbytes only in 3.9+
+    iv = bytes([random.randint(0, 255) for _ in range(12)])
+
+    cipher = AES.new(shared_secret, AES.MODE_GCM, nonce=iv, mac_len=16)
+    with open(args.input, 'rb') as f:
+        pt = f.read()
+        ct, tag = cipher.encrypt_and_digest(pt)
+
+    with open(args.output, 'wb') as f:
+        f.write(ct)
+        f.write(iv)
+        f.write(tag)
+    with open(args.output_hex, 'wt') as f:
+        f.write(ct.hex())
+        f.write(iv.hex())
+        f.write(tag.hex())
+        f.write("\n")
+
+    #logging.debug(f'out(enc_efuseinfo)= {ct.hex()}{iv.hex()}{tag.hex()}')
+
+    sys.exit(0)
diff --git a/tools/ampk/gen_keypair.py b/tools/ampk/gen_keypair.py
new file mode 100755
index 0000000..e38fe11
--- /dev/null
+++ b/tools/ampk/gen_keypair.py
@@ -0,0 +1,69 @@
+#!/usr/bin/env python3
+
+# Copyright Amlogic
+
+import argparse
+import logging
+import os
+import sys
+
+from Cryptodome.PublicKey import ECC
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        description='\n\
+    Generates partner an ECC keypair.\n\
+    gen_keypair.py -k PRIV_KEY -u PUB_KEY')
+
+    parser.add_argument('-k', '--private-key',
+                        metavar='vendor_priv.pem',
+                        required=True,
+                        help='Output private key file',
+                        )
+    parser.add_argument('-u', '--public-key',
+                        metavar='vendor_pub.pem',
+                        required=True,
+                        help='output public key file',
+                        )
+
+    parser.add_argument('-v', '--verbose',
+                        action='store_true',
+                        help='Enable verbose log')
+
+    args = parser.parse_args()
+
+    if args.verbose:
+        logging.basicConfig(level=logging.DEBUG)
+    else:
+        logging.basicConfig(level=logging.INFO)
+
+    logging.info('Running {}'.format(sys.argv))
+
+    return args
+
+if __name__ == '__main__':
+    args = parse_args()
+
+    logging.debug(f'{args.private_key=}')
+    logging.debug(f'{args.public_key=}')
+
+    if os.path.exists(args.private_key):
+        logging.error('Private key {} already exists'.format(args.private_key))
+        sys.exit(1)
+    if os.path.exists(args.public_key):
+        logging.error('Public key {} already exists'.format(args.public_key))
+        sys.exit(1)
+
+    key = ECC.generate(curve='p256')
+
+
+    logging.info('Writing private key to {}'.format(args.private_key))
+    with open(args.private_key, 'wt') as f:
+        f.write(key.export_key(format='PEM'))
+
+    logging.info('Writing public key to {}'.format(args.public_key))
+    with open(args.public_key, 'wt') as f:
+        f.write(key.public_key().export_key(format='PEM'))
+
+    sys.exit(0)
diff --git a/tools/ampk/gen_keypair.sh b/tools/ampk/gen_keypair.sh
new file mode 100755
index 0000000..6b1e150
--- /dev/null
+++ b/tools/ampk/gen_keypair.sh
@@ -0,0 +1,22 @@
+#!/bin/bash
+
+if [ $# -ne 2 ]; then
+    echo "Usage: $0 vendor_priv.pem vendor_pub.pem"
+    exit 1
+fi
+
+priv=$1 #user_priv.pem
+pub=$2  #user_pub.pem
+
+if [ -f "$priv" ]; then
+    echo "Error: Private key $priv already exists"
+    exit 1
+fi
+if [ -f "$pub" ]; then
+    echo "Error: Public key $pub already exists"
+    exit 1
+fi
+
+openssl ecparam -name prime256v1 -genkey -noout -out "$priv"
+openssl ec -in "$priv"  -pubout -out "$pub"
+