Security Audit Script for Pyhton [Scan the "PYTHON" Code for Vulnerabilties]

security.py

import argparse
import json
import re
import sys
from dataclasses import asdict, dataclass
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional, Tuple


@dataclass
class SecurityIssue:
    file: str
    line: int
    category: str
    severity: str
    description: str
    pattern: str
    line_content: str
    recommendation: str
    confidence: str  # high, medium, low


class SecurityAuditor:

    def __init__(self, exclude_dirs: Optional[List[str]] = None,
                 exclude_files: Optional[List[str]] = None):
        self.exclude_dirs = set(exclude_dirs or ['.git', '__pycache__',
                                                  '.venv', 'venv', 'env',
                                                  'node_modules', '.idea',
                                                  '.vscode', 'build', 'dist'])
        self.exclude_files = set(exclude_files or [])
        self.issues: List[SecurityIssue] = []
        self.stats = {
            'files_scanned': 0,
            'lines_scanned': 0,
            'scan_start': None,
            'scan_end': None
        }

        # Define severity levels for each category
        self.severity_map = {
            'Hardcoded Secrets': 'CRITICAL',
            'SQL Injection': 'CRITICAL',
            'Command Injection': 'CRITICAL',
            'Weak Cryptography': 'HIGH',
            'Path Traversal': 'HIGH',
            'XSS Vulnerabilities': 'HIGH',
            'Insecure Deserialization': 'CRITICAL',
            'Debug Mode Enabled': 'MEDIUM',
            'Insecure Headers': 'MEDIUM',
            'Missing Validation': 'MEDIUM',
            'Deprecated/Insecure Functions': 'MEDIUM',
            'Environment Issues': 'LOW',
            'Code Quality Issues': 'LOW'
        }

        self.patterns = self._init_patterns()

    def _init_patterns(self) -> Dict[str, List[Tuple[str, str, str]]]:
        return {
            'Hardcoded Secrets': [
                
                (r'(?i)(api[_-]?key|secret[_-]?key|access[_-]?token|refresh[_-]?token|password|passwd|pwd)\s*[=:]\s*[\"\'][^\"\']{8,}[\"\']',
                 'Hardcoded credentials found in source code',
                 'Use environment variables or secure secret management systems'),

                (r'(?i)(aws[_-]?(secret|access)|azure[_-]?key|gcp[_-]?key)\s*[=:]\s*[\"\'][^\"\']+[\"\']',
                 'Cloud service credentials found',
                 'Use cloud provider SDKs with IAM roles or secure credential storage'),

                (r'(?i)(token|secret|key)\s*[=:]\s*(os\.environ\.get|os\.getenv)\([\"\'](?!SECRET_|TOKEN_|KEY_)',
                 'Potential hardcoded secret key name',
                 'Use generic names for environment variables to avoid leaking purpose'),
            ],

            'SQL Injection': [
                (r'(?i)cursor\.(execute|executemany)\([^)]*[\"\'].*%[^w].*[\"\']',
                 'Potential SQL injection with string formatting',
                 'Use parameterized queries with ? or %s placeholders'),

                (r'(?i)\.execute\(f[\"\'].*{.*}.*[\"\']',
                 'Potential SQL injection with f-strings',
                 'Use parameterized queries instead of string interpolation'),

                (r'(?i)(raw\(|raw_query|from_raw)',
                 'Use of raw SQL without parameterization',
                 'Use ORM methods or parameterized queries'),
            ],

            'Command Injection': [
                (r'(?i)(os\.system|os\.popen|subprocess\.(call|Popen|run))\([^)]*(f[\"\']|[\"\'].*\+.*[\"\'])',
                 'Potential command injection with string concatenation',
                 'Use shlex.quote() or pass arguments as list'),

                (r'eval\s*\(|exec\s*\(|compile\s*\(',
                 'Use of eval/exec/compile with user input',
                 'Avoid executing dynamic code; use safer alternatives'),

                (r'(?i)pickle\.loads?\(|marshal\.loads?\(|yaml\.load\(',
                 'Insecure deserialization',
                 'Use pickle.loads() only with trusted data or use yaml.safe_load()'),
            ],

            'XSS Vulnerabilities': [
                (r'(?i)(response\.write|print|echo)\s*\([^)]*(request\.|form\.|query\.|args\.)',
                 'Potential XSS vulnerability',
                 'Escape user input with html.escape() or use template engines'),

                (r'(?i)(<script>|javascript:|on\w+\s*=)',
                 'Potential inline JavaScript',
                 'Use Content Security Policy and proper escaping'),
            ],

            'Path Traversal': [
                (r'(?i)(open|file|openfile)\([^)]*\.\./',
                 'Potential path traversal',
                 'Use os.path.abspath() and validate paths'),

                (r'(?i)\.\./.*\.\./',
                 'Directory traversal pattern',
                 'Sanitize and validate file paths'),
            ],

            'Weak Cryptography': [
                (r'hashlib\.(md5|sha1)\(',
                 'Use of weak hash functions',
                 'Use hashlib.sha256() or better'),

                (r'random\.(random|randint|choice)\([^)]*\)',
                 'Use of insecure random for cryptographic purposes',
                 'Use secrets module or random.SystemRandom()'),

                (r'cryptography\.fernet\.InvalidToken',
                 'Missing proper exception handling for cryptography',
                 'Handle InvalidToken exception securely'),
            ],

            'Debug Mode Enabled': [
                (r'(?i)(debug|DEBUG)\s*=\s*(True|1)',
                 'Debug mode enabled in production code',
                 'Set debug=False in production'),

                (r'(?i)app\.run\([^)]*debug\s*=\s*True',
                 'Flask debug mode enabled',
                 'Use app.run(debug=False) in production'),
            ],

            'Insecure Headers': [
                (r'(?i)Access-Control-Allow-Origin\s*:\s*\*',
                 'Overly permissive CORS policy',
                 'Restrict CORS to specific domains'),

                (r'(?i)(X-Content-Type-Options|X-Frame-Options|Content-Security-Policy)\s*:',
                 'Missing or misconfigured security headers',
                 'Ensure proper security headers are set'),
            ],

            'Missing Validation': [
                (r'(?i)(input\(|raw_input\()',
                 'Direct user input without validation',
                 'Validate and sanitize all user input'),

                (r'(?i)\.get\([^)]*\)\s*(?:\.|$)',
                 'Potential missing validation of dictionary access',
                 'Add proper validation and error handling'),
            ],

            'Insecure Deserialization': [
                (r'(?i)(pickle|marshal|yaml)\.load\(',
                 'Insecure deserialization function',
                 'Use safe deserialization methods or validate input'),

                (r'__reduce__|__setstate__',
                 'Potential pickle exploitation',
                 'Avoid unpickling untrusted data'),
            ],

            'Environment Issues': [
                (r'(?i)(localhost|127\.0\.0\.1):(8000|8080|3000)',
                 'Hardcoded localhost URLs',
                 'Use configuration files or environment variables'),

                (r'(?i)\.pyc$|\.pyo$',
                 'Python bytecode files in repository',
                 'Add *.pyc and *.pyo to .gitignore'),
            ],

            'Deprecated/Insecure Functions': [
                (r'(?i)md5\(|sha1\(|crypt\(',
                 'Deprecated/insecure cryptographic function',
                 'Use modern, secure alternatives'),

                (r'execfile\(|reload\(|xreadlines\(\)',
                 'Deprecated Python function',
                 'Use modern Python 3 equivalents'),
            ],
        }

    def should_exclude(self, file_path: Path) -> bool:

        for part in file_path.parts:
            if part in self.exclude_dirs:
                return True

        for pattern in self.exclude_files:
            if re.match(pattern, str(file_path)):
                return True

        return False

    def assess_confidence(self, pattern: str, line_content: str) -> str:
        line_lower = line_content.lower()

        high_confidence = [
            r'password\s*=',
            r'api[_-]?key\s*=',
            r'eval\(',
            r'os\.system\(',
            r'hashlib\.md5\(',
        ]

        medium_confidence = [
            r'debug\s*=\s*True',
            r'\.execute\(f',
            r'random\.random\(',
        ]

        for hc in high_confidence:
            if re.search(hc, line_lower):
                return 'high'

        for mc in medium_confidence:
            if re.search(mc, line_lower):
                return 'medium'

        return 'low'

    def scan_file(self, file_path: Path) -> None:
        try:
            with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                content = f.read()
                lines = content.split('\n')
                self.stats['lines_scanned'] += len(lines)

            for category, pattern_list in self.patterns.items():
                for pattern, description, recommendation in pattern_list:
                    for i, line in enumerate(lines, 1):
                        line_stripped = line.strip()
                        if not line_stripped or line_stripped.startswith('#'):
                            continue

                        if re.search(pattern, line, re.IGNORECASE):
                            confidence = self.assess_confidence(pattern, line)

                            issue = SecurityIssue(
                                file=str(file_path),
                                line=i,
                                category=category,
                                severity=self.severity_map.get(category, 'MEDIUM'),
                                description=description,
                                pattern=pattern,
                                line_content=line_stripped[:200],  # Truncate long lines
                                recommendation=recommendation,
                                confidence=confidence
                            )
                            self.issues.append(issue)

        except (IOError, PermissionError, UnicodeDecodeError) as e:
            print(f"⚠️  Warning: Could not read {file_path}: {e}")

    def scan_directory(self, directory: Path) -> None:
        self.stats['scan_start'] = datetime.now()

        try:
            for file_path in directory.rglob('*.py'):
                if not self.should_exclude(file_path):
                    self.stats['files_scanned'] += 1
                    self.scan_file(file_path)

                    # Progress indicator
                    if self.stats['files_scanned'] % 50 == 0:
                        print(f"📁 Scanned {self.stats['files_scanned']} files...")

        except KeyboardInterrupt:
            print("\n\nScan interrupted by user.")
            sys.exit(1)

        self.stats['scan_end'] = datetime.now()

    def generate_report(self, output_format: str = 'text', output_file: Optional[str] = None) -> None:

        by_severity = {}
        for issue in self.issues:
            if issue.severity not in by_severity:
                by_severity[issue.severity] = []
            by_severity[issue.severity].append(issue)

        duration = (self.stats['scan_end'] - self.stats['scan_start']).total_seconds() if self.stats['scan_end'] else 0

        if output_format == 'json':
            report = {
                'metadata': {
                    'scan_date': datetime.now().isoformat(),
                    'duration_seconds': duration,
                    'files_scanned': self.stats['files_scanned'],
                    'lines_scanned': self.stats['lines_scanned'],
                    'total_issues': len(self.issues)
                },
                'issues': [asdict(issue) for issue in self.issues],
                'summary': {
                    severity: len(issues) for severity, issues in by_severity.items()
                }
            }

            report_json = json.dumps(report, indent=2, default=str)

            if output_file:
                with open(output_file, 'w') as f:
                    f.write(report_json)
            else:
                print(report_json)

        elif output_format == 'csv':
            import csv
            if output_file:
                with open(output_file, 'w', newline='') as f:
                    writer = csv.writer(f)
                    writer.writerow(['File', 'Line', 'Category', 'Severity', 'Description',
                                    'Confidence', 'Recommendation', 'Code Snippet'])
                    for issue in self.issues:
                        writer.writerow([
                            issue.file, issue.line, issue.category, issue.severity,
                            issue.description, issue.confidence, issue.recommendation,
                            issue.line_content
                        ])
            else:
                # Print CSV to stdout
                writer = csv.writer(sys.stdout)
                writer.writerow(['File', 'Line', 'Category', 'Severity', 'Description',
                                'Confidence', 'Recommendation', 'Code Snippet'])
                for issue in self.issues:
                    writer.writerow([
                        issue.file, issue.line, issue.category, issue.severity,
                        issue.description, issue.confidence, issue.recommendation,
                        issue.line_content
                    ])

        else:
            output = []
            output.append("=" * 80)
            output.append("🔍 SECURITY AUDIT REPORT")
            output.append("=" * 80)
            output.append(f"📅 Scan Date: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
            output.append(f"⏱️  Duration: {duration:.2f} seconds")
            output.append(f"📄 Files Scanned: {self.stats['files_scanned']}")
            output.append(f"📝 Lines Scanned: {self.stats['lines_scanned']}")
            output.append(f"⚠️  Total Issues Found: {len(self.issues)}")
            output.append("")

            # Summary by severity
            output.append("📊 SEVERITY SUMMARY:")
            output.append("-" * 40)
            for severity in ['CRITICAL', 'HIGH', 'MEDIUM', 'LOW']:
                if severity in by_severity:
                    output.append(f"  {severity}: {len(by_severity[severity])} issues")

            output.append("")

            for severity in ['CRITICAL', 'HIGH', 'MEDIUM', 'LOW']:
                if severity in by_severity and by_severity[severity]:
                    output.append(f"🚨 {severity} ISSUES ({len(by_severity[severity])}):")
                    output.append("-" * 60)

                    by_category = {}
                    for issue in by_severity[severity]:
                        if issue.category not in by_category:
                            by_category[issue.category] = []
                        by_category[issue.category].append(issue)

                    for category, issues in by_category.items():
                        output.append(f"\n  📂 {category}:")
                        for issue in issues:
                            output.append(f"    📍 {issue.file}:{issue.line}")
                            output.append(f"       Confidence: {issue.confidence.upper()}")
                            output.append(f"       Description: {issue.description}")
                            output.append(f"       Code: {issue.line_content}")
                            output.append(f"       Recommendation: {issue.recommendation}")
                            output.append("")
                    output.append("")

            if not self.issues:
                output.append("✅ No security issues found!")
                output.append("")
                output.append("💡 Tips:")
                output.append("  • Consider running dependency audit with: pip-audit")
                output.append("  • Use bandit for additional static analysis: bandit -r .")
                output.append("  • Check for secrets with: detect-secrets scan")

            output.append("=" * 80)
            output.append("📋 Security audit complete!")

            report_text = "\n".join(output)

            if output_file:
                with open(output_file, 'w') as f:
                    f.write(report_text)
            else:
                print(report_text)


def main():
    parser = argparse.ArgumentParser(
        description='Security Audit Tool for Python Code',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  %(prog)s                          # Scan current directory
  %(prog)s /path/to/project         # Scan specific directory
  %(prog)s --format json            # Output as JSON
  %(prog)s --output report.json     # Save report to file
  %(prog)s --exclude-dir venv,test  # Exclude directories
        """
    )

    parser.add_argument('path', nargs='?', default='.',
                       help='Directory to scan (default: current directory)')
    parser.add_argument('--format', choices=['text', 'json', 'csv'], default='text',
                       help='Output format (default: text)')
    parser.add_argument('--output', '-o', help='Output file (default: stdout)')
    parser.add_argument('--exclude-dir', help='Comma-separated directories to exclude')
    parser.add_argument('--exclude-file', help='Comma-separated file patterns to exclude')
    parser.add_argument('--verbose', '-v', action='store_true',
                       help='Verbose output')

    args = parser.parse_args()

    # Parse exclude directories
    exclude_dirs = []
    if args.exclude_dir:
        exclude_dirs = [d.strip() for d in args.exclude_dir.split(',')]

    exclude_files = []
    if args.exclude_file:
        exclude_files = [f.strip() for f in args.exclude_file.split(',')]

    # Create auditor and scan
    auditor = SecurityAuditor(exclude_dirs=exclude_dirs, exclude_files=exclude_files)

    scan_path = Path(args.path)
    if not scan_path.exists():
        print(f"❌ Error: Path '{args.path}' does not exist")
        sys.exit(1)

    if args.verbose:
        print(f"🔍 Starting security audit of: {scan_path.absolute()}")
        print(f"📁 Excluding directories: {exclude_dirs}")

    auditor.scan_directory(scan_path)

    auditor.generate_report(
        output_format=args.format,
        output_file=args.output
    )

    # Exit with appropriate code
    critical_issues = sum(1 for issue in auditor.issues if issue.severity == 'CRITICAL')
    if critical_issues > 0:
        if args.verbose:
            print(f"\n❌ Found {critical_issues} critical issues. Audit failed.")
        sys.exit(1)
    else:
        sys.exit(0)


if __name__ == "__main__":
    main()


"""
# Basic scan
python security_audit.py

# Scan specific directory
python security_audit.py /path/to/project

# Generate JSON report
python security_audit.py --format json --output report.json

# Exclude directories
python security_audit.py --exclude-dir venv,test,node_modules

# Verbose output
python security_audit.py -v

# Use in CI/CD pipeline (exits with code 1 if critical issues found)
python security_audit.py --format json && echo "Security audit passed"

"""