tests: cover valid and invalid TOTP authentication flows

Author: root

vertica_python/tests/integration_tests/test_authentication.py

@@ -133,51 +133,115 @@ def test_oauth_access_token(self):
 from io import StringIO
 import sys
 
-# Example static base32 secret
-TOTP_SECRET = "O5D7DQICJTM34AZROWHSAO4O53ELRJN3"
 
 @pytest.fixture
-def totp_code():
-    return pyotp.TOTP(TOTP_SECRET).now()
-
-def test_totp_connection_string(totp_code):
-    """Test: TOTP passed via connection info (like Go test WithTOTPInConnStr)"""
-    try:
-        conn_info = {
-            'user': "dbadmin",
-            'host': "localhost",
-            'port': 5433,
-            'database': "dbadmin",
-            'totp': totp_code,
-            'ssl': False,
-        }
-
-        with vertica_python.connect(**conn_info) as conn:
-            cur = conn.cursor()
-            cur.execute("SELECT version()")
-            result = cur.fetchone()
-            print("[ConnStr] Connected to Vertica Version:", result[0])
-    except Exception as e:
-        pytest.skip(f"Skipping test_totp_connection_string due to connection error: {e}")
-
-def test_totp_from_stdin(monkeypatch, totp_code):
-    """Test: Simulate user entering TOTP at runtime (like Go test WithTOTPFromStdin)"""
-    try:
-        monkeypatch.setattr('sys.stdin', StringIO(totp_code + "\n"))
-
-        conn_info = {
-            'user': "dbadmin",
-            'host': "localhost",
-            'port': 5433,
-            'database': "dbadmin",
-            # Note: TOTP not passed here, simulating prompt from driver
-            'ssl': False,
-        }
-
-        with vertica_python.connect(**conn_info) as conn:
-            cur = conn.cursor()
-            cur.execute("SELECT version()")
-            result = cur.fetchone()
-            print("[Stdin] Connected to Vertica Version:", result[0])
-    except Exception as e:
-        pytest.skip(f"Skipping test_totp_from_stdin due to connection error: {e}")
+# Positive TOTP Test (Like SHA512 format)
+def test_TOTP(self):
+    with self._connect() as conn:
+        cur = conn.cursor()
+
+        cur.execute("DROP USER IF EXISTS totp_user")
+        cur.execute("DROP AUTHENTICATION IF EXISTS totp_auth CASCADE")
+
+        try:
+            # Create user with MFA
+            cur.execute("CREATE USER totp_user IDENTIFIED BY 'password' ENFORCEMFA")
+
+            # Grant authentication
+            # Note: METHOD is 'trusted' or 'password' depending on how MFA is enforced in Vertica
+            cur.execute("CREATE AUTHENTICATION totp_auth METHOD 'password' HOST '0.0.0.0/0'")
+            cur.execute("GRANT AUTHENTICATION totp_auth TO totp_user")
+
+            # Generate TOTP
+            import pyotp
+            TOTP_SECRET = "O5D7DQICJTM34AZROWHSAO4O53ELRJN3"
+            totp_code = pyotp.TOTP(TOTP_SECRET).now()
+
+            # Set connection info
+            self._conn_info['user'] = 'totp_user'
+            self._conn_info['password'] = 'password'
+            self._conn_info['totp'] = totp_code
+
+            # Try connection
+            with self._connect() as totp_conn:
+                c = totp_conn.cursor()
+                c.execute("SELECT 1")
+                res = c.fetchone()
+                self.assertEqual(res[0], 1)
+
+        finally:
+            cur.execute("DROP USER IF EXISTS totp_user")
+            cur.execute("DROP AUTHENTICATION IF EXISTS totp_auth CASCADE")
+
+# Negative Test: Missing TOTP
+def test_TOTP_missing_code(self):
+    with self._connect() as conn:
+        cur = conn.cursor()
+
+        cur.execute("DROP USER IF EXISTS totp_user")
+        cur.execute("DROP AUTHENTICATION IF EXISTS totp_auth CASCADE")
+
+        try:
+            cur.execute("CREATE USER totp_user IDENTIFIED BY 'password' ENFORCEMFA")
+            cur.execute("CREATE AUTHENTICATION totp_auth METHOD 'password' HOST '0.0.0.0/0'")
+            cur.execute("GRANT AUTHENTICATION totp_auth TO totp_user")
+
+            self._conn_info['user'] = 'totp_user'
+            self._conn_info['password'] = 'password'
+            self._conn_info.pop('totp', None)  # No TOTP
+
+            err_msg = "TOTP was requested but not provided"
+            self.assertConnectionFail(err_msg=err_msg)
+
+        finally:
+            cur.execute("DROP USER IF EXISTS totp_user")
+            cur.execute("DROP AUTHENTICATION IF EXISTS totp_auth CASCADE")
+
+# Negative Test: Invalid TOTP Format
+def test_TOTP_invalid_format(self):
+    with self._connect() as conn:
+        cur = conn.cursor()
+
+        cur.execute("DROP USER IF EXISTS totp_user")
+        cur.execute("DROP AUTHENTICATION IF EXISTS totp_auth CASCADE")
+
+        try:
+            cur.execute("CREATE USER totp_user IDENTIFIED BY 'password' ENFORCEMFA")
+            cur.execute("CREATE AUTHENTICATION totp_auth METHOD 'password' HOST '0.0.0.0/0'")
+            cur.execute("GRANT AUTHENTICATION totp_auth TO totp_user")
+
+            self._conn_info['user'] = 'totp_user'
+            self._conn_info['password'] = 'password'
+            self._conn_info['totp'] = "123"   # Invalid
+
+            err_msg = "Invalid TOTP format"
+            self.assertConnectionFail(err_msg=err_msg)
+
+        finally:
+            cur.execute("DROP USER IF EXISTS totp_user")
+            cur.execute("DROP AUTHENTICATION IF EXISTS totp_auth CASCADE")
+
+# Negative Test: Wrong TOTP (Valid format, wrong value)
+def test_TOTP_wrong_code(self):
+    with self._connect() as conn:
+        cur = conn.cursor()
+
+        cur.execute("DROP USER IF EXISTS totp_user")
+        cur.execute("DROP AUTHENTICATION IF EXISTS totp_auth CASCADE")
+
+        try:
+            cur.execute("CREATE USER totp_user IDENTIFIED BY 'password' ENFORCEMFA")
+            cur.execute("CREATE AUTHENTICATION totp_auth METHOD 'password' HOST '0.0.0.0/0'")
+            cur.execute("GRANT AUTHENTICATION totp_auth TO totp_user")
+
+            self._conn_info['user'] = 'totp_user'
+            self._conn_info['password'] = 'password'
+            self._conn_info['totp'] = "999999"   # Wrong OTP
+
+            err_msg = "Invalid TOTP"
+            self.assertConnectionFail(err_msg=err_msg)
+
+        finally:
+            cur.execute("DROP USER IF EXISTS totp_user")
+            cur.execute("DROP AUTHENTICATION IF EXISTS totp_auth CASCADE")
+