Plugin Development

def checks(*args):
    """
    Specify which AST node types the test should analyze.
    
    Parameters:
    - *args: str, AST node type names ('Call', 'Import', 'Str', etc.)
    
    Usage:
    @checks('Call', 'Attribute') - Run test on function calls and attribute access
    @checks('Import', 'ImportFrom') - Run test on import statements
    """

def test_id(id_val):
    """
    Assign unique identifier to security test.
    
    Parameters:
    - id_val: str, unique test identifier (e.g., 'B101', 'B999')
    
    Usage:
    @test_id('B999') - Assign test ID B999
    """

def takes_config(name=None):
    """
    Indicate test accepts configuration data.
    
    Parameters:
    - name: str, configuration section name (optional)
    
    Usage:
    @takes_config('hardcoded_password') - Access config['hardcoded_password']
    @takes_config() - Access general test configuration
    """

def accepts_baseline(*args):
    """
    Mark formatter as supporting baseline results.
    Used for output formatters that can handle baseline comparison.
    
    Parameters:
    - *args: Additional baseline configuration options
    """

# Assert and Debug Tests
B101 = "assert_used"                    # Use of assert detected
B201 = "flask_debug_true"               # Flask app run in debug mode

# Code Injection Tests  
B102 = "exec_used"                      # Use of exec detected
B301 = "pickle_load"                    # Pickle library usage
B506 = "yaml_load"                      # Use of yaml.load

# Shell Injection Tests
B601 = "paramiko_calls"                 # Paramiko shell commands
B602 = "subprocess_popen_with_shell_equals_true"
B603 = "subprocess_without_shell_equals_true" 
B604 = "any_other_function_with_shell_equals_true"
B605 = "start_process_with_a_shell"
B606 = "start_process_with_no_shell"
B607 = "start_process_with_partial_path"

# SQL Injection Tests
B608 = "hardcoded_sql_expressions"      # SQL string formatting

# Cryptography Tests
B101 = "hashlib_insecure_functions"     # Weak hash functions
B326 = "weak_cryptographic_key"         # Weak crypto keys

# File System Tests
B103 = "set_bad_file_permissions"       # Insecure file permissions
B108 = "hardcoded_tmp_directory"        # Hardcoded temp paths

# Network Security Tests
B104 = "hardcoded_bind_all_interfaces"  # Binding to all interfaces
B501 = "request_with_no_cert_validation" # No certificate validation
B502 = "ssl_with_bad_version"           # Weak SSL/TLS versions

# Password/Secret Tests
B105 = "hardcoded_password_string"      # Hardcoded password strings
B106 = "hardcoded_password_funcarg"     # Password in function args
B107 = "hardcoded_password_default"     # Password in defaults

from bandit.core import test_properties as test
import bandit

@test.checks('Call')
@test.test_id('B999')
def detect_dangerous_function(context):
    """
    Detect calls to a specific dangerous function.
    
    Parameters:
    - context: Context object with call information
    
    Returns:
    Issue object if vulnerability found, None otherwise
    """
    if context.call_function_name == 'dangerous_function':
        return bandit.Issue(
            severity=bandit.HIGH,
            confidence=bandit.HIGH,
            text="Call to dangerous_function() detected",
            cwe=94  # Code Injection CWE
        )

@test.checks('Import', 'ImportFrom', 'Call')
@test.test_id('B998')
def detect_crypto_misuse(context):
    """Detect cryptographic misuse across imports and calls."""
    
    # Check imports
    weak_crypto_modules = ['md5', 'sha1', 'des']
    for module in weak_crypto_modules:
        if context.is_module_being_imported(module):
            return bandit.Issue(
                severity=bandit.MEDIUM,
                confidence=bandit.HIGH,
                text=f"Import of weak cryptographic module: {module}",
                cwe=326
            )
    
    # Check function calls
    if context.call_function_name_qual in ['hashlib.md5', 'hashlib.sha1']:
        return bandit.Issue(
            severity=bandit.MEDIUM,
            confidence=bandit.HIGH,
            text="Use of weak hash function",
            cwe=326
        )

