From d7d9e7e465818115cb7d60ad15247f09a8c8c816 Mon Sep 17 00:00:00 2001
From: Md Sadre Alam <quic_mdalam@quicinc.com>
Date: Thu, 9 May 2024 16:38:02 +0530
Subject: [PATCH] tools: Add image encryption script for aes xts mode

Signed-off-by: Md Sadre Alam <quic_mdalam@quicinc.com>
Change-Id: I85048912f86272d061e4f50989c4a97a758aa3db
---
 tools/encrypt_image_aes_xts.py | 73 ++++++++++++++++++++++++++++++++++
 1 file changed, 73 insertions(+)
 create mode 100644 tools/encrypt_image_aes_xts.py

diff --git a/tools/encrypt_image_aes_xts.py b/tools/encrypt_image_aes_xts.py
new file mode 100644
index 0000000000..d844a101b5
--- /dev/null
+++ b/tools/encrypt_image_aes_xts.py
@@ -0,0 +1,73 @@
+/*
+ * Copyright (c) 2024 Qualcomm Innovation Center, Inc. All rights reserved.
+ * SPDX-License-Identifier: GPL-2.0-only
+*/
+
+#!/usr/bin/env python3
+
+"""
+This script encrypts the provided image in blocks of 512.
+Currently only AES XTS is supported. For other modes use
+standard OpenSSL tool.
+
+Usage: python3 encrypt_image.py <input image> <key> <output image name>
+
+"""
+
+import os
+import sys
+from typing import BinaryIO
+
+# Import the pre-requisites
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives.ciphers import Cipher
+from cryptography.hazmat.primitives.ciphers.algorithms import AES
+from cryptography.hazmat.primitives.ciphers.modes import XTS
+
+sectorsize = 512
+sec_dbl = sectorsize * 2
+
+def xts_encrypt(rawfile: BinaryIO, key: bytes, encrypted_file: BinaryIO):
+
+    algorithm = AES(key)
+    sec_cnt = 0
+
+    while True:
+        rawsector = rawfile.read(sectorsize)
+        if not rawsector:
+            break
+
+        # Pad data to 512
+        if len(rawsector) < sectorsize:
+            raw_mul = (sectorsize - len(rawsector))
+            rawsector += (b'\x00' * raw_mul)
+        xts_tweak_bytes = sec_cnt.to_bytes(length=16, byteorder='little')
+        encryptor = Cipher(algorithm=algorithm, mode=XTS(xts_tweak_bytes), backend=default_backend()).encryptor()
+        xts_output = encryptor.update(rawsector) + encryptor.finalize()
+
+        encrypted_file.write(xts_output)
+        sec_cnt += 1
+
+    plain_size = sec_cnt * sectorsize
+    pad_check = (plain_size % (sec_dbl * 4))
+    if (pad_check != 0):
+        pad_length = ((plain_size // sec_dbl) + 4) * sec_dbl
+        pad_length = (pad_length - plain_size)
+        encrypted_file.write(b'\x00' * pad_length)
+
+if __name__ == '__main__':
+
+    if len(sys.argv) != 4:
+      print("Usage: python3 encrypt_image.py <input image> <key> <output image name>")
+      sys.exit(1)
+
+    plain_image = sys.argv[1]
+    cipher_image = sys.argv[3]
+    enc_key =   sys.argv[2]
+
+    with open(enc_key, 'rb') as f:
+        rootfs_key = f.read()
+
+    with open(plain_image, 'rb') as rootfs_img, open(cipher_image, 'wb') as rootfs_enc:
+        xts_encrypt(rootfs_img, rootfs_key, rootfs_enc)
+