From 60eaadf0a18dd714a9f4cf483445043b59388265 Mon Sep 17 00:00:00 2001
From: Przemyslaw Bida <przemyslaw.bida@nordicsemi.no>
Date: Mon, 12 Aug 2024 09:55:15 +0200
Subject: [PATCH] [tcat] Add timeout while connecting over BLE.

Adding timeout while handling ble connection establishement in TCAT.
---
 tools/tcat_ble_client/ble/ble_stream.py        |  4 ++++
 tools/tcat_ble_client/ble/ble_stream_secure.py | 10 ++++++++--
 tools/tcat_ble_client/cli/base_commands.py     | 10 ++++++----
 3 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/tools/tcat_ble_client/ble/ble_stream.py b/tools/tcat_ble_client/ble/ble_stream.py
index 3e22da8f49aa..4c20b2d89483 100644
--- a/tools/tcat_ble_client/ble/ble_stream.py
+++ b/tools/tcat_ble_client/ble/ble_stream.py
@@ -91,3 +91,7 @@ async def recv(self, bufsize, recv_timeout=0.2):
         self.__receive_buffer = self.__receive_buffer[bufsize:]
         logger.debug(f'retrieved {message}')
         return message
+
+    async def disconnect(self):
+        if self.client.is_connected:
+            await self.client.disconnect()
\ No newline at end of file
diff --git a/tools/tcat_ble_client/ble/ble_stream_secure.py b/tools/tcat_ble_client/ble/ble_stream_secure.py
index ab8858c295dc..de2b42d8ae04 100644
--- a/tools/tcat_ble_client/ble/ble_stream_secure.py
+++ b/tools/tcat_ble_client/ble/ble_stream_secure.py
@@ -34,6 +34,7 @@
 
 from tlv.tlv import TLV
 from tlv.tcat_tlv import TcatTLVType
+from time import time
 import utils
 
 logger = logging.getLogger(__name__)
@@ -60,7 +61,7 @@ def load_cert(self, certfile='', keyfile='', cafile=''):
         if cafile:
             self.ssl_context.load_verify_locations(cafile=cafile)
 
-    async def do_handshake(self):
+    async def do_handshake(self, timeout=30.0):
         is_debug = logger.getEffectiveLevel() <= logging.DEBUG
         self.ssl_object = self.ssl_context.wrap_bio(
             incoming=self.incoming,
@@ -68,7 +69,8 @@ async def do_handshake(self):
             server_side=False,
             server_hostname=None,
         )
-        while True:
+        start = time()
+        while (time() - start) < timeout:
             try:
                 if not is_debug:
                     print('.', end='')
@@ -97,6 +99,10 @@ async def do_handshake(self):
                 if output:
                     self.incoming.write(output)
                 await asyncio.sleep(0.02)
+        else:
+            print('TLS Connection timed out.')
+            return False
+        return True
 
     async def send(self, bytes):
         self.ssl_object.write(bytes)
diff --git a/tools/tcat_ble_client/cli/base_commands.py b/tools/tcat_ble_client/cli/base_commands.py
index 9c57e4ea8d6d..2cc472ed8e40 100644
--- a/tools/tcat_ble_client/cli/base_commands.py
+++ b/tools/tcat_ble_client/cli/base_commands.py
@@ -208,8 +208,10 @@ async def execute_default(self, args, context):
             keyfile=path.join(cert_path, 'commissioner_key.pem'),
             cafile=path.join(cert_path, 'ca_cert.pem'),
         )
-
         print('Setting up secure channel...')
-        await ble_sstream.do_handshake()
-        print('Done')
-        context['ble_sstream'] = ble_sstream
+        if await ble_sstream.do_handshake():
+            print('Done')
+            context['ble_sstream'] = ble_sstream
+        else:
+            print('Secure channel not established.')
+            await ble_stream.disconnect()
\ No newline at end of file