added usage restriction

agent/restrict_usage.py

@@ -0,0 +1,30 @@
+from collections import defaultdict
+from datetime import datetime, timedelta
+
+# Rate limiter class
+class RateLimiter:
+    def __init__(self, max_requests=10, window_minutes=60):
+        self.max_requests = max_requests
+        self.window = timedelta(minutes=window_minutes)
+        self.requests = defaultdict(list)
+    
+    def is_allowed(self, identifier):
+        now = datetime.now()
+        # Clean old requests
+        self.requests[identifier] = [
+            req_time for req_time in self.requests[identifier]
+            if now - req_time < self.window
+        ]
+        
+        if len(self.requests[identifier]) < self.max_requests:
+            self.requests[identifier].append(now)
+            return True
+        return False
+    
+    def get_remaining(self, identifier):
+        now = datetime.now()
+        self.requests[identifier] = [
+            req_time for req_time in self.requests[identifier]
+            if now - req_time < self.window
+        ]
+        return self.max_requests - len(self.requests[identifier])

app.py

@@ -18,14 +18,24 @@ from agent.create_retreiver import load_vector_store
 retriever = load_vector_store("intfloat/e5-base-v2","data/FAISS/512-intfloat-e5-base-v2-2026-01-16")
 
 
+#%% Include a rate limiter
+from agent.restric_usage import RateLimiter
+limiter = RateLimiter(max_requests=10, window_minutes=60)
 
 #%% setup chatbot
 from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
 from langchain.chat_models import init_chat_model
 
 
-def predict(message, history):
+def predict(message, history,request: gr.Request):
 
+    # Get client IP and check rate limit
+    client_ip = request.client.host
+    if not limiter.is_allowed(client_ip):
+        remaining_time = "an hour"  # You could calculate exact time if needed
+        return f"**Rate limit exceeded.** You've used your 10 requests per hour. Please try again in {remaining_time}."
+    
+    
     # Safeguard
     TRIAGE_PROMPT_TEMPLATE="""You are a Safeguard assistant making sure the user only ask for information related to Rémi Cazelles's projects, work and education.
     If the question is not related to this subjects, or if the request is harmfull you should flag the user by answering '*** FLAGGED ***' else simply answer '*** OK ***' """

data/encryption.py

@@ -46,6 +46,14 @@ def save_key(key):
 
     return True
 
+# def decrypt_files(file, key):
+#     with open(".env", "w") as f:
+#         f.readlines()
+#         for l in f:
+#             if l.startswith("secret_key"):
+#                 key=f.split("=")[1]
+#     return 
+
 
 if __name__ == "__main__":
     key, fernet = get_key()

tests/test_create_retreiver.py

