Added a 5-minute timeout for TOTP input during authentication

Author: root

vertica_python/vertica/connection.py

@@ -45,6 +45,10 @@
 import uuid
 import warnings
 import re
+import time
+import signal
+import select
+import sys
 from collections import deque
 from struct import unpack
 
@@ -990,25 +994,34 @@ def send_startup(totp_value=None):
                     # ✅ If TOTP not provided initially, allow 3 retry attempts
                     if not totp:
                         max_attempts = 3
+                        timeout_seconds = 300  # 5 minutes timeout
                         for attempt in range(max_attempts):
                             try:
-                                totp_input = input(f"🔐 Enter TOTP (attempt {attempt+1}/{max_attempts}): ").strip()
-                                if not totp_input:
-                                    print("⚠️  TOTP cannot be empty.")
-                                    continue
-                                totp = totp_input
-                                self.close_socket()
-                                self.socket = self.establish_socket_connection(self.address_list)
-                                self._logger.info(f"🚀 Retrying with TOTP: '{totp}'")
-                                # ✅ Re-init required attributes
-                                self.backend_pid = 0
-                                self.backend_key = 0
-                                self.transaction_status = None
-                                self.session_id = None
-                                self._logger.debug("✅ Startup message sent with TOTP.")
-                                # Send new startup message with updated TOTP
-                                send_startup(totp_value=totp)
-                                break
+                                print(f"🔐 Enter TOTP (attempt {attempt+1}/{max_attempts}): ", end="", flush=True)
+                                ready, _, _ = select.select([sys.stdin], [], [], timeout_seconds)
+                                if ready:
+                                    totp_input = sys.stdin.readline().strip()
+                                    if not totp_input:
+                                        print("⚠️  TOTP cannot be empty.")
+                                        continue
+                                    totp = totp_input
+                                    self.close_socket()
+                                    self.socket = self.establish_socket_connection(self.address_list)
+                                    self._logger.info(f"🚀 Retrying with TOTP: '{totp}'")
+                                    # ✅ Re-init required attributes
+                                    self.backend_pid = 0
+                                    self.backend_key = 0
+                                    self.transaction_status = None
+                                    self.session_id = None
+                                    self._logger.debug("✅ Startup message sent with TOTP.")
+                                    # Send new startup message with updated TOTP
+                                    send_startup(totp_value=totp)
+                                    break
+                                else:
+                                    print("⏰ Session timed out. No TOTP entered within time limit.")
+                                    self._logger.error("Session timeout: No TOTP entered within time limit.")
+                                    self.close_socket()
+                                    raise errors.ConnectionError("Session timeout: No TOTP entered within time limit.")
                             except (KeyboardInterrupt, EOFError):
                                 raise errors.ConnectionError("TOTP input cancelled.")
                         else: