add tee task workflow

README.md

@@ -25,6 +25,7 @@ Requirements: poetry
 ```bash
 pip install poetry
 poetry install
+cd ./meca-contracts/src/ganach && npm install
 ```
 
 Requirements: node.js 20.11.1 and npm (tested with 8.5.5)

src/pymeca/contracts_abi/MecaScheduler.json

@@ -198,6 +198,37 @@
         "stateMutability": "view",
         "type": "function"
     },
+    {
+        "inputs": [
+            {
+                "internalType": "bytes32",
+                "name": "taskId",
+                "type": "bytes32"
+            }
+        ],
+        "name": "getTeeTask",
+        "outputs": [
+            {
+                "components": [
+                    {
+                        "internalType": "bytes32",
+                        "name": "encryptedInputHash",
+                        "type": "bytes32"
+                    },
+                    {
+                        "internalType": "bytes32[2]",
+                        "name": "enclavePublicKey",
+                        "type": "bytes32[2]"
+                    }
+                ],
+                "internalType": "struct MecaSchedulerAbstractContract.TeeTask",
+                "name": "",
+                "type": "tuple"
+            }
+        ],
+        "stateMutability": "view",
+        "type": "function"
+    },
     {
         "inputs": [],
         "name": "getTowerContract",
@@ -261,6 +292,42 @@
         "stateMutability": "nonpayable",
         "type": "function"
     },
