Refactor cookie init into shared helper and expand test coverage.

Author: jonasbardino

tests/test_mig_shared_auth.py

@@ -42,6 +42,9 @@
 
 TEST_USER_DN = \
     '/C=DK/ST=NA/L=NA/O=Test Org/OU=NA/CN=Test User/emailAddress=test@example.com'
+GDP_USER_DN = f'{TEST_USER_DN}/GDP=projectx'
+TEST_CLIENT_PREFIX = \
+    '2e1c3d78bddf637ed6b83067c15ac9b9893545ff6a549519178cb4a252ed38b5'
 
 
 class MigSharedAuth__twofactor(MigTestCase):
@@ -50,6 +53,19 @@ class MigSharedAuth__twofactor(MigTestCase):
     def _provide_configuration(self):
         return 'testconfig'
 
+    def _mimic_cookie_init(self, session_key):
+        """Mimic twofactor session cookie setup"""
+        cookie = http.cookies.SimpleCookie()
+        session_start = time.time()
+        cookie['2FA_Auth'] = session_key
+        cookie['2FA_Auth']['path'] = '/'
+        # NOTE: SimpleCookie translates expires ttl to actual date from now
+        cookie['2FA_Auth']['expires'] = twofactor_cookie_ttl
+        cookie['2FA_Auth']['secure'] = True
+        cookie['2FA_Auth']['httponly'] = True
+        environ = {'HTTP_COOKIE': cookie}
+        return environ
+
     def before_each(self):
         """Setup test environment before each test method"""
         ensure_dirs_exist(self.configuration.user_cache)
@@ -58,6 +74,18 @@ def before_each(self):
         ensure_dirs_exist(self.configuration.mrsl_files_dir)
         ensure_dirs_exist(self.configuration.resource_pending)
         ensure_dirs_exist(self.configuration.twofactor_home)
+        self.configuration.site_enable_gdp = False  # Disable GDP by default
+
+    def test_twofactor_available(self):
+        """Test pyotp availability detection"""
+        # Positive case (assuming pyotp is installed)
+        self.assertTrue(auth.twofactor_available(self.configuration))
+
+        # Simulate missing pyotp
+        original_pyotp = auth.pyotp
+        auth.pyotp = None
+        self.assertFalse(auth.twofactor_available(self.configuration))
+        auth.pyotp = original_pyotp
 
     def test_get_twofactor_secrets_generates_valid_key(self):
         """Test get_twofactor_secrets generates and stores valid key"""
@@ -165,16 +193,8 @@ def test_twofactor_session_lifecycle(self):
                          'session should have correct TTL')
 
         # Mimic cookie init
-        cookie = http.cookies.SimpleCookie()
-        session_start = time.time()
-        cookie['2FA_Auth'] = session_key
-        cookie['2FA_Auth']['path'] = '/'
-        # NOTE: SimpleCookie translates expires ttl to actual date from now
-        cookie['2FA_Auth']['expires'] = twofactor_cookie_ttl
-        cookie['2FA_Auth']['secure'] = True
-        cookie['2FA_Auth']['httponly'] = True
+        environ = self._mimic_cookie_init(session_key)
 
-        environ = {'HTTP_COOKIE': cookie}
         # Expire session
         expire_result = auth.expire_twofactor_session(self.configuration,
                                                       TEST_USER_DN, environ)
@@ -184,6 +204,169 @@ def test_twofactor_session_lifecycle(self):
         self.assertNotIn(session_key, sessions_after,
                          'session should be removed')
 
