Updated CodeSafe with new allow_attributes and set immediate_termination to be True by default for safer function calling.

Author: Infinitode

README.md

@@ -12,6 +12,16 @@ An open-source Python library for code encryption, decryption, and safe evaluati
 > [!NOTE]
 > **CodeSafe** is intended to quickly encrypt/decrypt code files, and run them (only for Python script files) while in their encrypted form, but not as a means for powerful encryption, just code obfuscation. We have also included a `safe_eval` function, that can safely evaluate expressions within a safe environment.
 
+### Changelog v0.0.3:
+- Added an `allow_attributes` parameter to `safe_eval` and set `immediate_termination` to be `True` by default for safer function calling.
+
+### Changelog v0.0.2:
+- Fixed function returns.
+- Added error handling to `CodeSafe`, removed some print statements with edits from `@0XC7R`.
+
+### Changelog v0.0.1:
+- Initial release
+
 ## Installation
 
 You can install CodeSafe using pip:

codesafe/__init__.py

@@ -1,18 +1,12 @@
 import ast
 import builtins
 from multiprocessing import Process, Queue
+import gc
 import base64
 import re
-import gc
 
 __all__ = ['safe_eval', 'EvaluationTimeoutError', 'UnsafeExpressionError', 'encrypt_to_file', 'encrypt', 'decrypt', 'run', 'decrypt_to_file']
 
-def cleanup_resources():
-    """
-    Clean up all open resources by closing them.
-    """
-    gc.collect()  # Force garbage collection
-
 class EvaluationTimeoutError(Exception):
     """Custom exception to handle timeouts during evaluation."""
     def __init__(self, error):
@@ -34,12 +28,13 @@ def _eval_in_process(expr: str, safe_globals: dict, queue: Queue):
         queue.put(e)
 
 def safe_eval(expr: str,
-              allowed_builtins: dict = None,
-              allowed_vars: dict = None,
+              allowed_builtins: dict = {},
+              allowed_vars: dict = {},
               timeout: float = 5,
-              restricted_imports: list = None,
-              allowed_function_calls: list = None,
-              immediate_termination: bool = False,
+              restricted_imports: list = [],
+              allowed_function_calls: list = [],
+              allow_attributes: bool = False,
+              immediate_termination: bool = True,
               file_access: bool = False,
               network_access: bool = False) -> object:
     """
@@ -52,7 +47,8 @@ def safe_eval(expr: str,
         timeout (float, optional): Time limit for evaluation in seconds. Defaults to 5.
         restricted_imports (list, optional): A list of restricted imports or modules. Defaults to [].
         allowed_function_calls (list, optional): A list of allowed function names to call. Defaults to [].
-        immediate_termination (bool, optional): Whether to forcibly terminate the evaluation if it exceeds the timeout. Defaults to False.
+        allow_attributes (bool, optional): Whether to allow access to safe attributes and methods (e.g., 'str.upper()'). Defaults to False.
+        immediate_termination (bool, optional): Whether to forcibly terminate the evaluation if it exceeds the timeout. Defaults to True.
         file_access (bool, optional): Whether to allow file access (open, etc.). Defaults to False.
         network_access (bool, optional): Whether to allow network access (requests, etc.). Defaults to False.
 
@@ -65,21 +61,16 @@ def safe_eval(expr: str,
         UnsafeExpressionError: If restricted imports or unsafe nodes are detected in the AST.
         SyntaxError: If the expression contains invalid syntax.
     """
-    allowed_builtins = allowed_builtins or {}
-    allowed_vars = allowed_vars or {}
-    restricted_imports = restricted_imports or []
-    allowed_function_calls = allowed_function_calls or []
 
     # Restrict file access by removing file-related functions from built-ins if file_access is False
-    safe_builtins = allowed_builtins.copy()
+    safe_builtins = {k: v for k, v in builtins.__dict__.items()}
+    safe_builtins.update(allowed_builtins)
+
     if not file_access:
         # Remove all file-related functions from built-ins
-        file_access_functions = {'open', 'os.remove', 'os.rename', 'os.rmdir', 'os.unlink',
-                                 'os.mkdir', 'os.makedirs', 'os.chmod', 'os.chown',
-                                 'os.truncate', 'os.path', 'shutil', 'pathlib', 'subprocess'}
-
-        # Filter out functions related to file access
-        safe_builtins = {k: v for k, v in builtins.__dict__.items() if k not in file_access_functions}
+        file_funcs = {'open'}
+        for func in file_funcs:
+            safe_builtins.pop(func, None)
 
     if not network_access:
         # Remove networking-related functions from built-ins
