#!/usr/bin/env python3
"""python_ast_security_audit.py

A safe Python 3 utility for static inspection of Python source files using the
Abstract Syntax Tree (AST). It helps security reviewers and engineers spot
common patterns that deserve manual review, such as:

- eval / exec / compile usage
- subprocess calls and shell usage
- dynamic imports
- broad exception handling
- hardcoded secrets or credentials in literals
- potentially unsafe deserialization sinks

Important:
- This script does NOT execute the target code.
- This script is not a full security scanner.
- Use it as one layer in a broader review process.

Examples:
    python python_ast_security_audit.py path/to/file.py
    python python_ast_security_audit.py path/to/repo --recursive
    python python_ast_security_audit.py path/to/file.py --json
    python python_ast_security_audit.py path/to/file.py --fail-on high
"""

from __future__ import annotations

import argparse
import ast
import json
import os
import sys
from dataclasses import dataclass, asdict
from pathlib import Path
from typing import Iterable, List, Optional


HIGH = "high"
MEDIUM = "medium"
LOW = "low"
SEVERITIES = {HIGH, MEDIUM, LOW}


@dataclass
class Finding:
    file: str
    line: int
    severity: str
    rule: str
    message: str
    code: Optional[str] = None


class AuditVisitor(ast.NodeVisitor):
    def __init__(self, filename: str) -> None:
        self.filename = filename
        self.findings: List[Finding] = []
        self.imports: set[str] = set()

    def add_finding(self, node: ast.AST, severity: str, rule: str, message: str) -> None:
        line = getattr(node, "lineno", 1)
        code = None
        self.findings.append(
            Finding(
                file=self.filename,
                line=line,
                severity=severity,
                rule=rule,
                message=message,
                code=code,
            )
        )

    def visit_Import(self, node: ast.Import) -> None:
        for alias in node.names:
            self.imports.add(alias.name.split(".")[0])
            if alias.name in {"subprocess", "pickle", "marshal", "socket", "os", "shlex"}:
                self.add_finding(
                    node,
                    LOW,
                    "risky-import",
                    f"Imported module '{alias.name}' may deserve review depending on usage.",
                )
        self.generic_visit(node)

    def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
        if node.module:
            self.imports.add(node.module.split(".")[0])
            if node.module in {"subprocess", "pickle", "marshal", "socket", "os", "shlex"}:
                self.add_finding(
                    node,
                    LOW,
                    "risky-import",
                    f"Imported from module '{node.module}' may deserve review depending on usage.",
                )
        self.generic_visit(node)

    def visit_Call(self, node: ast.Call) -> None:
        name = self._resolve_call_name(node)

        if name in {"eval", "exec", "compile"}:
            self.add_finding(node, HIGH, "dynamic-execution", f"Use of {name}() should be manually reviewed.")

        if name in {"input"}:
            self.add_finding(node, LOW, "interactive-input", "Interactive input found; verify it is not used in unsafe flows.")

        if name in {"subprocess.run", "subprocess.Popen", "subprocess.call", "subprocess.check_output"}:
            shell_true = any(self._is_shell_true(kw) for kw in node.keywords)
            severity = HIGH if shell_true else MEDIUM
            msg = f"Subprocess call detected ({name})."
            if shell_true:
                msg += " shell=True increases risk and should be reviewed carefully."
            self.add_finding(node, severity, "subprocess-use", msg)

        if name in {"pickle.loads", "pickle.load", "yaml.load", "marshal.loads", "marshal.load"}:
            self.add_finding(node, HIGH, "deserialization-sink", f"Potentially unsafe deserialization sink detected: {name}.")

        if name in {"__import__", "importlib.import_module"}:
            self.add_finding(node, MEDIUM, "dynamic-import", f"Dynamic import detected: {name}. Review source of module name.")

        if name in {"open"}:
            self._check_open_call(node)

        self.generic_visit(node)

    def visit_ExceptHandler(self, node: ast.ExceptHandler) -> None:
        if node.type is None:
            self.add_finding(node, MEDIUM, "broad-exception", "Bare except detected; it may hide security-relevant failures.")
        self.generic_visit(node)

    def visit_Constant(self, node: ast.Constant) -> None:
        if isinstance(node.value, str) and self._looks_like_secret(node.value):
            self.add_finding(node, MEDIUM, "possible-secret", "String literal looks like a secret, token, or credential.")
        self.generic_visit(node)

    def _resolve_call_name(self, node: ast.Call) -> str:
        func = node.func
        if isinstance(func, ast.Name):
            return func.id
        if isinstance(func, ast.Attribute):
            parts = []
            current: ast.AST = func
            while isinstance(current, ast.Attribute):
                parts.append(current.attr)
                current = current.value
            if isinstance(current, ast.Name):
                parts.append(current.id)
            return ".".join(reversed(parts))
        return ""

    def _is_shell_true(self, kw: ast.keyword) -> bool:
        return kw.arg == "shell" and isinstance(kw.value, ast.Constant) and kw.value.value is True

    def _check_open_call(self, node: ast.Call) -> None:
        path_arg = node.args[0] if node.args else None
        if isinstance(path_arg, ast.BinOp):
            self.add_finding(node, MEDIUM, "path-construction", "File path is built dynamically; review for path traversal or unsafe concatenation.")
        elif isinstance(path_arg, ast.JoinedStr):
            self.add_finding(node, MEDIUM, "path-construction", "f-string used in file path; review for untrusted input.")

    def _looks_like_secret(self, value: str) -> bool:
        lowered = value.lower()
        markers = ["api_key", "apikey", "secret", "token", "passwd", "password", "bearer "]
        return any(marker in lowered for marker in markers) and len(value) > 12