+    def test_load_twofactor_key_missing(self):
+        """Test handling of missing twofactor key"""
+        self._provision_test_user(self, TEST_USER_DN)
+        result = auth.load_twofactor_key(TEST_USER_DN, self.configuration,
+                                         allow_missing=True)
+        self.assertIsNone(result, 'should return None for missing key')
+
+    def test_generate_session_prefix(self):
+        """Test session prefix generation"""
+        prefix = auth.generate_session_prefix(self.configuration, TEST_USER_DN)
+        self.assertEqual(prefix, TEST_CLIENT_PREFIX)
+
+    def test_gdp_client_id_transformation(self):
+        """Test GDP client ID normalization"""
+        self.configuration.site_enable_gdp = True
+        from mig.shared.gdp.all import get_base_client_id
+        base_dn = get_base_client_id(self.configuration, GDP_USER_DN,
+                                     expand_oid_alias=False)
+
+        # With GDP enabled, ID should be transformed to base project
+        session_key = auth.generate_session_key(self.configuration,
+                                                GDP_USER_DN)
+        base_prefix = auth.generate_session_prefix(self.configuration, base_dn)
+        self.assertTrue(session_key.startswith(base_prefix))
+
+    def test_client_twofactor_session_parsing(self):
+        """Test cookie parsing for session ID extraction"""
+        cookie = http.cookies.SimpleCookie()
+        cookie['2FA_Auth'] = 'test_session_id'
+        environ = {'HTTP_COOKIE': cookie.output(header='')}
+
+        session_id = auth.client_twofactor_session(self.configuration,
+                                                   TEST_USER_DN, environ)
+        self.assertEqual(session_id, 'test_session_id',
+                         'should extract session ID from cookies')
+
+    def test_multiple_active_sessions(self):
+        """Test handling of multiple concurrent sessions"""
+        client_dir = self._provision_test_user(self, TEST_USER_DN)
+
+        # Create 3 sessions from different addresses
+        sessions = []
+        for addr in ['192.168.0.1', '10.0.0.1', '172.16.0.1']:
+            session_key = auth.generate_session_key(self.configuration,
+                                                    TEST_USER_DN)
+            auth.save_twofactor_session(
+                self.configuration, TEST_USER_DN, session_key,
+                addr, 'TestAgent', time.time()
+            )
+            sessions.append(session_key)
+
+        # List all sessions
+        all_sessions = auth.list_twofactor_sessions(self.configuration,
+                                                    TEST_USER_DN)
+        self.assertEqual(len(all_sessions), 3,
+                         'should list all active sessions')
+
+        # Filter by address
+        filtered = auth.list_twofactor_sessions(
+            self.configuration, TEST_USER_DN, user_addr='10.0.0.1')
+        self.assertEqual(len(filtered), 1,
+                         'should filter sessions by address')
+
+    def test_expired_session_cleanup(self):
+        """Test automatic exclusion of expired sessions"""
+        client_dir = self._provision_test_user(self, TEST_USER_DN)
+        valid_key = auth.generate_session_key(self.configuration,
+                                              TEST_USER_DN)
+        expired_key = auth.generate_session_key(self.configuration,
+                                                TEST_USER_DN)
+
+        # Create valid session (expires future)
+        auth.save_twofactor_session(
+            self.configuration, TEST_USER_DN, valid_key,
+            '127.0.0.1', 'TestAgent', time.time()
+        )
+
+        # Create expired session
+        auth.save_twofactor_session(
+            self.configuration, TEST_USER_DN, expired_key,
+            '127.0.0.1', 'TestAgent', time.time() - twofactor_cookie_ttl - 100
+        )
+
+        # Only valid session should appear in active listings
+        active = auth.active_twofactor_session(
+            self.configuration, TEST_USER_DN, '127.0.0.1')
+        self.assertEqual(active['session_key'], valid_key,
+                         'should only return active sessions')
+
+    def test_custom_token_interval(self):
+        """Test token verification with custom interval"""
+        self._provision_test_user(self, TEST_USER_DN)
+
+        # Save custom interval
+        client_dir = client_id_dir(TEST_USER_DN)
+        interval_path = os.path.join(
+            self.configuration.user_settings,
+            client_dir,
+            'twofactor_interval'
+        )
+        with open(interval_path, 'w') as fh:
+            fh.write('60')  # 1 minute interval
+
+        # Generate key with custom interval
+        b32_key = auth.reset_twofactor_key(
+            TEST_USER_DN, self.configuration, interval=60)
+
+        totp = auth.get_totp(TEST_USER_DN, b32_key, self.configuration)
+        valid_token = totp.now()
+
+        # Verify works with custom interval
+        result = auth.verify_twofactor_token(
+            self.configuration, TEST_USER_DN, b32_key, valid_token)
+        self.assertTrue(result, 'should accept token with custom interval')
+
+    def test_strict_address_session_handling(self):
+        """Test session handling with strict address enforcement"""
+        self.configuration.site_twofactor_strict_address = True
+        session_key = auth.generate_session_key(self.configuration,
+                                                TEST_USER_DN)
+        user_addr = '192.168.1.100'
+
+        auth.save_twofactor_session(
+            self.configuration, TEST_USER_DN, session_key,
+            user_addr, 'TestAgent', time.time()
+        )
+
+        # Should have address-linked file
+        addr_file = os.path.join(
+            self.configuration.twofactor_home,
+            f"{user_addr}_{session_key}"
+        )
+        self.assertTrue(os.path.exists(addr_file),
+                        'should create address-linked session file')
+
+        # Mimic cookie init
+        environ = self._mimic_cookie_init(session_key)
+
+        # Expire should remove both files
+        auth.expire_twofactor_session(self.configuration, TEST_USER_DN,
+                                      environ, user_addr=user_addr)
+        self.assertFalse(os.path.exists(addr_file),
+                         'should remove address-linked session file')
+        self.assertFalse(os.path.exists(
+            os.path.join(self.configuration.twofactor_home, session_key)
+        ), 'should remove main session file')
+
+    def test_check_twofactor_active_missing_session(self):
+        """Test session check when session cookie is missing"""
+        self._provision_test_user(self, TEST_USER_DN)
+        # Empty environment
+        result = auth.check_twofactor_active(self.configuration,
+                                             TEST_USER_DN, '127.0.0.1', {})
+        self.assertFalse(result, 'should reject missing session')
+
+        # Invalid cookie
+        cookie = http.cookies.SimpleCookie()
+        cookie['invalid'] = 'value'
+        environ = {'HTTP_COOKIE': cookie.output(header='')}
+        result = auth.check_twofactor_active(self.configuration,
+                                             TEST_USER_DN, '127.0.0.1', environ)
+        self.assertFalse(result, 'should reject invalid session cookie')
+
 
 if __name__ == '__main__':
     testmain()