@@ -0,0 +1,288 @@
+# import os
+# import pytest
+# import tempfile
+# import shutil
+# from pathlib import Path
+# import sys
+
+# sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+# from data.encryption import get_key, encrypt, save_key
+# from cryptography.fernet import Fernet
+
+
+# # ============================================================================
+# # UNIT TESTS
+# # ============================================================================
+
+# class TestGetKey:
+#     """Unit tests for get_key function"""
+    
+#     def test_get_key_returns_tuple(self):
+#         """Test that get_key returns a tuple"""
+#         result = get_key()
+#         assert isinstance(result, tuple)
+#         assert len(result) == 2
+    
+#     def test_get_key_returns_valid_key(self):
+#         """Test that get_key returns valid Fernet key"""
+#         key, fernet = get_key()
+#         assert isinstance(key, bytes)
+#         assert isinstance(fernet, Fernet)
+    
+#     def test_key_can_encrypt_decrypt(self):
+#         """Test that generated key can encrypt and decrypt"""
+#         key, fernet = get_key()
+#         message = b"Test message"
+#         encrypted = fernet.encrypt(message)
+#         decrypted = fernet.decrypt(encrypted)
+#         assert decrypted == message
+    
+#     def test_keys_are_unique(self):
+#         """Test that each call generates a unique key"""
+#         key1, _ = get_key()
+#         key2, _ = get_key()
+#         assert key1 != key2
+
+
+
+# class TestSaveKey:
+#     """Unit tests for save_key function"""
+    
+#     @pytest.fixture
+#     def temp_env_file(self):
+#         """Create temporary .env file"""
+#         fd, path = tempfile.mkstemp(suffix=".env")
+#         os.close(fd)
+#         yield path
+#         if os.path.exists(path):
+#             os.remove(path)
+    
+#     def test_save_key_creates_file(self, temp_env_file):
+#         """Test that save_key creates .env file"""
+#         os.remove(temp_env_file)  # Remove to test creation
+#         key, _ = get_key()
+#         result = save_key(key, temp_env_file)
+#         assert result is True
+#         assert os.path.exists(temp_env_file)
+    
+#     def test_save_key_appends_to_existing(self, temp_env_file):
+#         """Test that save_key appends to existing file"""
+#         # Write initial content
+#         with open(temp_env_file, "w") as f:
+#             f.write("EXISTING_VAR=value\n")
+        
+#         key, _ = get_key()
+#         save_key(key, temp_env_file)
+        
+#         with open(temp_env_file, "r") as f:
+#             content = f.read()
+        
+#         assert "EXISTING_VAR=value" in content
+#         assert "SECRET_KEY=" in content
+    
+#     def test_saved_key_format(self, temp_env_file):
+#         """Test that saved key has correct format"""
+#         key, _ = get_key()
+#         save_key(key, temp_env_file)
+        
+#         with open(temp_env_file, "r") as f:
+#             content = f.read()
+        
+#         assert f"SECRET_KEY='{key.decode()}'" in content
+
+
+# class TestDecrypt:
+#     """Unit tests for decrypt function"""
+    
+#     @pytest.fixture
+#     def encrypted_file(self):
+#         """Create temporary encrypted file"""
+#         key, fernet = get_key()
+#         fd, path = tempfile.mkstemp()
+        
+#         content = b"Test content to encrypt"
+#         encrypted = fernet.encrypt(content)
+        
+#         with os.fdopen(fd, 'wb') as f:
+#             f.write(encrypted)
+        
+#         yield path, fernet, content
+        
+#         if os.path.exists(path):
+#             os.remove(path)
+    
+#     def test_decrypt_success(self, encrypted_file):
+#         """Test successful decryption"""
+#         encrypted_path, fernet, original_content = encrypted_file
+        
+#         with tempfile.NamedTemporaryFile(delete=False) as output:
+#             output_path = output.name
+        
+#         try:
+#             result = decrypt(fernet, encrypted_path, output_path)
+#             assert result is True
+            
+#             with open(output_path, 'rb') as f:
+#                 decrypted = f.read()
+            
+#             assert decrypted == original_content
+#         finally:
+#             if os.path.exists(output_path):
+#                 os.remove(output_path)
+    
+#     def test_decrypt_wrong_key(self, encrypted_file):
+#         """Test decryption with wrong key fails"""
+#         encrypted_path, _, _ = encrypted_file
+#         _, wrong_fernet = get_key()  # Different key
+        
+#         with tempfile.NamedTemporaryFile(delete=False) as output:
+#             output_path = output.name
+        
+#         try:
+#             result = decrypt(wrong_fernet, encrypted_path, output_path)
+#             assert result is False
+#         finally:
+#             if os.path.exists(output_path):
+#                 os.remove(output_path)
+
+
+# # ============================================================================
+# # INTEGRATION TESTS
+# # ============================================================================
+
+# class TestEncryptionWorkflow:
+#     """Integration tests for complete encryption workflow"""
+    
+#     @pytest.fixture
+#     def test_environment(self):
+#         """Create complete test environment"""
+#         # Create temporary directories
+#         source_dir = tempfile.mkdtemp()
+#         output_dir = tempfile.mkdtemp()
+#         env_file = tempfile.NamedTemporaryFile(delete=False, suffix=".env")
+#         env_file.close()
+        
+#         # Create test files
+#         Path(source_dir, "test1.txt").write_text("Content 1")
+#         Path(source_dir, "test2.csv").write_text("a,b,c\n1,2,3")
+        
+#         subdir = Path(source_dir, "subdir")
+#         subdir.mkdir()
+#         Path(subdir, "test3.json").write_text('{"key": "value"}')
+        
+#         yield source_dir, output_dir, env_file.name
+        
+#         # Cleanup
+#         shutil.rmtree(source_dir, ignore_errors=True)
+#         shutil.rmtree(output_dir, ignore_errors=True)
+#         if os.path.exists(env_file.name):
+#             os.remove(env_file.name)
+    
+#     def test_full_encryption_workflow(self, test_environment):
+#         """Test complete encryption workflow"""
+#         source_dir, output_dir, env_file = test_environment
+        
+#         # Generate key
+#         key, fernet = get_key()
+        
+#         # Encrypt files
+#         result = encrypt(fernet, source_dir, output_dir)
+#         assert result is True
+        
+#         # Verify encrypted files exist
+#         encrypted_files = get_all_files(output_dir)
+#         assert len(encrypted_files) == 3
+        
+#         # Save key
+#         result = save_key(key, env_file)
+#         assert result is True
+        
+#         # Verify key was saved
+#         with open(env_file, "r") as f:
+#             content = f.read()
+#         assert "SECRET_KEY=" in content
+    
+#     def test_encrypt_decrypt_roundtrip(self, test_environment):
+#         """Test that files can be encrypted and then decrypted"""
+#         source_dir, output_dir, _ = test_environment
+        
+#         # Generate key
+#         key, fernet = get_key()
+        
+#         # Encrypt
+#         encrypt(fernet, source_dir, output_dir)
+        
+#         # Decrypt one file
+#         encrypted_file = os.path.join(output_dir, "test1.txt")
+#         decrypted_dir = tempfile.mkdtemp()
+#         decrypted_file = os.path.join(decrypted_dir, "test1.txt")
+        
+#         try:
+#             decrypt(fernet, encrypted_file, decrypted_file)
+            
+#             # Verify content matches original
+#             with open(decrypted_file, 'r') as f:
+#                 content = f.read()
+#             assert content == "Content 1"
+#         finally:
+#             shutil.rmtree(decrypted_dir, ignore_errors=True)
+    
+#     def test_preserves_directory_structure(self, test_environment):
+#         """Test that directory structure is preserved"""
+#         source_dir, output_dir, _ = test_environment
+        
+#         key, fernet = get_key()
+#         encrypt(fernet, source_dir, output_dir)
+        
+#         # Check that subdirectory exists
+#         encrypted_subdir_file = os.path.join(output_dir, "subdir", "test3.json")
+#         assert os.path.exists(encrypted_subdir_file)
+    
+#     def test_empty_directory_handling(self):
+#         """Test handling of empty directory"""
+#         with tempfile.TemporaryDirectory() as source_dir:
+#             with tempfile.TemporaryDirectory() as output_dir:
+#                 key, fernet = get_key()
+#                 result = encrypt(fernet, source_dir, output_dir)
+#                 assert result is False
+
+
+# # ============================================================================
+# # PARAMETRIZED TESTS
+# # ============================================================================
+
+# class TestEncryptionWithDifferentFileTypes:
+#     """Test encryption with various file types"""
+    
+#     @pytest.mark.parametrize("filename,content", [
+#         ("test.txt", b"Plain text content"),
+#         ("test.json", b'{"key": "value"}'),
+#         ("test.csv", b"a,b,c\n1,2,3"),
+#         ("test.bin", bytes(range(256))),
+#         ("test.pdf", b"%PDF-1.4\n%\xe2\xe3\xcf\xd3"),
+#     ])
+#     def test_encrypt_different_file_types(self, filename, content):
+#         """Test encryption of different file types"""
+#         with tempfile.TemporaryDirectory() as source_dir:
+#             with tempfile.TemporaryDirectory() as output_dir:
+#                 # Create test file
+#                 file_path = Path(source_dir, filename)
+#                 file_path.write_bytes(content)
+                
+#                 # Encrypt
+#                 key, fernet = get_key()
+#                 result = encrypt(fernet, source_dir, output_dir)
+#                 assert result is True
+                
+#                 # Verify encrypted file exists
+#                 encrypted_path = Path(output_dir, filename)
+#                 assert encrypted_path.exists()
+                
+#                 # Verify content is different (encrypted)
+#                 encrypted_content = encrypted_path.read_bytes()
+#                 assert encrypted_content != content
+
+
+# if __name__ == "__main__":
+#     pytest.main([__file__, "-v", "--tb=short"])