@test.checks('Str')
@test.test_id('B997')
@test.takes_config('sensitive_strings')
def detect_sensitive_strings(context, config):
    """
    Detect sensitive string patterns from configuration.
    
    Configuration format in bandit.yaml:
    sensitive_strings:
      patterns:
        - "password"
        - "secret"
        - "api_key"
      min_length: 8
    """
    if not context.string_literal_value:
        return
    
    # Get configuration
    patterns = config.get('patterns', [])
    min_length = config.get('min_length', 6)
    
    string_value = context.string_literal_value.lower()
    
    if len(string_value) >= min_length:
        for pattern in patterns:
            if pattern.lower() in string_value:
                return bandit.Issue(
                    severity=bandit.MEDIUM,
                    confidence=bandit.LOW,
                    text=f"Potentially sensitive string containing '{pattern}'",
                    cwe=200  # Information Exposure
                )

@test.checks('Call') 
@test.test_id('B996')
def detect_unsafe_deserialization(context):
    """Detect unsafe deserialization patterns."""
    
    unsafe_functions = [
        'pickle.loads',
        'pickle.load', 
        'cPickle.loads',
        'cPickle.load',
        'dill.loads',
        'yaml.load'
    ]
    
    if context.call_function_name_qual in unsafe_functions:
        # Analyze arguments for user input
        severity = bandit.MEDIUM
        confidence = bandit.MEDIUM
        
        if context.call_args:
            first_arg = context.call_args[0]
            
            # Check for direct user input patterns
            if hasattr(first_arg, 'id') and first_arg.id in ['input', 'raw_input']:
                severity = bandit.HIGH
                confidence = bandit.HIGH
            
            # Check for network input patterns  
            elif (hasattr(first_arg, 'attr') and 
                  any(net_attr in first_arg.attr for net_attr in ['recv', 'read', 'readline'])):
                severity = bandit.HIGH
                confidence = bandit.MEDIUM
        
        return bandit.Issue(
            severity=severity,
            confidence=confidence,
            text=f"Unsafe deserialization with {context.call_function_name_qual}",
            cwe=502  # Deserialization of Untrusted Data
        )

@test.checks('Call')
@test.test_id('B995')
def detect_command_injection(context):
    """Detect potential command injection vulnerabilities."""
    
    try:
        if context.call_function_name_qual not in ['os.system', 'os.popen', 'subprocess.call']:
            return
        
        if not context.call_args:
            return
        
        first_arg = context.call_args[0]
        
        # String literal analysis
        if hasattr(first_arg, 's'):
            command = first_arg.s
            dangerous_chars = ['&', '|', ';', '`', '$', '(', ')']
            
            if any(char in command for char in dangerous_chars):
                return bandit.Issue(
                    severity=bandit.HIGH,
                    confidence=bandit.MEDIUM,
                    text="Command with shell metacharacters detected",
                    cwe=78
                )
        
        # Variable/expression analysis
        elif hasattr(first_arg, 'id'):
            # Variable is being used - lower confidence
            return bandit.Issue(
                severity=bandit.MEDIUM,
                confidence=bandit.LOW,
                text="Command execution with variable input",
                cwe=78
            )
    
    except AttributeError:
        # Handle cases where AST nodes don't have expected attributes
        return None

@test.accepts_baseline()
def custom_formatter(manager, fileobj, sev_level, conf_level, lines):
    """
    Custom output formatter for security reports.
    
    Parameters:
    - manager: BanditManager instance with scan results
    - fileobj: File-like object for output
    - sev_level: str, minimum severity level to include
    - conf_level: str, minimum confidence level to include  
    - lines: bool, include line numbers in output
    """
    
    # Get filtered issues
    issues = manager.get_issue_list(sev_level, conf_level)
    
    # Custom output format
    fileobj.write("=== Custom Security Report ===\n")
    fileobj.write(f"Total Issues: {len(issues)}\n\n")
    
    # Group by severity
    by_severity = {}
    for issue in issues:
        severity = issue.severity
        if severity not in by_severity:
            by_severity[severity] = []
        by_severity[severity].append(issue)
    
    # Output by severity level
    for severity in ['HIGH', 'MEDIUM', 'LOW']:
        if severity in by_severity:
            fileobj.write(f"\n{severity} SEVERITY ISSUES:\n")
            fileobj.write("-" * 40 + "\n")
            
            for issue in by_severity[severity]:
                fileobj.write(f"File: {issue.fname}\n")
                if lines:
                    fileobj.write(f"Line: {issue.lineno}\n")
                fileobj.write(f"Test: {issue.test_id}\n")
                fileobj.write(f"Issue: {issue.text}\n")
                fileobj.write(f"CWE: {issue.cwe}\n\n")

[entry_points]
bandit.plugins =
    my_custom_test = mypackage.security_tests:my_custom_test
    
bandit.formatters = 
    custom = mypackage.formatters:custom_formatter