#!/usr/bin/env python3
"""
Pre-commit hook to enforce safe_requests usage for SSRF protection.

This prevents direct usage of requests.get/post/Session which bypasses
SSRF validation. Use safe_get/safe_post/SafeSession from the security module.

See: https://cwe.mitre.org/data/definitions/918.html (CWE-918: SSRF)
"""

import ast
import sys


# Files/patterns where direct requests usage is allowed
ALLOWED_PATTERNS = {
    "security/safe_requests.py",  # The wrapper itself must use raw requests
    "tests/",  # Test files may need to mock/test raw requests
    "test_",  # Test files
    "_test.py",  # Test files
    "examples/",  # Example scripts run client-side, not server-side
}

# HTTP methods that should use safe_get/safe_post
UNSAFE_METHODS = {"get", "post", "put", "delete", "patch", "head", "options"}


class RequestsChecker(ast.NodeVisitor):
    """AST visitor to detect unsafe requests usage."""

    def __init__(self, filename: str):
        self.filename = filename
        self.errors = []

    def visit_Call(self, node):
        # Skip if file is in allowed list
        if self._is_file_allowed():
            return self.generic_visit(node)

        # Pattern 1: requests.get(), requests.post(), etc.
        if (
            isinstance(node.func, ast.Attribute)
            and node.func.attr in UNSAFE_METHODS
            and isinstance(node.func.value, ast.Name)
            and node.func.value.id == "requests"
        ):
            method = node.func.attr
            if method in ("get", "post"):
                safe_method = f"safe_{method}"
                suggestion = (
                    f"Use {safe_method}() from security module instead."
                )
            else:
                suggestion = (
                    "Use SafeSession() from security module for HTTP requests."
                )

            self.errors.append(
                (
                    node.lineno,
                    f"Direct requests.{method}() bypasses SSRF protection. "
                    f"{suggestion}",
                )
            )

        # Pattern 2: requests.Session()
        elif (
            isinstance(node.func, ast.Attribute)
            and node.func.attr == "Session"
            and isinstance(node.func.value, ast.Name)
            and node.func.value.id == "requests"
        ):
            self.errors.append(
                (
                    node.lineno,
                    "Direct requests.Session() bypasses SSRF protection. "
                    "Use SafeSession() from security module instead.",
                )
            )

        self.generic_visit(node)

    def _is_file_allowed(self) -> bool:
        """Check if this file is allowed to use requests directly."""
        for pattern in ALLOWED_PATTERNS:
            if pattern in self.filename:
                return True
        return False


def check_file(filename: str) -> bool:
    """Check a single Python file for unsafe requests usage."""
    if not filename.endswith(".py"):
        return True

    try:
        with open(filename, "r", encoding="utf-8") as f:
            content = f.read()
    except Exception as e:
        print(f"Error reading {filename}: {e}")
        return False

    try:
        tree = ast.parse(content, filename=filename)
        checker = RequestsChecker(filename)
        checker.visit(tree)

        if checker.errors:
            print(f"\n{filename}:")
            for line_num, error in checker.errors:
                print(f"  Line {line_num}: {error}")
            return False

    except SyntaxError:
        # Skip files with syntax errors (let other tools handle that)
        pass
    except Exception as e:
        print(f"Error parsing {filename}: {e}")
        return False

    return True


def main():
    """Main function to check all provided Python files."""
    if len(sys.argv) < 2:
        print("Usage: check-safe-requests.py <file1> <file2> ...")
        sys.exit(1)

    files_to_check = sys.argv[1:]
    has_errors = False

    for filename in files_to_check:
        if not check_file(filename):
            has_errors = True

    if has_errors:
        print("\n" + "=" * 70)
        print("SSRF Protection: Unsafe requests usage detected!")
        print("=" * 70)
        print("\nTo fix, replace direct requests calls with safe alternatives:")
        print("  - requests.get()     -> safe_get()")
        print("  - requests.post()    -> safe_post()")
        print("  - requests.Session() -> SafeSession()")
        print("\nImport from security module:")
        print("  from ...security import safe_get, safe_post, SafeSession")
        print("\nFor localhost/internal services, use:")
        print("  safe_get(url, allow_localhost=True)")
        print("  safe_get(url, allow_private_ips=True)")
        print("\nSee: https://cwe.mitre.org/data/definitions/918.html")
        sys.exit(1)
    else:
        sys.exit(0)


if __name__ == "__main__":
    main()