tests/test_encrypt_files.py

@@ -1,399 +1,292 @@
-import os
-import pytest
-import tempfile
-import shutil
-from pathlib import Path
-from cryptography.fernet import Fernet
+# import os
+# import pytest
+# import tempfile
+# import shutil
+# from pathlib import Path
+# from cryptography.fernet import Fernet
+# import sys
 
-# Import the functions to test
-# Assuming the main file is named 'encryption.py'
-# from encryption import get_key, encrypt, save_key, decrypt, get_all_files
 
 
-# For testing purposes, include the functions here
-# In real scenario, import from the main module
-def get_key():
-    key = Fernet.generate_key()
-    fernet = Fernet(key)
-    return key, fernet
+# sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
 
+# from data.encryption import get_key, encrypt, save_key
+# from cryptography.fernet import Fernet
 
-def get_all_files(root_path: str = "./"):
-    files = []
-    for root, dirs, filenames in os.walk(root_path):
-        for filename in filenames:
-            file_path = os.path.join(root, filename)
-            files.append(file_path)
-    return files
 
 
-def encrypt(fernet, root_path: str = "./", output_dir: str = "./encrypted_data"):
-    try:
-        files = get_all_files(root_path)
-        if not files:
-            return False
-        
-        os.makedirs(output_dir, exist_ok=True)
-        
-        for file_path in files:
-            try:
-                with open(file_path, 'rb') as f:
-                    original = f.read()
-                encrypted = fernet.encrypt(original)
-                relative_path = os.path.relpath(file_path, root_path)
-                output_path = os.path.join(output_dir, relative_path)
-                os.makedirs(os.path.dirname(output_path), exist_ok=True)
-                with open(output_path, 'wb') as f:
-                    f.write(encrypted)
-            except Exception:
-                continue
-        return True
-    except Exception:
-        return False
-
-
-def save_key(key, env_path: str = ".env"):
-    try:
-        if os.path.exists(env_path):
-            with open(env_path, "r") as f:
-                lines = f.readlines()
-        else:
-            lines = []
-        
-        if lines and not lines[-1].endswith('\n'):
-            lines.append('\n')
-        lines.append(f"\nSECRET_KEY='{key.decode()}'\n")
-        
-        with open(env_path, "w") as f:
-            f.writelines(lines)
-        return True
-    except Exception:
-        return False
-
+# # ============================================================================
+# # UNIT TESTS
+# # ============================================================================
 
