Security-aware Always Activate The Venv for Python

.zshrc

# Standard zshrc stuff like path, etc.
# ...

# Enable using the two files above:

# Functions
source venv-auto-activate.sh

# Venv security whitelist/blocklist
alias venv-security='uv run -q --no-project ~/scripts/venv-security.py'
alias vnvsec='uv run -q --no-project ~/scripts/venv-security.py'

venv-auto-activate.sh

# Virtual Environment Auto-Activation
# ===================================
# Automatically activates/deactivates Python virtual environments
# when changing directories
# See full write up at https://mkennedy.codes/posts/always-activate-the-venv-a-shell-script/

# Auto-activate virtual environment for any project with a venv directory
function chpwd() {
    # Function to find venv directory in current path or parent directories
    # Prefers 'venv' over '.venv' if both exist
    local find_venv() {
        local dir="$PWD"
        while [[ "$dir" != "/" ]]; do
            if [[ -d "$dir/venv" && -f "$dir/venv/bin/activate" ]]; then
                echo "$dir/venv"
                return 0
            elif [[ -d "$dir/.venv" && -f "$dir/.venv/bin/activate" ]]; then
                echo "$dir/.venv"
                return 0
            fi
            dir="$(dirname "$dir")"
        done
        return 1
    }
    
    local venv_path
    venv_path=$(find_venv)
    
    if [[ -n "$venv_path" ]]; then
        # Normalize paths for comparison (handles symlinks and path differences)
        # Use zsh :A modifier to resolve paths without triggering chpwd recursively
        local normalized_venv_path="${venv_path:A}"
        local normalized_current_venv=""
        if [[ -n "${VIRTUAL_ENV:-}" ]]; then
            normalized_current_venv="${VIRTUAL_ENV:A}"
        fi
        
        # We found a venv, check if it's already active
        if [[ "$normalized_current_venv" != "$normalized_venv_path" ]]; then
            # Deactivate current venv if different
            if [[ -n "${VIRTUAL_ENV:-}" ]] && type deactivate >/dev/null 2>&1; then
                deactivate
            fi
            # Security check: only activate trusted venvs
            if uv run -q --no-project ~/scripts/venv-security.py check "$normalized_venv_path"; then
                source "$venv_path/bin/activate"
                local project_name=$(basename "$(dirname "$venv_path")")
                echo "🐍 Activated virtual environment \033[95m$project_name\033[0m."
            fi
        fi
    else
        # No venv found, deactivate if we have one active
        if [[ -n "${VIRTUAL_ENV:-}" ]] && type deactivate >/dev/null 2>&1; then
            local project_name=$(basename "$(dirname "${VIRTUAL_ENV}")")
            deactivate
            echo "🔒 Deactivated virtual environment \033[95m$project_name\033[0m."
        elif [[ -n "${VIRTUAL_ENV:-}" ]]; then
            # VIRTUAL_ENV is set but deactivate function is not available
            # This can happen when opening a new shell with VIRTUAL_ENV from previous session
            unset VIRTUAL_ENV
        fi
    fi
}

# Run the chpwd function when the shell starts
chpwd

venv-security.py

#!/usr/bin/env python3
"""
Venv security whitelist/blocklist manager.

Usage:
    venv-security check <venv-path>   # Check and prompt if unknown
    venv-security trust <venv-path>   # Add to trusted list
    venv-security block <venv-path>   # Add to blocked list
    venv-security list                # Show current config
    venv-security remove <venv-path>  # Remove from both lists

Exit codes:
    0 = trusted (activate)
    1 = blocked or declined (don't activate)
    2 = error
"""

import json
import sys
from pathlib import Path

CONFIG_DIR = Path.home() / ".config" / "venv-security"
CONFIG_FILE = CONFIG_DIR / "config.json"

def load_config() -> dict:
    if not CONFIG_FILE.exists():
        return {"trusted": [], "blocked": []}
    try:
        return json.loads(CONFIG_FILE.read_text())
    except (json.JSONDecodeError, IOError):
        return {"trusted": [], "blocked": []}

def save_config(config: dict) -> None:
    CONFIG_DIR.mkdir(parents=True, exist_ok=True)
    CONFIG_FILE.write_text(json.dumps(config, indent=2) + "\n")

def resolve_path(path: str) -> str:
    return str(Path(path).resolve())

def check_venv(venv_path: str) -> int:
    """Check venv status, prompt if unknown. Returns exit code."""
    path = resolve_path(venv_path)
    config = load_config()

    if path in config["trusted"]:
        return 0

    if path in config["blocked"]:
        return 1

    # Unknown venv - prompt user
    print(f"\n🔒 Unknown virtual environment detected:")
    print(f"   {path}\n")

    try:
        response = input("Trust and activate this venv? [y/N/block] ").strip().lower()
    except (EOFError, KeyboardInterrupt):
        print()
        return 1

    if response in ("y", "yes"):
        config["trusted"].append(path)
        save_config(config)
        print(f"✓ Added to trusted list")
        return 0
    elif response in ("b", "block"):
        config["blocked"].append(path)
        save_config(config)
        print(f"✗ Added to blocked list (won't ask again)")
        return 1
    else:
        print(f"– Skipped (will ask again next time)")
        return 1

def trust_venv(venv_path: str) -> int:
    """Add venv to trusted list."""
    path = resolve_path(venv_path)
    config = load_config()

    # Remove from blocked if present
    if path in config["blocked"]:
        config["blocked"].remove(path)

    if path not in config["trusted"]:
        config["trusted"].append(path)
        save_config(config)
        print(f"✓ Trusted: {path}")
    else:
        print(f"Already trusted: {path}")
    return 0

def block_venv(venv_path: str) -> int:
    """Add venv to blocked list."""
    path = resolve_path(venv_path)
    config = load_config()

    # Remove from trusted if present
    if path in config["trusted"]:
        config["trusted"].remove(path)

    if path not in config["blocked"]:
        config["blocked"].append(path)
        save_config(config)
        print(f"✗ Blocked: {path}")
    else:
        print(f"Already blocked: {path}")
    return 0

def remove_venv(venv_path: str) -> int:
    """Remove venv from both lists."""
    path = resolve_path(venv_path)
    config = load_config()
    removed = False

    if path in config["trusted"]:
        config["trusted"].remove(path)
        removed = True
    if path in config["blocked"]:
        config["blocked"].remove(path)
        removed = True

    if removed:
        save_config(config)
        print(f"Removed: {path}")
    else:
        print(f"Not found in any list: {path}")
    return 0

def list_config() -> int:
    """Display current configuration."""
    config = load_config()

    print("Trusted venvs:")
    if config["trusted"]:
        for p in sorted(config["trusted"]):
            print(f"  ✓ {p}")
    else:
        print("  (none)")

    print("\nBlocked venvs:")
    if config["blocked"]:
        for p in sorted(config["blocked"]):
            print(f"  ✗ {p}")
    else:
        print("  (none)")

    return 0

def main() -> int:
    if len(sys.argv) < 2:
        print(__doc__)
        return 2

    command = sys.argv[1]

    if command == "list":
        return list_config()

    if len(sys.argv) < 3 and command != "list":
        print(f"Usage: venv-security {command} <venv-path>")
        return 2

    venv_path = sys.argv[2] if len(sys.argv) > 2 else ""

    commands = {
        "check": check_venv,
        "trust": trust_venv,
        "block": block_venv,
        "remove": remove_venv,
    }

    if command not in commands:
        print(f"Unknown command: {command}")
        print(__doc__)
        return 2

    return commands[command](venv_path)

if __name__ == "__main__":
    sys.exit(main())