Skip to content

Commit 1395e62

Browse files
sharmagotsharmagot
authored
Add strict TOTP validation to block invalid authentication requests (#578)
1 parent 494fe29 commit 1395e62

File tree

2 files changed

+298
-129
lines changed

2 files changed

+298
-129
lines changed

‎vertica_python/tests/integration_tests/test_authentication.py‎

Lines changed: 235 additions & 114 deletions
Original file line numberDiff line numberDiff line change
@@ -123,121 +123,242 @@ def test_oauth_access_token(self):
123123
cur.execute("SELECT authentication_method FROM sessions WHERE session_id=(SELECT current_session())")
124124
res = cur.fetchone()
125125
self.assertEqual(res[0], 'OAuth')
126-
# -------------------------------
127-
# TOTP Authentication Test for Vertica-Python Driver
128-
# -------------------------------
129-
import os
130-
import pyotp
131-
from io import StringIO
132-
import sys
133-
134-
135-
# Positive TOTP Test (Like SHA512 format)
136-
def totp_positive_scenario(self):
137-
with self._connect() as conn:
138-
cur = conn.cursor()
139-
140-
cur.execute("DROP USER IF EXISTS totp_user")
141-
cur.execute("DROP AUTHENTICATION IF EXISTS totp_auth CASCADE")
142-
143-
try:
144-
# Create user with MFA
145-
cur.execute("CREATE USER totp_user IDENTIFIED BY 'password' ENFORCEMFA")
146-
147-
# Grant authentication
148-
# Note: METHOD is 'trusted' or 'password' depending on how MFA is enforced in Vertica
149-
cur.execute("CREATE AUTHENTICATION totp_auth METHOD 'password' HOST '0.0.0.0/0'")
150-
cur.execute("GRANT AUTHENTICATION totp_auth TO totp_user")
151-
152-
# Generate TOTP
153-
TOTP_SECRET = "O5D7DQICJTM34AZROWHSAO4O53ELRJN3"
154-
totp_code = pyotp.TOTP(TOTP_SECRET).now()
155-
156-
# Set connection info
157-
self._conn_info['user'] = 'totp_user'
158-
self._conn_info['password'] = 'password'
159-
self._conn_info['totp'] = totp_code
160-
161-
# Try connection
162-
with self._connect() as totp_conn:
163-
c = totp_conn.cursor()
164-
c.execute("SELECT 1")
165-
res = c.fetchone()
166-
self.assertEqual(res[0], 1)
167-
168-
finally:
169-
cur.execute("DROP USER IF EXISTS totp_user")
170-
cur.execute("DROP AUTHENTICATION IF EXISTS totp_auth CASCADE")
171-
172-
# Negative Test: Missing TOTP
173-
def totp_missing_code_scenario(self):
174-
with self._connect() as conn:
175-
cur = conn.cursor()
176-
177-
cur.execute("DROP USER IF EXISTS totp_user")
178-
cur.execute("DROP AUTHENTICATION IF EXISTS totp_auth CASCADE")
179-
126+
127+
def test_totp_connection(self):
128+
"""
129+
Steps:
130+
1) Admin pre-cleanup and MFA user/auth creation with ENFORCEMFA
131+
2) Attempt user connection to capture enrollment error and extract TOTP secret
132+
3) Generate valid TOTP and verify:
133+
- success with TOTP in connection options
134+
- success via stdin prompt
135+
4) Verify failures for invalid/blank/long/alphanumeric codes via options and stdin
136+
"""
137+
import re
138+
import os
139+
import sys
140+
import pyotp
141+
from ... import connect
142+
from ... import errors
143+
144+
test_user = 'mfa_user'
145+
test_password = 'pwd'
146+
147+
# Admin connection, setup MFA artifacts
148+
with self._connect() as admin:
149+
cur = admin.cursor()
150+
151+
# Pre-cleanup (ignore failures)
152+
cleanup_pre = [
153+
f"DROP USER IF EXISTS {test_user};",
154+
"DROP AUTHENTICATION pw_local_mfa CASCADE;",
155+
"DROP AUTHENTICATION pw_ipv4_mfa CASCADE;",
156+
"DROP AUTHENTICATION pw_ipv6_mfa CASCADE;",
157+
]
158+
for q in cleanup_pre:
159+
try:
160+
cur.execute(q)
161+
except Exception:
162+
pass
163+
164+
# Create user + ENFORCEMFA authentications and grant
165+
dbname = self._conn_info['database']
166+
create_stmts = [
167+
f"CREATE USER {test_user} IDENTIFIED BY '{test_password}';",
168+
f"GRANT ALL PRIVILEGES ON DATABASE {dbname} TO {test_user};",
169+
f"GRANT ALL ON SCHEMA public TO {test_user};",
170+
"CREATE AUTHENTICATION pw_local_mfa METHOD 'password' LOCAL ENFORCEMFA;",
171+
"CREATE AUTHENTICATION pw_ipv4_mfa METHOD 'password' HOST '0.0.0.0/0' ENFORCEMFA;",
172+
"CREATE AUTHENTICATION pw_ipv6_mfa METHOD 'password' HOST '::/0' ENFORCEMFA;",
173+
f"GRANT AUTHENTICATION pw_local_mfa TO {test_user};",
174+
f"GRANT AUTHENTICATION pw_ipv4_mfa TO {test_user};",
175+
f"GRANT AUTHENTICATION pw_ipv6_mfa TO {test_user};",
176+
]
180177
try:
181-
cur.execute("CREATE USER totp_user IDENTIFIED BY 'password' ENFORCEMFA")
182-
cur.execute("CREATE AUTHENTICATION totp_auth METHOD 'password' HOST '0.0.0.0/0'")
183-
cur.execute("GRANT AUTHENTICATION totp_auth TO totp_user")
184-
185-
self._conn_info['user'] = 'totp_user'
186-
self._conn_info['password'] = 'password'
187-
self._conn_info.pop('totp', None) # No TOTP
188-
189-
err_msg = "TOTP was requested but not provided"
190-
self.assertConnectionFail(err_msg=err_msg)
191-
192-
finally:
193-
cur.execute("DROP USER IF EXISTS totp_user")
194-
cur.execute("DROP AUTHENTICATION IF EXISTS totp_auth CASCADE")
195-
196-
# Negative Test: Invalid TOTP Format
197-
def totp_invalid_format_scenario(self):
198-
with self._connect() as conn:
199-
cur = conn.cursor()
200-
201-
cur.execute("DROP USER IF EXISTS totp_user")
202-
cur.execute("DROP AUTHENTICATION IF EXISTS totp_auth CASCADE")
203-
178+
for q in create_stmts:
179+
cur.execute(q)
180+
except Exception as e:
181+
# Older server versions may not support ENFORCEMFA in CREATE AUTHENTICATION
182+
# Perform cleanup and skip gracefully to keep CI green
183+
try:
184+
for q in [
185+
f"DROP USER IF EXISTS {test_user};",
186+
"DROP AUTHENTICATION pw_local_mfa CASCADE;",
187+
"DROP AUTHENTICATION pw_ipv4_mfa CASCADE;",
188+
"DROP AUTHENTICATION pw_ipv6_mfa CASCADE;",
189+
]:
190+
try:
191+
cur.execute(q)
192+
except Exception:
193+
pass
194+
finally:
195+
import pytest
196+
pytest.skip("ENFORCEMFA not supported on this server version; skipping TOTP flow test.")
197+
198+
# Ensure cleanup after test
199+
def _final_cleanup():
204200
try:
205-
cur.execute("CREATE USER totp_user IDENTIFIED BY 'password' ENFORCEMFA")
206-
cur.execute("CREATE AUTHENTICATION totp_auth METHOD 'password' HOST '0.0.0.0/0'")
207-
cur.execute("GRANT AUTHENTICATION totp_auth TO totp_user")
208-
209-
self._conn_info['user'] = 'totp_user'
210-
self._conn_info['password'] = 'password'
211-
self._conn_info['totp'] = "123" # Invalid
212-
213-
err_msg = "Invalid TOTP format"
214-
self.assertConnectionFail(err_msg=err_msg)
215-
216-
finally:
217-
cur.execute("DROP USER IF EXISTS totp_user")
218-
cur.execute("DROP AUTHENTICATION IF EXISTS totp_auth CASCADE")
219-
220-
# Negative Test: Wrong TOTP (Valid format, wrong value)
221-
def totp_wrong_code_scenario(self):
222-
with self._connect() as conn:
223-
cur = conn.cursor()
224-
225-
cur.execute("DROP USER IF EXISTS totp_user")
226-
cur.execute("DROP AUTHENTICATION IF EXISTS totp_auth CASCADE")
201+
with self._connect() as admin2:
202+
c2 = admin2.cursor()
203+
for q in [
204+
f"DROP USER IF EXISTS {test_user};",
205+
"DROP AUTHENTICATION pw_local_mfa CASCADE;",
206+
"DROP AUTHENTICATION pw_ipv4_mfa CASCADE;",
207+
"DROP AUTHENTICATION pw_ipv6_mfa CASCADE;",
208+
]:
209+
try:
210+
c2.execute(q)
211+
except Exception:
212+
pass
213+
except Exception:
214+
pass
215+
216+
# Step 3: Attempt to connect as MFA user to capture enrollment error and TOTP secret
217+
mfa_conn_info = dict(self._conn_info)
218+
mfa_conn_info['user'] = test_user
219+
mfa_conn_info['password'] = test_password
220+
221+
secret = None
222+
# Feed a blank line to stdin to avoid a long interactive prompt
223+
original_stdin = sys.stdin
224+
try:
225+
rfd, wfd = os.pipe()
226+
os.write(wfd, ("\n").encode('utf-8'))
227+
os.close(wfd)
228+
sys.stdin = os.fdopen(rfd)
227229

228230
try:
229-
cur.execute("CREATE USER totp_user IDENTIFIED BY 'password' ENFORCEMFA")
230-
cur.execute("CREATE AUTHENTICATION totp_auth METHOD 'password' HOST '0.0.0.0/0'")
231-
cur.execute("GRANT AUTHENTICATION totp_auth TO totp_user")
232-
233-
self._conn_info['user'] = 'totp_user'
234-
self._conn_info['password'] = 'password'
235-
self._conn_info['totp'] = "999999" # Wrong OTP
236-
237-
err_msg = "Invalid TOTP"
238-
self.assertConnectionFail(err_msg=err_msg)
239-
240-
finally:
241-
cur.execute("DROP USER IF EXISTS totp_user")
242-
cur.execute("DROP AUTHENTICATION IF EXISTS totp_auth CASCADE")
243-
231+
# Expect failure that includes the TOTP secret in error text
232+
with connect(**mfa_conn_info) as _:
233+
# Unexpected success
234+
self.fail('Expected MFA enrollment error was not thrown')
235+
except errors.ConnectionError as e:
236+
msg = str(e)
237+
# Match text like: Your TOTP secret key is "YEUDLX65RD3S5FBW64IBM5W6E6GVWUVJ"
238+
m = re.search(r"(?i)TOTP secret key is\s+\"([A-Z2-7=]+)\"", msg)
239+
if m:
240+
secret = m.group(1)
241+
else:
242+
# If environment doesn't provide enrollment message, skip the flow gracefully
243+
_final_cleanup()
244+
self.skipTest('TOTP enrollment secret not provided by server; skipping MFA flow scenario.')
245+
finally:
246+
sys.stdin = original_stdin
247+
248+
# Step 4: Generate valid TOTP
249+
totp_code = pyotp.TOTP(secret).now()
250+
251+
# Scenario 1: Valid TOTP in connection options
252+
try:
253+
mfa_conn_info['totp'] = totp_code
254+
with connect(**mfa_conn_info) as conn1:
255+
cur1 = conn1.cursor()
256+
cur1.execute('SELECT version()')
257+
_ = cur1.fetchone()
258+
finally:
259+
mfa_conn_info.pop('totp', None)
260+
261+
# Scenario 2: Valid TOTP via stdin
262+
original_stdin = sys.stdin
263+
try:
264+
rfd, wfd = os.pipe()
265+
os.write(wfd, (totp_code + "\n").encode('utf-8'))
266+
os.close(wfd)
267+
sys.stdin = os.fdopen(rfd)
268+
269+
with connect(**mfa_conn_info) as conn2:
270+
cur2 = conn2.cursor()
271+
cur2.execute('SELECT 1')
272+
self.assertEqual(cur2.fetchone()[0], 1)
273+
finally:
274+
sys.stdin = original_stdin
275+
276+
# Scenario 3: Invalid TOTP in options (syntactically valid but wrong value)
277+
try:
278+
mfa_conn_info['totp'] = '123456'
279+
with self.assertRaises(errors.ConnectionError):
280+
with connect(**mfa_conn_info):
281+
pass
282+
finally:
283+
mfa_conn_info.pop('totp', None)
284+
285+
# Scenario 4: Invalid TOTP via stdin (syntactically valid but wrong)
286+
original_stdin = sys.stdin
287+
try:
288+
rfd, wfd = os.pipe()
289+
os.write(wfd, ("123456\n").encode('utf-8'))
290+
os.close(wfd)
291+
sys.stdin = os.fdopen(rfd)
292+
with self.assertRaises(errors.ConnectionError):
293+
with connect(**mfa_conn_info):
294+
pass
295+
finally:
296+
sys.stdin = original_stdin
297+
298+
# Scenario 5: Blank TOTP in options (client-side validation)
299+
try:
300+
mfa_conn_info['totp'] = ''
301+
with self.assertRaises(errors.ConnectionError):
302+
with connect(**mfa_conn_info):
303+
pass
304+
finally:
305+
mfa_conn_info.pop('totp', None)
306+
307+
# Scenario 6: Blank TOTP via stdin (client-side validation)
308+
original_stdin = sys.stdin
309+
try:
310+
rfd, wfd = os.pipe()
311+
os.write(wfd, ("\n").encode('utf-8'))
312+
os.close(wfd)
313+
sys.stdin = os.fdopen(rfd)
314+
with self.assertRaises(errors.ConnectionError):
315+
with connect(**mfa_conn_info):
316+
pass
317+
finally:
318+
sys.stdin = original_stdin
319+
320+
# Scenario 7: Long TOTP in options (client-side validation)
321+
try:
322+
mfa_conn_info['totp'] = '1234567'
323+
with self.assertRaises(errors.ConnectionError):
324+
with connect(**mfa_conn_info):
325+
pass
326+
finally:
327+
mfa_conn_info.pop('totp', None)
328+
329+
# Scenario 8: Long TOTP via stdin (client-side validation)
330+
original_stdin = sys.stdin
331+
try:
332+
rfd, wfd = os.pipe()
333+
os.write(wfd, ("1234567\n").encode('utf-8'))
334+
os.close(wfd)
335+
sys.stdin = os.fdopen(rfd)
336+
with self.assertRaises(errors.ConnectionError):
337+
with connect(**mfa_conn_info):
338+
pass
339+
finally:
340+
sys.stdin = original_stdin
341+
342+
# Scenario 9: Alphanumeric TOTP in options (client-side validation)
343+
try:
344+
mfa_conn_info['totp'] = '12AB34'
345+
with self.assertRaises(errors.ConnectionError):
346+
with connect(**mfa_conn_info):
347+
pass
348+
finally:
349+
mfa_conn_info.pop('totp', None)
350+
351+
# Scenario 10: Alphanumeric TOTP via stdin (client-side validation)
352+
original_stdin = sys.stdin
353+
try:
354+
rfd, wfd = os.pipe()
355+
os.write(wfd, ("12AB34\n").encode('utf-8'))
356+
os.close(wfd)
357+
sys.stdin = os.fdopen(rfd)
358+
with self.assertRaises(errors.ConnectionError):
359+
with connect(**mfa_conn_info):
360+
pass
361+
finally:
362+
sys.stdin = original_stdin
363+
364+
_final_cleanup()

0 commit comments

Comments
 (0)