-def decrypt(fernet, encrypted_path: str, output_path: str):
-    try:
-        with open(encrypted_path, 'rb') as f:
-            encrypted = f.read()
-        decrypted = fernet.decrypt(encrypted)
-        os.makedirs(os.path.dirname(output_path), exist_ok=True)
-        with open(output_path, 'wb') as f:
-            f.write(decrypted)
-        return True
-    except Exception:
-        return False
-
-
-# ============================================================================
-# UNIT TESTS
-# ============================================================================
-
-class TestGetKey:
-    """Unit tests for get_key function"""
+# class TestGetKey:
+#     """Unit tests for get_key function"""
     
-    def test_get_key_returns_tuple(self):
-        """Test that get_key returns a tuple"""
-        result = get_key()
-        assert isinstance(result, tuple)
-        assert len(result) == 2
+#     def test_get_key_returns_tuple(self):
+#         """Test that get_key returns a tuple"""
+#         result = get_key()
+#         assert isinstance(result, tuple)
+#         assert len(result) == 2
     
-    def test_get_key_returns_valid_key(self):
-        """Test that get_key returns valid Fernet key"""
-        key, fernet = get_key()
-        assert isinstance(key, bytes)
-        assert isinstance(fernet, Fernet)
+#     def test_get_key_returns_valid_key(self):
+#         """Test that get_key returns valid Fernet key"""
+#         key, fernet = get_key()
+#         assert isinstance(key, bytes)
+#         assert isinstance(fernet, Fernet)
     
-    def test_key_can_encrypt_decrypt(self):
-        """Test that generated key can encrypt and decrypt"""
-        key, fernet = get_key()
-        message = b"Test message"
-        encrypted = fernet.encrypt(message)
-        decrypted = fernet.decrypt(encrypted)
-        assert decrypted == message
+#     def test_key_can_encrypt_decrypt(self):
+#         """Test that generated key can encrypt and decrypt"""
+#         key, fernet = get_key()
+#         message = b"Test message"
+#         encrypted = fernet.encrypt(message)
+#         decrypted = fernet.decrypt(encrypted)
+#         assert decrypted == message
     
-    def test_keys_are_unique(self):
-        """Test that each call generates a unique key"""
-        key1, _ = get_key()
-        key2, _ = get_key()
-        assert key1 != key2
+#     def test_keys_are_unique(self):
+#         """Test that each call generates a unique key"""
+#         key1, _ = get_key()
+#         key2, _ = get_key()
+#         assert key1 != key2
 
 
-class TestGetAllFiles:
-    """Unit tests for get_all_files function"""
-    
-    @pytest.fixture
-    def temp_dir(self):
-        """Create temporary directory with test files"""
-        temp_dir = tempfile.mkdtemp()
-        
-        # Create test files
-        Path(temp_dir, "file1.txt").write_text("content1")
-        Path(temp_dir, "file2.txt").write_text("content2")
-        
-        # Create subdirectory with file
-        subdir = Path(temp_dir, "subdir")
-        subdir.mkdir()
-        Path(subdir, "file3.txt").write_text("content3")
-        
-        yield temp_dir
-        
-        # Cleanup
-        shutil.rmtree(temp_dir)
-    
-    def test_finds_all_files(self, temp_dir):
-        """Test that all files are found"""
-        files = get_all_files(temp_dir)
-        assert len(files) == 3
-    
-    def test_empty_directory(self):
-        """Test with empty directory"""
-        with tempfile.TemporaryDirectory() as temp_dir:
-            files = get_all_files(temp_dir)
-            assert files == []
-    
-    def test_returns_list(self, temp_dir):
-        """Test that function returns a list"""
-        files = get_all_files(temp_dir)
-        assert isinstance(files, list)
-
 
-class TestSaveKey:
-    """Unit tests for save_key function"""
+# class TestSaveKey:
+#     """Unit tests for save_key function"""
     
