lib/unit: Add Unit ID ECHD & AES function.

Signed-off-by: tinyu <[email protected]>

Author: TinyuZhao

docs/en/units/id.rst

@@ -222,4 +222,40 @@ Methods
 
         |set_certificate_signing_request.png|
 
+.. method:: IDUnit.aes_ecb_encrypt(position, data) -> bytearray
 
+    Performs AES-ECB mode encryption on data. The AES key is stored in slot 9 of the chip.
+
+    :param int position: Position for encryption operation (used after left shift by 6)
+    :param bytearray data: Data to be encrypted, must be 16 bytes
+
+    - Return: ``bytearray``: Returns 16 bytes of encrypted data
+
+
+.. method:: IDUnit.aes_ecb_decrypt(position, data) -> bytearray
+
+    Performs AES-ECB mode decryption on data. The AES key is stored in slot 9 of the chip.
+
+    :param int position: Position for decryption operation (used after left shift by 6)
+    :param bytearray data: Data to be decrypted, must be 16 bytes
+
+    - Return: ``bytearray``: Returns 16 bytes of decrypted data
+
+
+.. method:: IDUnit.aes_gfm_encrypt(data) -> bytearray
+
+    Performs Galois Field Multiplication (GFM) encryption operation in AES-GCM mode.
+
+    :param bytearray data: Data to be encrypted, must be 16 bytes
+
+    - Return: ``bytearray``: Returns 16 bytes of encrypted data
+
+.. method:: IDUnit.ecdh_stored_key(slot_num, mode, external_public_key=None) -> bytearray
+
+    Performs an ECDH key exchange operation using stored keys.
+
+    :param int slot_num: The key slot number
+    :param int mode: ECDH operation mode
+    :param bytearray external_public_key: External public key, must be 64 bytes. If not provided, a new public key will be generated
+
+    - Return: ``bytearray``: Returns 32 bytes shared secret when mode is 0x0C or 0x08;

m5stack/libs/driver/atecc608b_tngtls/atecc.py

@@ -13,12 +13,12 @@
 #
 # This library is distributed in the hope that it will be useful,
 # but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 # Lesser General Public License for more details.
 #
 # You should have received a copy of the GNU Lesser General Public
 # License along with this library; if not, write to the Free Software
-# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
+# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
 
 
 import time
@@ -409,7 +409,9 @@ def _read(self, zone: int, address: int, buffer: bytearray) -> None:
         time.sleep(0.001)
         self.idle()
 
-    def _send_command(self, opcode: int, param_1: int, param_2: int = 0x00, data="") -> None:
+    def _send_command(
+        self, opcode: int, param_1: int, param_2: int = 0x00, data: bytearray = bytearray(b"")
+    ) -> None:
         #! Sends a security command packet over i2c.
         #! assembling command packet
         command_packet = bytearray(8 + len(data))
@@ -451,7 +453,7 @@ def _get_response(self, buf, length: int = None, retries: int = 20) -> int:
             raise RuntimeError("Failed to read data from chip")
         if self._debug:
             print("\tReceived: ", [hex(i) for i in response])
-        crc = response[-2] | (response[-1] << 8)
+        crc = (response[-1] << 8) | response[-2]
         crc2 = self._at_crc(response[0:-2])
         if crc != crc2:
             raise RuntimeError("CRC Mismatch")

m5stack/libs/unit/id.py

@@ -3,10 +3,11 @@
 # SPDX-License-Identifier: MIT
 
 from machine import I2C
-from .pahub import PAHUBUnit
-from .unit_helper import UnitError
+from unit.pahub import PAHUBUnit
+from unit.unit_helper import UnitError
 import binascii
 import struct
+import time
 
 from driver.atecc608b_tngtls.atecc import ATECC
 from driver.atecc608b_tngtls.atecc_cert_util import CSR
@@ -24,7 +25,6 @@ def __init__(self, i2c: I2C | PAHUBUnit, address: int = _I2C_ADDR) -> None:
         self._i2c_addr = address
         if address not in self._i2c.scan():
             raise UnitError("ID unit maybe not found in Grove")
-
         self.atecc = ATECC(self._i2c, self._i2c_addr)
 
     def get_revision_number(self) -> int:
@@ -55,14 +55,34 @@ def uniform(self, min: float = 0, max: float = 0) -> float:
         min, max = (max, min) if min > max else (min, max)
         return min + (max - min) * self.random()
 
-    def get_generate_key(self, slot_num: int, private_key: bool = False) -> bytearray:
+    def get_generate_key(self, slot_num: int, private_key: bool = False):
         #! Generates a private or public key.
         assert 0 <= slot_num <= 4, "Provided slot must be between 0 and 4."
         # Create a new key
         key = bytearray(64)
         self.atecc.gen_key(key, slot_num, private_key)
         return key
 
+    def ecdh_stored_key(self, slot_num: int, mode: int, external_public_key: bytearray = None):
+        if external_public_key is None:
+            external_public_key = bytearray(64)
+            self.atecc.gen_key(external_public_key, 0, private_key=False)
+
+        if len(external_public_key) != 64:
+            raise ValueError("External public key must be exactly 64 bytes.")
+
+        self.atecc._send_command(
+            opcode=0x43, param_1=mode, param_2=slot_num, data=external_public_key
+        )
+
+        if mode in (0x0C, 0x08):
+            time.sleep(1)
+            res = bytearray(32)
+            self.atecc._get_response(res)
+            return res
+        else:
+            return None
+
     def get_ecdsa_sign(self, slot: int, message: str | list | bytearray) -> bytearray | None:
         #! Generates and returns a signature using the ECDSA algorithm.
         if len(message) == 32:
@@ -120,3 +140,66 @@ def set_certificate_signing_request(
         )
         with open(file_path, "w+") as pem_file:
             pem_file.write(pem_content)
+
+    # ref: ATECC608A-TFLXTLS 4.4.1 Table 4-5
+    def read_data_zone(self, data_zone: int):
+        self.atecc._send_command(
+            opcode=0x02,
+            param_1=0x00,
+            param_2=data_zone,
+        )
+        time.sleep(0.1)
+        temp = bytearray(4)
+        self.atecc._get_response(temp)
+        return temp
+
+    def write_data_zone(self, data_zone: int, data: bytearray):
+        self.atecc._send_command(
+            opcode=0x12,
+            param_1=0x02,
+            param_2=data_zone,
+            data=data,
+        )
+
+    def aes_ecb_encrypt(self, position: int, data: bytearray):
+        if len(data) != 16:
+            raise ValueError("Data must be 16 bytes")
+        self.atecc._send_command(
+            opcode=0x51,
+            param_1=(position << 6),
+            param_2=0x09,  # ATECC608B AES Key only stored in slot 9
+            data=data,
+        )
+        time.sleep(1)
+        res = bytearray(16)
+        self.atecc._get_response(res)
+        return res
+
+    def aes_ecb_decrypt(self, position: int, data: bytearray):
+        if len(data) != 16:
+            raise ValueError("Data must be 16 bytes")
+        self.atecc._send_command(
+            opcode=0x51,
+            param_1=(position << 6) + 1,
+            param_2=0x09,  # ATECC608B AES Key only stored in slot 9
+            data=data,
+        )
+        time.sleep(1)
+        res = bytearray(16)
+        self.atecc._get_response(res)
+        return res
+
+    def get_h_field(self):
+        zero_block = bytearray([0x00] * 16)
+        return self.aes_ecb_encrypt(0x00, zero_block)
+
+    def aes_gfm_encrypt(self, data: bytearray):
+        if len(data) != 16:
+            raise ValueError("Data must be 16 bytes")
+        data = self.get_h_field() + data
+        time.sleep(1)
+        self.atecc._send_command(opcode=0x51, param_1=0x03, param_2=0x00, data=data)
+        time.sleep(1)
+        res = bytearray(16)
+        self.atecc._get_response(res)
+        return res