security.py

Overview

security.py provides functionality to analyze source code for potentially unsafe or dangerous constructs, focusing primarily on Python code. It leverages Python's Abstract Syntax Tree (AST) capabilities to perform static code analysis and detect risky code patterns such as dangerous imports, function calls, and code constructs that may lead to security vulnerabilities or unsafe execution.

The primary component is the SecurePythonAnalyzer class, which extends ast.NodeVisitor to traverse the AST nodes of Python code and flag suspicious elements. The module also exposes a utility function analyze_code_security to analyze given code snippets in supported languages, currently only Python, returning safety status and detailed findings.

This module is intended to be part of a larger system where code safety and security are critical, such as a code execution platform or automated code review system (e.g., InfiniFlow). It helps prevent execution of malicious or harmful code by identifying unsafe constructs before runtime.

Classes and Functions

Class: SecurePythonAnalyzer

An AST-based analyzer designed to detect unsafe Python code patterns by visiting nodes in the AST and flagging dangerous imports, calls, and other suspicious constructs.

Constants

• DANGEROUS_IMPORTS: Set[str]
  A set of module names considered risky when imported. Examples: "os", "subprocess", "pickle", "ctypes", etc.

• DANGEROUS_CALLS: Set[str]
  Function or method names that represent potentially harmful operations. Examples: "eval", "exec", "open", "system", "pickle.load", etc.

Attributes

• unsafe_items: List[Tuple[str, int]]
  Stores tuples of (description, line number) for every detected unsafe code element.

Methods

All methods override ast.NodeVisitor visit methods to analyze different AST node types.

• init(self)
  Initializes the analyzer with an empty list of findings.

• visit_Import(self, node: ast.Import)
  Checks all imported modules against DANGEROUS_IMPORTS. Records any matches.

• visit_ImportFrom(self, node: ast.ImportFrom)
  Checks from ... import ... statements for dangerous modules.

• visit_Call(self, node: ast.Call)
  Detects calls to dangerous functions by checking the function name.

• visit_Attribute(self, node: ast.Attribute)
  Inspect attribute access to detect calls like os.system.

• visit_BinOp(self, node: ast.BinOp)
  Flags suspicious binary operations, especially string concatenations that could form dangerous commands.

• visit_FunctionDef(self, node: ast.FunctionDef)
  Detects if a function is defined with a dangerous name (e.g., redefining eval).

• visit_Assign(self, node: ast.Assign)
  Checks assignment targets for dangerous variable names.

• visit_Lambda(self, node: ast.Lambda)
  Detects lambda expressions invoking dangerous functions.

• visit_ListComp(self, node: ast.ListComp)
  Detects list comprehensions involving dangerous calls.

• visit_DictComp(self, node: ast.DictComp)
  Detects dictionary comprehensions involving dangerous calls in keys or values.

• visit_SetComp(self, node: ast.SetComp)
  Detects set comprehensions involving dangerous calls.

• visit_Yield(self, node: ast.Yield)
  Checks yield statements that return results of dangerous function calls.

Usage Example

import ast

code = """
import os
def safe_func():
    print("Hello")

def eval(x):
    return x
"""

analyzer = SecurePythonAnalyzer()
tree = ast.parse(code)
analyzer.visit(tree)

print(analyzer.unsafe_items)
# Output might be:
# [('Import: os', 2), ('Function Definition: eval', 6)]

Function: analyze_code_security

def analyze_code_security(code: str, language: SupportLanguage) -> Tuple[bool, List[Tuple[str, int]]]:

Analyzes the provided source code string and returns whether it is considered safe along with a list of detected issues.

Parameters

• code (str): Source code to analyze.

• language (SupportLanguage): Enum value representing the programming language of the source code.

Returns

• Tuple[bool, List[Tuple[str, int]]]
  A tuple where the first element is a boolean indicating if the code is safe (True if no unsafe items detected), and the second element is a list of tuples describing unsafe elements with their line numbers.

Behavior

• For Python code (SupportLanguage.PYTHON):

  • Parses the code into an AST.

  • Uses SecurePythonAnalyzer to visit the AST nodes.

  • Returns True if no unsafe items found, otherwise False with list of issues.

  • Catches parsing errors and returns them as unsafe issues.

• For unsupported languages:

  • Logs a warning.

  • Returns True with a note that manual review is recommended.

Usage Example

from models.enums import SupportLanguage

code_snippet = """
import subprocess
subprocess.call('rm -rf /')
"""

is_safe, issues = analyze_code_security(code_snippet, SupportLanguage.PYTHON)
print(is_safe)  # False
print(issues)   # [('Import: subprocess', 2), ('Call: call', 3)]

Important Implementation Details

• The analysis is purely static and based on AST node inspection without executing the code.

• The DANGEROUS_IMPORTS and DANGEROUS_CALLS sets are curated to cover modules and functions commonly associated with security risks such as arbitrary code execution, file system access, and network operations.

• The analyzer visits a broad range of AST nodes (imports, calls, function definitions, assignments, comprehensions, lambdas, yields) to catch various ways unsafe code can manifest.

• For function calls and attribute accesses, the analyzer matches exact names and simple attribute chains (e.g., subprocess.call).

• The module currently only supports Python code for detailed analysis. Other languages are marked as safe by default but flagged for manual review.

• The module logs errors and warnings via a centralized logger from core.logger, integrating with the application's logging infrastructure.

Integration and Interactions

• With models.enums.SupportLanguage: Uses the SupportLanguage enum to identify the language of the input code and decide the analysis approach.

• With core.logger: Uses logging to report errors and warnings during analysis.

• As part of a larger code execution or review system: This module can be integrated into pipelines where user-provided or dynamically loaded code must be vetted for safety before execution.

• Potentially interacts with other security modules or policy enforcement components by providing verdicts and detailed findings on code safety.

Structure Diagram

classDiagram
    class SecurePythonAnalyzer {
        -DANGEROUS_IMPORTS: set
        -DANGEROUS_CALLS: set
        -unsafe_items: List~Tuple~str,int~~
        +__init__()
        +visit_Import(node)
        +visit_ImportFrom(node)
        +visit_Call(node)
        +visit_Attribute(node)
        +visit_BinOp(node)
        +visit_FunctionDef(node)
        +visit_Assign(node)
        +visit_Lambda(node)
        +visit_ListComp(node)
        +visit_DictComp(node)
        +visit_SetComp(node)
        +visit_Yield(node)
    }

    class analyze_code_security {
        +analyze_code_security(code: str, language: SupportLanguage) -> Tuple~bool, List~Tuple~str, int~~~
    }

    analyze_code_security ..> SecurePythonAnalyzer : uses

Summary

The security.py module is a critical security component that performs static analysis on Python code to detect potentially unsafe operations. Its design leverages Python's AST for deep inspection and flags dangerous imports, calls, and constructs. The module provides a clean and extensible way to integrate code safety checks into larger systems, helping mitigate risks associated with executing untrusted code.