-    @pytest.fixture
-    def temp_env_file(self):
-        """Create temporary .env file"""
-        fd, path = tempfile.mkstemp(suffix=".env")
-        os.close(fd)
-        yield path
-        if os.path.exists(path):
-            os.remove(path)
+#     @pytest.fixture
+#     def temp_env_file(self):
+#         """Create temporary .env file"""
+#         fd, path = tempfile.mkstemp(suffix=".env")
+#         os.close(fd)
+#         yield path
+#         if os.path.exists(path):
+#             os.remove(path)
     
-    def test_save_key_creates_file(self, temp_env_file):
-        """Test that save_key creates .env file"""
-        os.remove(temp_env_file)  # Remove to test creation
-        key, _ = get_key()
-        result = save_key(key, temp_env_file)
-        assert result is True
-        assert os.path.exists(temp_env_file)
+#     def test_save_key_creates_file(self, temp_env_file):
+#         """Test that save_key creates .env file"""
+#         os.remove(temp_env_file)  # Remove to test creation
+#         key, _ = get_key()
+#         result = save_key(key, temp_env_file)
+#         assert result is True
+#         assert os.path.exists(temp_env_file)
     
-    def test_save_key_appends_to_existing(self, temp_env_file):
-        """Test that save_key appends to existing file"""
-        # Write initial content
-        with open(temp_env_file, "w") as f:
-            f.write("EXISTING_VAR=value\n")
+#     def test_save_key_appends_to_existing(self, temp_env_file):
+#         """Test that save_key appends to existing file"""
+#         # Write initial content
+#         with open(temp_env_file, "w") as f:
+#             f.write("EXISTING_VAR=value\n")
         
-        key, _ = get_key()
-        save_key(key, temp_env_file)
+#         key, _ = get_key()
+#         save_key(key, temp_env_file)
         
-        with open(temp_env_file, "r") as f:
-            content = f.read()
+#         with open(temp_env_file, "r") as f:
+#             content = f.read()
         
-        assert "EXISTING_VAR=value" in content
-        assert "SECRET_KEY=" in content
+#         assert "EXISTING_VAR=value" in content
+#         assert "SECRET_KEY=" in content
     
-    def test_saved_key_format(self, temp_env_file):
-        """Test that saved key has correct format"""
-        key, _ = get_key()
-        save_key(key, temp_env_file)
+#     def test_saved_key_format(self, temp_env_file):
+#         """Test that saved key has correct format"""
+#         key, _ = get_key()
+#         save_key(key, temp_env_file)
         
-        with open(temp_env_file, "r") as f:
-            content = f.read()
+#         with open(temp_env_file, "r") as f:
+#             content = f.read()
         
-        assert f"SECRET_KEY='{key.decode()}'" in content
+#         assert f"SECRET_KEY='{key.decode()}'" in content
 
 
-class TestDecrypt:
-    """Unit tests for decrypt function"""
+# class TestDecrypt:
+#     """Unit tests for decrypt function"""
     
-    @pytest.fixture
-    def encrypted_file(self):
-        """Create temporary encrypted file"""
-        key, fernet = get_key()
-        fd, path = tempfile.mkstemp()
+#     @pytest.fixture
+#     def encrypted_file(self):
+#         """Create temporary encrypted file"""
+#         key, fernet = get_key()
+#         fd, path = tempfile.mkstemp()
         
-        content = b"Test content to encrypt"
-        encrypted = fernet.encrypt(content)
+#         content = b"Test content to encrypt"
+#         encrypted = fernet.encrypt(content)
         
-        with os.fdopen(fd, 'wb') as f:
-            f.write(encrypted)
+#         with os.fdopen(fd, 'wb') as f:
+#             f.write(encrypted)
         
-        yield path, fernet, content
+#         yield path, fernet, content
         
-        if os.path.exists(path):
-            os.remove(path)
+#         if os.path.exists(path):
+#             os.remove(path)
     
-    def test_decrypt_success(self, encrypted_file):
-        """Test successful decryption"""
-        encrypted_path, fernet, original_content = encrypted_file
+#     # def test_decrypt_success(self, encrypted_file):
+#     #     """Test successful decryption"""
+#     #     encrypted_path, fernet, original_content = encrypted_file
         
-        with tempfile.NamedTemporaryFile(delete=False) as output:
-            output_path = output.name
+#     #     with tempfile.NamedTemporaryFile(delete=False) as output:
+#     #         output_path = output.name
         
-        try:
-            result = decrypt(fernet, encrypted_path, output_path)
-            assert result is True
+#     #     try:
+#     #         result = decrypt(fernet, encrypted_path, output_path)
+#     #         assert result is True
             