@@ -98,10 +89,9 @@ def safe_eval(expr: str,
     try:
         parsed_expr = ast.parse(expr, mode='eval')
     except SyntaxError as e:
-        cleanup_resources()
         raise SyntaxError(f"Invalid syntax: {e}")
 
-    _check_ast(parsed_expr, restricted_imports, allowed_function_calls)
+    _check_ast(parsed_expr, restricted_imports, allowed_function_calls, allow_attributes)
 
     queue = Queue()
     process = Process(target=_eval_in_process, args=(expr, safe_globals, queue))
@@ -113,36 +103,40 @@ def safe_eval(expr: str,
         if immediate_termination:
             process.terminate()  # Terminate the process immediately
             process.join()  # Ensure the process has finished
-            cleanup_resources()
+            gc.collect()
             raise EvaluationTimeoutError("Evaluation timed out and was terminated.")
         else:
-            cleanup_resources()
+            gc.collect()
             raise EvaluationTimeoutError("Evaluation timed out.")
 
     # Check for results or exceptions
     if not queue.empty():
         result = queue.get()
         if isinstance(result, Exception):
-            cleanup_resources()
+            gc.collect()
             raise result
     else:
-        cleanup_resources()
+        gc.collect()
         raise EvaluationTimeoutError("No result returned from the evaluation process.")
 
     return result
 
-def _check_ast(parsed_expr, restricted_imports, allowed_function_calls):
+def _check_ast(parsed_expr, restricted_imports, allowed_function_calls, allow_attributes):
     """
     Check the AST for unsafe operations such as imports and function calls.
 
     Parameters:
         parsed_expr (ast.AST): The parsed AST expression.
         restricted_imports (list): A list of restricted import statements.
         allowed_function_calls (list): A list of allowed function calls.
+        allow_attributes (bool): Whether to allow attribute access.
 
     Raises:
         UnsafeExpressionError: If unsafe operations are detected in the expression.
     """
+    # Attributes starting with '__' are always blocked if attributes are allowed.
+    blocked_attrs = {'__globals__', '__closure__', '__code__', '__subclasses__', '__init__', '__class__', '__bases__'}
+
     for node in ast.walk(parsed_expr):
         # Block all imports if restricted
         if isinstance(node, (ast.Import, ast.ImportFrom)):
@@ -157,8 +151,11 @@ def _check_ast(parsed_expr, restricted_imports, allowed_function_calls):
                 raise UnsafeExpressionError(f"Function call to '{node.func.id}' is not allowed.")
 
         # Prevent attribute access (e.g., accessing os.system or other potentially harmful attributes)
-        if isinstance(node, ast.Attribute):
-            raise UnsafeExpressionError("Attribute access is restricted for security reasons.")
+        if isinstance(node, ast.Attribute) and not allow_attributes:
+            raise UnsafeExpressionError("Attribute access is disabled.")
+        elif isinstance(node, ast.Attribute) and allow_attributes:
+            if node.attr.startswith('__') and node.attr in blocked_attrs:
+                raise UnsafeExpressionError(f"Access to dunder attribute '{node.attr}' is not allowed.")
 
 # Custom key (visible but obfuscated)
 SHIFT = 3  # Caesar cipher shift value

test.py

@@ -48,6 +48,24 @@ def test_custom_sum(self):
 
         self.assertEqual(result, 30)
 
+    def test_safe_attribute_access(self):
+        expression = "'hello'.upper()"
+        result = safe_eval(expression, allow_attributes=True)
+        self.assertEqual(result, "HELLO")
+
+    def test_unsafe_attribute_access(self):
+        expression = "().__class__.__bases__[0]"
+        with self.assertRaises(UnsafeExpressionError):
+            safe_eval(expression, allow_attributes=True)
+
+    def test_attribute_access_disabled(self):
+        expression = "'hello'.upper()"
+        with self.assertRaises(UnsafeExpressionError):
+            safe_eval(
+                expression,
+                allow_attributes=False  # Explicitly disable
+            )
+
     def test_security_error_on_unsafe_function(self):
         expression = "os.getcwd()"  # Attempting to call an unsafe function
         try: