test_authentication.py

# Copyright (c) 2018-2024 Open Text.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

from .base import VerticaPythonIntegrationTestCase


class AuthenticationTestCase(VerticaPythonIntegrationTestCase):
    def setUp(self):
        super(AuthenticationTestCase, self).setUp()
        self._user = self._conn_info['user']
        self._password = self._conn_info['password']

    def tearDown(self):
        self._conn_info['user'] = self._user
        self._conn_info['password'] = self._password
        if 'oauth_access_token' in self._conn_info:
            del self._conn_info['oauth_access_token']
        super(AuthenticationTestCase, self).tearDown()

    def test_SHA512(self):
        with self._connect() as conn:
            cur = conn.cursor()
            cur.execute("DROP USER IF EXISTS sha512_user")
            cur.execute("DROP AUTHENTICATION IF EXISTS testIPv4hostHash CASCADE")
            cur.execute("DROP AUTHENTICATION IF EXISTS testIPv6hostHash CASCADE")
            cur.execute("DROP AUTHENTICATION IF EXISTS testlocalHash CASCADE")
            try:
                cur.execute("CREATE USER sha512_user IDENTIFIED BY 'password'")
                cur.execute("ALTER USER sha512_user SECURITY_ALGORITHM 'SHA512'")
                cur.execute("CREATE AUTHENTICATION testIPv4hostHash METHOD 'hash' HOST '0.0.0.0/0'")
                cur.execute("GRANT AUTHENTICATION testIPv4hostHash TO sha512_user")
                cur.execute("CREATE AUTHENTICATION testIPv6hostHash METHOD 'hash' HOST '::/0'")
                cur.execute("GRANT AUTHENTICATION testIPv6hostHash TO sha512_user")
                cur.execute("CREATE AUTHENTICATION testlocalHash METHOD 'hash' LOCAL")
                cur.execute("GRANT AUTHENTICATION testlocalHash TO sha512_user")

                self._conn_info['user'] = 'sha512_user'
                self._conn_info['password'] = 'password'

                expire_msg = 'The password for user sha512_user has expired'
                self.assertConnectionFail(err_msg=expire_msg)

                # Test SHA512 connection
                cur.execute("ALTER USER sha512_user IDENTIFIED BY 'password'")
                self.assertConnectionSuccess()

                # Switch to MD5 for hashing and storing the user password
                cur.execute("ALTER USER sha512_user SECURITY_ALGORITHM 'MD5'")
                self.assertConnectionFail(err_msg=expire_msg)

                # Test MD5 connection
                cur.execute("ALTER USER sha512_user IDENTIFIED BY 'password'")
                self.assertConnectionSuccess()
            finally:
                # Must clean up authentication methods within this session, no
                # matter whether an exception has occurred or not, otherwise
                # those authentication methods may affect existing users
                cur.execute("DROP USER IF EXISTS sha512_user")
                cur.execute("DROP AUTHENTICATION IF EXISTS testIPv4hostHash CASCADE")
                cur.execute("DROP AUTHENTICATION IF EXISTS testIPv6hostHash CASCADE")
                cur.execute("DROP AUTHENTICATION IF EXISTS testlocalHash CASCADE")

    def test_password_expire(self):
        with self._connect() as conn:
            cur = conn.cursor()
            cur.execute("DROP USER IF EXISTS pw_expire_user")
            cur.execute("DROP AUTHENTICATION IF EXISTS testIPv4hostHash CASCADE")
            cur.execute("DROP AUTHENTICATION IF EXISTS testIPv6hostHash CASCADE")
            cur.execute("DROP AUTHENTICATION IF EXISTS testlocalHash CASCADE")
            try:
                cur.execute("CREATE USER pw_expire_user IDENTIFIED BY 'password'")
                cur.execute("CREATE AUTHENTICATION testIPv4hostHash METHOD 'hash' HOST '0.0.0.0/0'")
                cur.execute("GRANT AUTHENTICATION testIPv4hostHash TO pw_expire_user")
                cur.execute("CREATE AUTHENTICATION testIPv6hostHash METHOD 'hash' HOST '::/0'")
                cur.execute("GRANT AUTHENTICATION testIPv6hostHash TO pw_expire_user")
                cur.execute("CREATE AUTHENTICATION testlocalHash METHOD 'hash' LOCAL")
                cur.execute("GRANT AUTHENTICATION testlocalHash TO pw_expire_user")

                # Test connection
                self._conn_info['user'] = 'pw_expire_user'
                self._conn_info['password'] = 'password'
                self.assertConnectionSuccess()

                # Expire the user's password immediately
                cur.execute("ALTER USER pw_expire_user PASSWORD EXPIRE")

                # Expect an error message
                self.assertConnectionFail(
                    err_msg='The password for user pw_expire_user has expired')
            finally:
                # Must clean up authentication methods within this session, no
                # matter whether an exception has occurred or not, otherwise
                # those authentication methods may affect existing users
                cur.execute("DROP USER IF EXISTS pw_expire_user")
                cur.execute("DROP AUTHENTICATION IF EXISTS testIPv4hostHash CASCADE")
                cur.execute("DROP AUTHENTICATION IF EXISTS testIPv6hostHash CASCADE")
                cur.execute("DROP AUTHENTICATION IF EXISTS testlocalHash CASCADE")

    def test_oauth_access_token(self):
        self.require_protocol_at_least(3 << 16 | 11)
        if not self._oauth_info['access_token']:
            self.skipTest('OAuth access token not set')
        if not self._oauth_info['user'] and not self._conn_info['database']:
            self.skipTest('Both database and oauth_user are not set')

        self._conn_info['user'] = self._oauth_info['user']
        self._conn_info['oauth_access_token'] = self._oauth_info['access_token']
        with self._connect() as conn:
            cur = conn.cursor()
            cur.execute("SELECT authentication_method FROM sessions WHERE session_id=(SELECT current_session())")
            res = cur.fetchone()
            self.assertEqual(res[0], 'OAuth')
    
    def test_totp_connection(self):
        """
        Steps:
        1) Admin pre-cleanup and MFA user/auth creation with ENFORCEMFA
        2) Attempt user connection to capture enrollment error and extract TOTP secret
        3) Generate valid TOTP and verify:
           - success with TOTP in connection options
           - success via stdin prompt
        4) Verify failures for invalid/blank/long/alphanumeric codes via options and stdin
        """
        import re
        import os
        import sys
        import pyotp
        from ... import connect
        from ... import errors

        test_user = 'mfa_user'
        test_password = 'pwd'

        # Admin connection, setup MFA artifacts
        with self._connect() as admin:
            cur = admin.cursor()

            # Pre-cleanup (ignore failures)
            cleanup_pre = [
                f"DROP USER IF EXISTS {test_user};",
                "DROP AUTHENTICATION pw_local_mfa CASCADE;",
                "DROP AUTHENTICATION pw_ipv4_mfa CASCADE;",
                "DROP AUTHENTICATION pw_ipv6_mfa CASCADE;",
            ]
            for q in cleanup_pre:
                try:
                    cur.execute(q)
                except Exception:
                    pass

            # Create user + ENFORCEMFA authentications and grant
            dbname = self._conn_info['database']
            create_stmts = [
                f"CREATE USER {test_user} IDENTIFIED BY '{test_password}';",
                f"GRANT ALL PRIVILEGES ON DATABASE {dbname} TO {test_user};",
                f"GRANT ALL ON SCHEMA public TO {test_user};",
                "CREATE AUTHENTICATION pw_local_mfa METHOD 'password' LOCAL ENFORCEMFA;",
                "CREATE AUTHENTICATION pw_ipv4_mfa METHOD 'password' HOST '0.0.0.0/0' ENFORCEMFA;",
                "CREATE AUTHENTICATION pw_ipv6_mfa METHOD 'password' HOST '::/0' ENFORCEMFA;",
                f"GRANT AUTHENTICATION pw_local_mfa TO {test_user};",
                f"GRANT AUTHENTICATION pw_ipv4_mfa TO {test_user};",
                f"GRANT AUTHENTICATION pw_ipv6_mfa TO {test_user};",
            ]
            try:
                for q in create_stmts:
                    cur.execute(q)
            except Exception as e:
                # Older server versions may not support ENFORCEMFA in CREATE AUTHENTICATION
                # Perform cleanup and skip gracefully to keep CI green
                try:
                    for q in [
                        f"DROP USER IF EXISTS {test_user};",
                        "DROP AUTHENTICATION pw_local_mfa CASCADE;",
                        "DROP AUTHENTICATION pw_ipv4_mfa CASCADE;",
                        "DROP AUTHENTICATION pw_ipv6_mfa CASCADE;",
                    ]:
                        try:
                            cur.execute(q)
                        except Exception:
                            pass
                finally:
                    import pytest
                    pytest.skip("ENFORCEMFA not supported on this server version; skipping TOTP flow test.")

        # Ensure cleanup after test
        def _final_cleanup():
            try:
                with self._connect() as admin2:
                    c2 = admin2.cursor()
                    for q in [
                        f"DROP USER IF EXISTS {test_user};",
                        "DROP AUTHENTICATION pw_local_mfa CASCADE;",
                        "DROP AUTHENTICATION pw_ipv4_mfa CASCADE;",
                        "DROP AUTHENTICATION pw_ipv6_mfa CASCADE;",
                    ]:
                        try:
                            c2.execute(q)
                        except Exception:
                            pass
            except Exception:
                pass

        # Step 3: Attempt to connect as MFA user to capture enrollment error and TOTP secret
        mfa_conn_info = dict(self._conn_info)
        mfa_conn_info['user'] = test_user
        mfa_conn_info['password'] = test_password

        secret = None
        # Feed a blank line to stdin to avoid a long interactive prompt
        original_stdin = sys.stdin
        try:
            rfd, wfd = os.pipe()
            os.write(wfd, ("\n").encode('utf-8'))
            os.close(wfd)
            sys.stdin = os.fdopen(rfd)

            try:
                # Expect failure that includes the TOTP secret in error text
                with connect(**mfa_conn_info) as _:
                    # Unexpected success
                    self.fail('Expected MFA enrollment error was not thrown')
            except errors.ConnectionError as e:
                msg = str(e)
                # Match text like: Your TOTP secret key is "YEUDLX65RD3S5FBW64IBM5W6E6GVWUVJ"
                m = re.search(r"(?i)TOTP secret key is\s+\"([A-Z2-7=]+)\"", msg)
                if m:
                    secret = m.group(1)
                else:
                    # If environment doesn't provide enrollment message, skip the flow gracefully
                    _final_cleanup()
                    self.skipTest('TOTP enrollment secret not provided by server; skipping MFA flow scenario.')
        finally:
            sys.stdin = original_stdin

        # Step 4: Generate valid TOTP
        totp_code = pyotp.TOTP(secret).now()

        # Scenario 1: Valid TOTP in connection options
        try:
            mfa_conn_info['totp'] = totp_code
            with connect(**mfa_conn_info) as conn1:
                cur1 = conn1.cursor()
                cur1.execute('SELECT version()')
                _ = cur1.fetchone()
        finally:
            mfa_conn_info.pop('totp', None)

        # Scenario 2: Valid TOTP via stdin
        original_stdin = sys.stdin
        try:
            rfd, wfd = os.pipe()
            os.write(wfd, (totp_code + "\n").encode('utf-8'))
            os.close(wfd)
            sys.stdin = os.fdopen(rfd)

            with connect(**mfa_conn_info) as conn2:
                cur2 = conn2.cursor()
                cur2.execute('SELECT 1')
                self.assertEqual(cur2.fetchone()[0], 1)
        finally:
            sys.stdin = original_stdin

        # Scenario 3: Invalid TOTP in options (syntactically valid but wrong value)
        try:
            mfa_conn_info['totp'] = '123456'
            with self.assertRaises(errors.ConnectionError):
                with connect(**mfa_conn_info):
                    pass
        finally:
            mfa_conn_info.pop('totp', None)

        # Scenario 4: Invalid TOTP via stdin (syntactically valid but wrong)
        original_stdin = sys.stdin
        try:
            rfd, wfd = os.pipe()
            os.write(wfd, ("123456\n").encode('utf-8'))
            os.close(wfd)
            sys.stdin = os.fdopen(rfd)
            with self.assertRaises(errors.ConnectionError):
                with connect(**mfa_conn_info):
                    pass
        finally:
            sys.stdin = original_stdin

        # Scenario 5: Blank TOTP in options (client-side validation)
        try:
            mfa_conn_info['totp'] = ''
            with self.assertRaises(errors.ConnectionError):
                with connect(**mfa_conn_info):
                    pass
        finally:
            mfa_conn_info.pop('totp', None)

        # Scenario 6: Blank TOTP via stdin (client-side validation)
        original_stdin = sys.stdin
        try:
            rfd, wfd = os.pipe()
            os.write(wfd, ("\n").encode('utf-8'))
            os.close(wfd)
            sys.stdin = os.fdopen(rfd)
            with self.assertRaises(errors.ConnectionError):
                with connect(**mfa_conn_info):
                    pass
        finally:
            sys.stdin = original_stdin

        # Scenario 7: Long TOTP in options (client-side validation)
        try:
            mfa_conn_info['totp'] = '1234567'
            with self.assertRaises(errors.ConnectionError):
                with connect(**mfa_conn_info):
                    pass
        finally:
            mfa_conn_info.pop('totp', None)

        # Scenario 8: Long TOTP via stdin (client-side validation)
        original_stdin = sys.stdin
        try:
            rfd, wfd = os.pipe()
            os.write(wfd, ("1234567\n").encode('utf-8'))
            os.close(wfd)
            sys.stdin = os.fdopen(rfd)
            with self.assertRaises(errors.ConnectionError):
                with connect(**mfa_conn_info):
                    pass
        finally:
            sys.stdin = original_stdin

        # Scenario 9: Alphanumeric TOTP in options (client-side validation)
        try:
            mfa_conn_info['totp'] = '12AB34'
            with self.assertRaises(errors.ConnectionError):
                with connect(**mfa_conn_info):
                    pass
        finally:
            mfa_conn_info.pop('totp', None)

        # Scenario 10: Alphanumeric TOTP via stdin (client-side validation)
        original_stdin = sys.stdin
        try:
            rfd, wfd = os.pipe()
            os.write(wfd, ("12AB34\n").encode('utf-8'))
            os.close(wfd)
            sys.stdin = os.fdopen(rfd)
            with self.assertRaises(errors.ConnectionError):
                with connect(**mfa_conn_info):
                    pass
        finally:
            sys.stdin = original_stdin

        _final_cleanup()