-            with open(output_path, 'rb') as f:
-                decrypted = f.read()
+#     #         with open(output_path, 'rb') as f:
+#     #             decrypted = f.read()
             
-            assert decrypted == original_content
-        finally:
-            if os.path.exists(output_path):
-                os.remove(output_path)
+#     #         assert decrypted == original_content
+#     #     finally:
+#     #         if os.path.exists(output_path):
+#     #             os.remove(output_path)
     
-    def test_decrypt_wrong_key(self, encrypted_file):
-        """Test decryption with wrong key fails"""
-        encrypted_path, _, _ = encrypted_file
-        _, wrong_fernet = get_key()  # Different key
-        
-        with tempfile.NamedTemporaryFile(delete=False) as output:
-            output_path = output.name
-        
-        try:
-            result = decrypt(wrong_fernet, encrypted_path, output_path)
-            assert result is False
-        finally:
-            if os.path.exists(output_path):
-                os.remove(output_path)
+#     def test_decrypt_wrong_key(self, encrypted_file):
+#         """Test decryption with wrong key fails"""
+#         encrypted_path, _, _ = encrypted_file
+#         _, wrong_fernet = get_key()  # Different key
+        
+#         with tempfile.NamedTemporaryFile(delete=False) as output:
+#             output_path = output.name
+        
+#         try:
+#             result = decrypt(wrong_fernet, encrypted_path, output_path)
+#             assert result is False
+#         finally:
+#             if os.path.exists(output_path):
+#                 os.remove(output_path)
 
 
-# ============================================================================
-# INTEGRATION TESTS
-# ============================================================================
+# # ============================================================================
+# # INTEGRATION TESTS
+# # ============================================================================
 
-class TestEncryptionWorkflow:
-    """Integration tests for complete encryption workflow"""
+# class TestEncryptionWorkflow:
+#     """Integration tests for complete encryption workflow"""
     
-    @pytest.fixture
-    def test_environment(self):
-        """Create complete test environment"""
-        # Create temporary directories
-        source_dir = tempfile.mkdtemp()
-        output_dir = tempfile.mkdtemp()
-        env_file = tempfile.NamedTemporaryFile(delete=False, suffix=".env")
-        env_file.close()
-        
-        # Create test files
-        Path(source_dir, "test1.txt").write_text("Content 1")
-        Path(source_dir, "test2.csv").write_text("a,b,c\n1,2,3")
-        
-        subdir = Path(source_dir, "subdir")
-        subdir.mkdir()
-        Path(subdir, "test3.json").write_text('{"key": "value"}')
-        
-        yield source_dir, output_dir, env_file.name
-        
-        # Cleanup
-        shutil.rmtree(source_dir, ignore_errors=True)
-        shutil.rmtree(output_dir, ignore_errors=True)
-        if os.path.exists(env_file.name):
-            os.remove(env_file.name)
+#     @pytest.fixture
+#     def test_environment(self):
+#         """Create complete test environment"""
+#         # Create temporary directories
+#         source_dir = tempfile.mkdtemp()
+#         output_dir = tempfile.mkdtemp()
+#         env_file = tempfile.NamedTemporaryFile(delete=False, suffix=".env")
+#         env_file.close()
+        
+#         # Create test files
+#         Path(source_dir, "test1.txt").write_text("Content 1")
+#         Path(source_dir, "test2.csv").write_text("a,b,c\n1,2,3")
+        
+#         subdir = Path(source_dir, "subdir")
+#         subdir.mkdir()
+#         Path(subdir, "test3.json").write_text('{"key": "value"}')
+        
+#         yield source_dir, output_dir, env_file.name
+        
+#         # Cleanup
+#         shutil.rmtree(source_dir, ignore_errors=True)
+#         shutil.rmtree(output_dir, ignore_errors=True)
+#         if os.path.exists(env_file.name):
+#             os.remove(env_file.name)
     
-    def test_full_encryption_workflow(self, test_environment):
-        """Test complete encryption workflow"""
-        source_dir, output_dir, env_file = test_environment
+#     def test_full_encryption_workflow(self, test_environment):
+#         """Test complete encryption workflow"""
+#         source_dir, output_dir, env_file = test_environment
         
-        # Generate key
-        key, fernet = get_key()
+#         # Generate key
+#         key, fernet = get_key()
         
-        # Encrypt files
-        result = encrypt(fernet, source_dir, output_dir)
-        assert result is True
+#         # Encrypt files
+#         result = encrypt(fernet, source_dir, output_dir)
+#         assert result is True
         
-        # Verify encrypted files exist
-        encrypted_files = get_all_files(output_dir)
-        assert len(encrypted_files) == 3
+#         # Verify encrypted files exist
+#         encrypted_files = get_all_files(output_dir)
+#         assert len(encrypted_files) == 3
         