+    {
+        "inputs": [
+            {
+                "internalType": "bytes32",
+                "name": "taskId",
+                "type": "bytes32"
+            },
+            {
+                "internalType": "bytes32",
+                "name": "encryptedInputHash",
+                "type": "bytes32"
+            }
+        ],
+        "name": "registerTeeTaskEncryptedInput",
+        "outputs": [],
+        "stateMutability": "nonpayable",
+        "type": "function"
+    },
+    {
+        "inputs": [
+            {
+                "internalType": "bytes32",
+                "name": "taskId",
+                "type": "bytes32"
+            },
+            {
+                "internalType": "bytes32[2]",
+                "name": "enclavePublicKey",
+                "type": "bytes32[2]"
+            }
+        ],
+        "name": "registerTeeTaskPubKey",
+        "outputs": [],
+        "stateMutability": "nonpayable",
+        "type": "function"
+    },
     {
         "inputs": [],
         "name": "schedulerFlag",

src/pymeca/host.py

@@ -527,3 +527,37 @@ def wrong_input_hash(
         tx_receipt = self._execute_transaction(transaction)
 
         return tx_receipt.status == 1
+
+    def register_tee_public_key(
+        self,
+        task_id: str,
+        tee_public_key: str
+    ) -> bool:
+        r"""
+        Register the tee public key of the task.
+
+        Args:
+            task_id : Task ID.
+            tee_public_key: the public key of the tee
+
+        Returns:
+            bool: if the register was successful
+        """
+        if not self.is_registered():
+            raise pymeca.utils.MecaError(
+                "The host is not registered"
+            )
+        transaction = self.get_scheduler_contract(
+        ).functions.registerTeeTaskPubKey(
+            taskId=self._bytes_from_hex(task_id),
+            enclavePublicKey=self._bytes_from_hex_public_key(tee_public_key)
+        ).build_transaction({
+            "from": self.account.address,
+            "nonce": self.w3.eth.get_transaction_count(
+                self.account.address
+            )
+        })
+
+        tx_receipt = self._execute_transaction(transaction)
+
+        return tx_receipt.status == 1

src/pymeca/pymeca.py

@@ -163,6 +163,31 @@ def running_task_from_tuple(
     }
 
 
+def tee_task_from_tuple(
+    tee_task_tuple: tuple
+) -> dict:
+    r"""
+    Transform a tee task tuple from the web3 result in
+    a dictionary
+
+    Args:
+        tee_task_tuple : tee task tuple
+
+    Returns:
+        tee_task : tee task dictionary
+        {
+            "encryptedInputHash"
+            "enclavePublicKey"
+        }
+    """
+    return {
+        "encryptedInputHash": "0x" + tee_task_tuple[0].hex().strip(),
+        "enclavePublicKey": "0x" + "".join([
+            x.hex().strip() for x in tee_task_tuple[1]
+        ]),
+    }
+
+
 def host_from_tuple(
     host_tuple: tuple
 ) -> dict:
@@ -437,6 +462,26 @@ def get_running_task(
         ).call()
         return running_task_from_tuple(tuple_running_task)
 
+    def get_tee_task(
+        self,
+        task_id: str
+    ) -> dict:
+        r"""
+        Get the tee task information for a task which
+        is running on the scheduler.
+
+        Args:
+            task_id : The task id on the scheduler
+
+        Returns:
+            dict : The tee task
+        """
+        tuple_tee_task = self.get_scheduler_contract(
+        ).functions.getTeeTask(
+            taskId=self._bytes_from_hex(task_id)
+        ).call()
+        return tee_task_from_tuple(tuple_tee_task)
+
     # host contract functions
     def get_host_register_fee(
         self

src/pymeca/user.py

@@ -112,6 +112,36 @@ def send_task_on_blockchain(
 
         return (tx_receipt.status == 1, task_id)
 
+    def register_tee_task_encrypted_input(
+        self,
+        task_id: str,
+        hash_encrypted_input: bytes
+    ) -> bool:
+        r"""
+        Register TEE task encrypted input.
+
+        Args:
+            task_id : Task ID.
+            encrypted_input : Encrypted input.
+
+        Returns:
+            bool : True if the encrypted input was registered successfully.
+        """
+        transaction = self.get_scheduler_contract(
+        ).functions.registerTeeTaskEncryptedInput(
+            taskId=self._bytes_from_hex(task_id),
+            encryptedInputHash=hash_encrypted_input
+        ).build_transaction({
+            "from": self.account.address,
+            "nonce": self.w3.eth.get_transaction_count(
+                self.account.address
+            )
+        })
+
+        tx_receipt = self._execute_transaction(transaction=transaction)
+
+        return tx_receipt.status == 1
+
     def finish_task(
         self,
         task_id: str,

tests/conftest.py

@@ -425,12 +425,28 @@ def initial_task(accounts):
     )
 
 
+@pytest.fixture(scope="session")
+def initial_tee_task(accounts):
+    return dict(
+        ipfsSha256="0x" + "2" * 64,
+        owner=Account.from_key(
+            accounts["meca_task"]["private_key"]
+        ).address,
+        fee=10,
+        computingType=2,
+        size=1024,
+        enclavePublicKey="0x" + "3" * 128,
+        encryptedInputHash="0x" + "4" * 64,
+    )
+
+
 @pytest.fixture(scope="session")
 def register_web3(
     simple_web3,
     initial_host,
     initial_tower,
-    initial_task
+    initial_task,
+    initial_tee_task
 ):
     r"""
     Return a function to start a register web3 instance for the
@@ -480,6 +496,13 @@ def _register_web3(
             computing_type=initial_task["computingType"],
             size=initial_task["size"]
         )
+        # register the tee task
+        actors["task_developer"].register_task(
+            ipfs_sha256=initial_tee_task["ipfsSha256"],
+            fee=initial_tee_task["fee"],
+            computing_type=initial_tee_task["computingType"],
+            size=initial_tee_task["size"]
+        )
 
         return (w3, server_process, addresses, actors)
 
@@ -514,10 +537,21 @@ def initial_host_task(initial_task):
     )
 
 
+# the initial tee task for the host
+@pytest.fixture(scope="class")
+def initial_host_tee_task(initial_tee_task):
+    return dict(
+        ipfsSha256=initial_tee_task["ipfsSha256"],
+        fee=10,
+        blockTimeout=3,
+    )
+
+
 @pytest.fixture(scope="class")
 def fill_web3(
     register_web3,
-    initial_host_task
+    initial_host_task,
+    initial_host_tee_task
 ):
     r"""
     Using the register environment it creates a flow for running a
@@ -552,6 +586,12 @@ def _fill_web3(
             block_timeout=initial_host_task["blockTimeout"],
             fee=initial_host_task["fee"]
         )
+        # the host add the tee task
+        actors["host"].add_task(
+            ipfs_sha256=initial_host_tee_task["ipfsSha256"],
+            block_timeout=initial_host_tee_task["blockTimeout"],
+            fee=initial_host_tee_task["fee"]
+        )
         # the host request to ge tin the tower list
         actors["host"].register_for_tower(
             tower_address=actors["tower"].account.address