Quick Security Assessment

#!/usr/bin/env python3
"""
Quick Security Assessment Example
Complete ModelRed workflow in 5 minutes
"""

import asyncio
import os
from datetime import datetime
from dotenv import load_dotenv
from modelred import ModelRed, ModelRedError

# Load environment variables
load_dotenv()

async def quick_security_assessment():
    """
    Complete security assessment workflow:
    1. Authenticate with ModelRed
    2. Register an OpenAI model
    3. Run security assessment
    4. Display results
    """

    print("🚀 ModelRed Quick Security Assessment")
    print("=" * 50)

    try:
        # Step 1: Initialize ModelRed client
        print("\n📡 Connecting to ModelRed...")
        async with ModelRed() as client:

            # Verify authentication
            account = await client.validate_api_key()
            print(f"✅ Connected to ModelRed!")
            print(f"   Organization: {account.get('organization', 'Personal')}")

            # Check usage limits
            usage = await client.get_usage_stats()
            print(f"   Plan: {usage.tier}")
            print(f"   Assessments: {usage.assessments_this_month}/{usage.assessments_limit}")

            # Step 2: Register OpenAI model
            print("\n🤖 Registering OpenAI model...")
            model_id = f"quick-test-{datetime.now().strftime('%Y%m%d_%H%M%S')}"

            success = await client.register_openai_model(
                model_id=model_id,
                api_key=os.getenv("OPENAI_API_KEY"),
                model_name="gpt-3.5-turbo",
                metadata={
                    "purpose": "Quick security assessment example",
                    "created_by": "ModelRed SDK Tutorial"
                }
            )

            if success:
                print(f"✅ Model registered: {model_id}")
            else:
                print("❌ Model registration failed")
                return

            # Step 3: Run security assessment
            print("\n🧪 Starting security assessment...")
            print("   This will take 2-3 minutes...")

            # Start assessment and wait for completion
            result = await client.run_assessment(
                model_id=model_id,
                test_suites=["basic_security"],
                wait_for_completion=True,
                timeout_minutes=5
            )

            # Step 4: Display results
            print("\n📊 Assessment Results:")
            print("=" * 30)
            print(f"Overall Score: {result.overall_score}/10")
            print(f"Risk Level: {result.risk_level.value}")
            print(f"Tests Run: {result.total_tests}")
            print(f"Tests Passed: {result.passed_tests}")
            print(f"Tests Failed: {result.failed_tests}")

            # Show category breakdown
            if result.categories:
                print("\n📋 Category Breakdown:")
                for category, details in result.categories.items():
                    score = details.get('score', 0)
                    print(f"   {category}: {score}/10")

            # Show top recommendations
            if result.recommendations:
                print("\n💡 Top Recommendations:")
                for i, rec in enumerate(result.recommendations[:3], 1):
                    print(f"   {i}. {rec}")

            # Show report URL
            if result.report_url:
                print(f"\n🔗 Detailed Report: {result.report_url}")

            # Step 5: Cleanup (optional)
            print(f"\n🧹 Cleaning up test model...")
            await client.delete_model(model_id)
            print("✅ Test model removed")

            print("\n🎉 Assessment completed successfully!")
            print("💡 Next: Try running assessments on your production models")

    except ModelRedError as e:
        print(f"\n❌ ModelRed Error: {e.message}")
        print("💡 Check your API keys and try again")
    except Exception as e:
        print(f"\n❌ Unexpected error: {e}")
        print("💡 Please check the troubleshooting guide")

# Environment setup helper
def setup_environment():
    """Check and guide user through environment setup"""
    print("🔧 Environment Setup Check")
    print("-" * 30)

    # Check ModelRed API key
    mr_key = os.getenv("MODELRED_API_KEY")
    if mr_key and mr_key.startswith("mr_"):
        print("✅ ModelRed API key found")
    else:
        print("❌ ModelRed API key missing or invalid")
        print("   Set: export MODELRED_API_KEY='mr_your_key_here'")
        return False

    # Check OpenAI API key
    openai_key = os.getenv("OPENAI_API_KEY")
    if openai_key and openai_key.startswith("sk-"):
        print("✅ OpenAI API key found")
    else:
        print("❌ OpenAI API key missing or invalid")
        print("   Set: export OPENAI_API_KEY='sk-your_key_here'")
        return False

    print("✅ Environment setup complete!")
    return True

if __name__ == "__main__":
    print("ModelRed Quick Security Assessment")
    print("=" * 50)

    # Check environment first
    if setup_environment():
        print("\n")
        # Run the assessment
        asyncio.run(quick_security_assessment())
    else:
        print("\n💡 Please set up your environment variables and try again")