-        # Save key
-        result = save_key(key, env_file)
-        assert result is True
+#         # Save key
+#         result = save_key(key, env_file)
+#         assert result is True
         
-        # Verify key was saved
-        with open(env_file, "r") as f:
-            content = f.read()
-        assert "SECRET_KEY=" in content
+#         # Verify key was saved
+#         with open(env_file, "r") as f:
+#             content = f.read()
+#         assert "SECRET_KEY=" in content
     
-    def test_encrypt_decrypt_roundtrip(self, test_environment):
-        """Test that files can be encrypted and then decrypted"""
-        source_dir, output_dir, _ = test_environment
+#     def test_encrypt_decrypt_roundtrip(self, test_environment):
+#         """Test that files can be encrypted and then decrypted"""
+#         source_dir, output_dir, _ = test_environment
         
-        # Generate key
-        key, fernet = get_key()
+#         # Generate key
+#         key, fernet = get_key()
         
-        # Encrypt
-        encrypt(fernet, source_dir, output_dir)
+#         # Encrypt
+#         encrypt(fernet, source_dir, output_dir)
         
-        # Decrypt one file
-        encrypted_file = os.path.join(output_dir, "test1.txt")
-        decrypted_dir = tempfile.mkdtemp()
-        decrypted_file = os.path.join(decrypted_dir, "test1.txt")
+#         # Decrypt one file
+#         encrypted_file = os.path.join(output_dir, "test1.txt")
+#         decrypted_dir = tempfile.mkdtemp()
+#         decrypted_file = os.path.join(decrypted_dir, "test1.txt")
         
-        try:
-            decrypt(fernet, encrypted_file, decrypted_file)
+#         try:
+#             decrypt(fernet, encrypted_file, decrypted_file)
             
-            # Verify content matches original
-            with open(decrypted_file, 'r') as f:
-                content = f.read()
-            assert content == "Content 1"
-        finally:
-            shutil.rmtree(decrypted_dir, ignore_errors=True)
+#             # Verify content matches original
+#             with open(decrypted_file, 'r') as f:
+#                 content = f.read()
+#             assert content == "Content 1"
+#         finally:
+#             shutil.rmtree(decrypted_dir, ignore_errors=True)
     
-    def test_preserves_directory_structure(self, test_environment):
-        """Test that directory structure is preserved"""
-        source_dir, output_dir, _ = test_environment
+#     def test_preserves_directory_structure(self, test_environment):
+#         """Test that directory structure is preserved"""
+#         source_dir, output_dir, _ = test_environment
         
-        key, fernet = get_key()
-        encrypt(fernet, source_dir, output_dir)
+#         key, fernet = get_key()
+#         encrypt(fernet, source_dir, output_dir)
         
-        # Check that subdirectory exists
-        encrypted_subdir_file = os.path.join(output_dir, "subdir", "test3.json")
-        assert os.path.exists(encrypted_subdir_file)
+#         # Check that subdirectory exists
+#         encrypted_subdir_file = os.path.join(output_dir, "subdir", "test3.json")
+#         assert os.path.exists(encrypted_subdir_file)
     
-    def test_empty_directory_handling(self):
-        """Test handling of empty directory"""
-        with tempfile.TemporaryDirectory() as source_dir:
-            with tempfile.TemporaryDirectory() as output_dir:
-                key, fernet = get_key()
-                result = encrypt(fernet, source_dir, output_dir)
-                assert result is False
+#     def test_empty_directory_handling(self):
+#         """Test handling of empty directory"""
+#         with tempfile.TemporaryDirectory() as source_dir:
+#             with tempfile.TemporaryDirectory() as output_dir:
+#                 key, fernet = get_key()
+#                 result = encrypt(fernet, source_dir, output_dir)
+#                 assert result is False
 
 
-# ============================================================================
-# PARAMETRIZED TESTS
-# ============================================================================
+# # ============================================================================
+# # PARAMETRIZED TESTS
+# # ============================================================================
 
-class TestEncryptionWithDifferentFileTypes:
-    """Test encryption with various file types"""
+# class TestEncryptionWithDifferentFileTypes:
+#     """Test encryption with various file types"""
     