def iter_python_files(paths: Iterable[Path], recursive: bool) -> Iterable[Path]:
    for path in paths:
        if path.is_file() and path.suffix == ".py":
            yield path
        elif path.is_dir() and recursive:
            yield from sorted(p for p in path.rglob("*.py") if p.is_file())


def analyze_file(path: Path) -> List[Finding]:
    try:
        source = path.read_text(encoding="utf-8")
    except OSError as exc:
        return [Finding(file=str(path), line=1, severity=LOW, rule="read-error", message=f"Could not read file: {exc}")]

    try:
        tree = ast.parse(source, filename=str(path))
    except SyntaxError as exc:
        return [Finding(file=str(path), line=exc.lineno or 1, severity=HIGH, rule="syntax-error", message=f"Syntax error: {exc.msg}")]

    visitor = AuditVisitor(str(path))
    visitor.visit(tree)
    return visitor.findings


def print_text(findings: List[Finding]) -> None:
    if not findings:
        print("No findings.")
        return
    for finding in findings:
        print(f"{finding.file}:{finding.line} [{finding.severity}] {finding.rule} - {finding.message}")


def print_json(findings: List[Finding]) -> None:
    print(json.dumps([asdict(f) for f in findings], indent=2, sort_keys=False))


def parse_args(argv: List[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Static AST-based Python security audit helper.")
    parser.add_argument("paths", nargs="+", help="One or more Python files or directories to scan.")
    parser.add_argument("--recursive", action="store_true", help="Recursively scan directories for .py files.")
    parser.add_argument("--json", action="store_true", help="Output findings as JSON.")
    parser.add_argument(
        "--fail-on",
        choices=sorted(SEVERITIES),
        default=None,
        help="Exit with code 2 if findings at or above the selected severity are present.",
    )
    return parser.parse_args(argv)


def severity_rank(severity: str) -> int:
    return {HIGH: 3, MEDIUM: 2, LOW: 1}.get(severity, 0)


def should_fail(findings: List[Finding], threshold: Optional[str]) -> bool:
    if threshold is None:
        return False
    threshold_rank = severity_rank(threshold)
    return any(severity_rank(f.severity) >= threshold_rank for f in findings)


def main(argv: List[str]) -> int:
    args = parse_args(argv)
    input_paths = [Path(p) for p in args.paths]

    targets = list(iter_python_files(input_paths, args.recursive))
    if not targets:
        print("No Python files found.", file=sys.stderr)
        return 1

    all_findings: List[Finding] = []
    for target in targets:
        all_findings.extend(analyze_file(target))

    if args.json:
        print_json(all_findings)
    else:
        print_text(all_findings)

    if should_fail(all_findings, args.fail_on):
        return 2
    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))