#!/bin/bash

# Default values
PACKAGE_NAME="litellm"
#PACKAGE_NAME="pip"  # use pip to test; every venv has pip
SEARCH_DIRS=("/home" "/opt" "/var/www")

# Parse command line flags
while getopts "p:d:h" opt; do
  case $opt in
    p) PACKAGE_NAME="$OPTARG" ;;
    d) SEARCH_DIRS+=("$OPTARG") ;;
    h) 
      echo "Usage: $0 [-p package_name] [-d additional_search_dir]"
      echo "Example: $0 -p pip -d /usr/local"
      exit 0 
      ;;
    \?) echo "Invalid option: -$OPTARG" >&2; exit 1 ;;
  esac
done

echo "Scanning for $PACKAGE_NAME installations..."
echo "Directories targeted for venv scan: ${SEARCH_DIRS[*]}"
echo "-----------------------------------"

# 1. Check system-wide Python environments
echo "Checking system-wide Python environments..."
for pip_cmd in pip pip3; do
    if command -v "$pip_cmd" &> /dev/null; then
        version=$("$pip_cmd" show "$PACKAGE_NAME" 2>/dev/null | grep "^Version:" | awk '{print $2}')
        if [ -n "$version" ]; then
            echo "[!] Found system-wide $PACKAGE_NAME version $version (via $pip_cmd)"
        fi
    fi
done

# 2. Search targeted directories for virtual environments
# Python package directories in site-packages often replace hyphens with underscores
PACKAGE_DIR_PREFIX=$(echo "$PACKAGE_NAME" | tr '-' '_')

for search_dir in "${SEARCH_DIRS[@]}"; do
    if [ ! -d "$search_dir" ]; then
        echo "Skipping $search_dir (Directory does not exist or is inaccessible)"
        continue
    fi
    
    echo "Scanning $search_dir for isolated virtual environments..."
    find "$search_dir" -type d \( -iname "${PACKAGE_NAME}-*.dist-info" -o -iname "${PACKAGE_DIR_PREFIX}-*.dist-info" \) 2>/dev/null | while read -r dist_info; do
        metadata_file="$dist_info/METADATA"
        if [ -f "$metadata_file" ]; then
            # Extract the version number directly from the METADATA file
            version=$(grep "^Version:" "$metadata_file" | awk '{print $2}')
            
            # Extract the base directory of the venv for cleaner output
            venv_dir=$(echo "$dist_info" | awk -F'/lib/python' '{print $1}')
            
            # Fallback if the path structure is non-standard
            if [ -z "$venv_dir" ] || [ "$venv_dir" = "$dist_info" ]; then
                venv_dir=$(dirname "$dist_info")
            fi
            
            echo "[!] Found $PACKAGE_NAME version $version in: $venv_dir"
        fi
    done
done

echo "-----------------------------------"
echo "Scan complete."
if [ "$PACKAGE_NAME" = "litellm" ]; then
    echo "Note: The known compromised versions of litellm are 1.82.7 and 1.82.8."
fi