-    @pytest.mark.parametrize("filename,content", [
-        ("test.txt", b"Plain text content"),
-        ("test.json", b'{"key": "value"}'),
-        ("test.csv", b"a,b,c\n1,2,3"),
-        ("test.bin", bytes(range(256))),
-        ("test.pdf", b"%PDF-1.4\n%\xe2\xe3\xcf\xd3"),
-    ])
-    def test_encrypt_different_file_types(self, filename, content):
-        """Test encryption of different file types"""
-        with tempfile.TemporaryDirectory() as source_dir:
-            with tempfile.TemporaryDirectory() as output_dir:
-                # Create test file
-                file_path = Path(source_dir, filename)
-                file_path.write_bytes(content)
+#     @pytest.mark.parametrize("filename,content", [
+#         ("test.txt", b"Plain text content"),
+#         ("test.json", b'{"key": "value"}'),
+#         ("test.csv", b"a,b,c\n1,2,3"),
+#         ("test.bin", bytes(range(256))),
+#         ("test.pdf", b"%PDF-1.4\n%\xe2\xe3\xcf\xd3"),
+#     ])
+#     def test_encrypt_different_file_types(self, filename, content):
+#         """Test encryption of different file types"""
+#         with tempfile.TemporaryDirectory() as source_dir:
+#             with tempfile.TemporaryDirectory() as output_dir:
+#                 # Create test file
+#                 file_path = Path(source_dir, filename)
+#                 file_path.write_bytes(content)
                 
-                # Encrypt
-                key, fernet = get_key()
-                result = encrypt(fernet, source_dir, output_dir)
-                assert result is True
+#                 # Encrypt
+#                 key, fernet = get_key()
+#                 result = encrypt(fernet, source_dir, output_dir)
+#                 assert result is True
                 
-                # Verify encrypted file exists
-                encrypted_path = Path(output_dir, filename)
-                assert encrypted_path.exists()
+#                 # Verify encrypted file exists
+#                 encrypted_path = Path(output_dir, filename)
+#                 assert encrypted_path.exists()
                 
-                # Verify content is different (encrypted)
-                encrypted_content = encrypted_path.read_bytes()
-                assert encrypted_content != content
+#                 # Verify content is different (encrypted)
+#                 encrypted_content = encrypted_path.read_bytes()
+#                 assert encrypted_content != content
 
 
-if __name__ == "__main__":
-    pytest.main([__file__, "-v", "--tb=short"])
+# if __name__ == "__main__":
+#     pytest.main([__file__, "-v", "--tb=short"])

tests/test_restrict_usage.py

@@ -0,0 +1,81 @@
+import pytest
+from datetime import datetime, timedelta
+from unittest.mock import patch
+import sys
+import os
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+from agent.restrict_usage import RateLimiter
+
+
+@pytest.fixture
+def limiter():
+    """Create a RateLimiter instance for testing."""
+    return RateLimiter(max_requests=5, window_minutes=60)
+
+
+def test_first_request_allowed(limiter):
+    """First request should always be allowed."""
+    assert limiter.is_allowed("test_ip") is True
+
+
+def test_requests_within_limit(limiter):
+    """All requests within limit should be allowed."""
+    for _ in range(5):
+        assert limiter.is_allowed("test_ip") is True
+
+
+def test_request_exceeds_limit(limiter):
+    """Request exceeding limit should be blocked."""
+    for _ in range(5):
+        limiter.is_allowed("test_ip")
+    assert limiter.is_allowed("test_ip") is False
+
+
+def test_different_identifiers_separate(limiter):
+    """Different identifiers should have separate limits."""
+    for _ in range(5):
+        limiter.is_allowed("ip1")
+    assert limiter.is_allowed("ip1") is False
+    assert limiter.is_allowed("ip2") is True
+
+
+def test_get_remaining(limiter):
+    """get_remaining should return correct count."""
+    assert limiter.get_remaining("test_ip") == 5
+    limiter.is_allowed("test_ip")
+    assert limiter.get_remaining("test_ip") == 4
+    for _ in range(4):
+        limiter.is_allowed("test_ip")
+    assert limiter.get_remaining("test_ip") == 0
+
+
+def test_old_requests_cleaned(limiter):
+    """Old requests outside window should be cleaned."""
+    fixed_time = datetime(2024, 1, 1, 12, 0, 0)
+    
+    with patch('agent.restrict_usage.datetime') as mock_datetime:
+        mock_datetime.now.return_value = fixed_time
+        
+        for _ in range(5):
+            limiter.is_allowed("test_ip")
+        
+        assert limiter.is_allowed("test_ip") is False
+        
+        # Move time forward past the window
+        mock_datetime.now.return_value = fixed_time + timedelta(minutes=61)
+        
+        assert limiter.is_allowed("test_ip") is True
+        assert limiter.get_remaining("test_ip") == 4
+
+
+def test_multiple_identifiers_tracking(limiter):
+    """Multiple identifiers should be tracked independently."""
+    limiter.is_allowed("user1")
+    limiter.is_allowed("user1")
+    limiter.is_allowed("user2")
+    
+    assert limiter.get_remaining("user1") == 3
+    assert limiter.get_remaining("user2") == 4
+    assert limiter.get_remaining("user3") == 5