upif/cli.py

"""
upif.cli
~~~~~~~~

Command Line Interface for UPIF.
Allows utilizing the security layer directly from the terminal.

Usage:
    upif scan "your prompt here"
    upif interactive      (Start a testing session)
    upif setup-ai         (Download Neural Model)
    upif activate KEY     (Activate Pro License)
    upif check            (System Status)
"""

import argparse
import sys
import json
import os
import subprocess
from upif import guard

def setup_ai_model():
    """Downloads the AI model using optimum-cli."""
    print("UPIF: AI Model Setup 🧠")
    print("-----------------------")
    
    # Check if 'optimum' is installed
    try:
        import optimum
    except ImportError:
        print("❌ Missing AI dependencies.")
        print("   Please run: pip install optimum[onnxruntime]")
        return

    # Define paths relative to the INSTALLED package
    # This works even when installed via pip in site-packages
    base_dir = os.path.dirname(os.path.abspath(__file__))
    target_dir = os.path.join(base_dir, "data")
    model_name = "ProtectAI/deberta-v3-base-prompt-injection"

    print(f"Target: {target_dir}")
    print("Downloading Model (400MB+)...")
    
    try:
        # Run export command
        cmd = [
            "optimum-cli", "export", "onnx",
            "--model", model_name,
            target_dir,
            "--task", "text-classification"
        ]
        subprocess.check_call(cmd)
        
        # Renaissance logic
        original = os.path.join(target_dir, "model.onnx")
        final = os.path.join(target_dir, "guard_model.onnx")
        if os.path.exists(original):
            if os.path.exists(final): os.remove(final)
            os.rename(original, final)
            print("✅ AI Brain Installed Successfully.")
    except Exception as e:
        print(f"❌ Failed: {e}")

def run_interactive():
    """REPL for testing patterns."""
    print("UPIF Interactive Firewall (Ctrl+C to exit)")
    print("------------------------------------------")
    while True:
        try:
            prompt = input(">> ")
            if not prompt: continue
            
            # Neural Only check logic
            heuristic = guard.process_input(prompt)
            if heuristic != prompt:
                print(f"🚫 BLOCKED (Heuristic): {heuristic}")
            elif "[BLOCKED_BY_AI]" in heuristic:
                print(f"🤖 BLOCKED (Neural): {heuristic}")
            else:
                print(f"✅ PASSED")
                
        except KeyboardInterrupt:
            print("\nExiting.")
            break

def main():
    parser = argparse.ArgumentParser(description="UPIF: Universal Prompt Injection Firewall CLI")
    subparsers = parser.add_subparsers(dest="command", help="Available commands")

    # Commands
    subparsers.add_parser("scan", help="Scan a single string").add_argument("text", help="Text to scan")
    subparsers.add_parser("activate", help="Activate License").add_argument("key", help="License Key")
    subparsers.add_parser("check", help="System Status")
    subparsers.add_parser("setup-ai", help="Download AI Model (Requires upif[pro])")
    subparsers.add_parser("interactive", help="Start interactive testing session")
    subparsers.add_parser("version", help="Show version")

    args = parser.parse_args()

    if args.command == "version":
        from upif import __version__
        print(f"UPIF v{__version__} (Gold Master)")
        
    elif args.command == "scan":
        result = guard.process_input(args.text)
        print(f"Result: {result}")
        
    elif args.command == "interactive":
        run_interactive()
        
    elif args.command == "setup-ai":
        setup_ai_model()

    elif args.command == "activate":
        print(f"Activating: {args.key}...")
        if guard.license_manager.activate(args.key):
            print("✅ Success!")
        else:
            print("❌ Failed.")
            sys.exit(1)

    elif args.command == "check":
        tier = guard.license_manager.get_tier()
        ai_status = "Active" if getattr(guard.neural_guard, "enabled", False) else "Disabled"
        print(f"Status: {tier} | Neural: {ai_status}")
    
    else:
        parser.print_help()

if __name__ == "__main__":
    main()