app.py

import gradio as gr
from transformers import pipeline

# Initialize pipelines
ocr = pipeline("image-to-text", model="microsoft/trocr-base-printed")
classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")

# Common phishing indicators
SUSPICIOUS_PHRASES = [
    "urgent", "immediately", "password", "account locked", "wire transfer",
    "bank verification", "click here", "verification code", "credit card",
    "suspended", "login now", "reset your password", "act now", "unusual activity",
    "security alert", "confirm your identity", "gift card", "lottery winner",
    "inheritance", "tax refund", "payment pending"
]

def extract_text_from_image(image):
    """Extract text from uploaded image using OCR."""
    if image is None:
        return ""
    try:
        result = ocr(image)
        return result[0]["generated_text"] if result else ""
    except Exception as e:
        return f"Error processing image: {str(e)}"

def analyze_text(text):
    """Analyze text for phishing indicators."""
    if not text.strip():
        return "No text to analyze.", "", gr.update(visible=False)
    
    # Zero-shot Classification
    candidate_labels = ["Phishing Email", "Legitimate Email"]
    result = classifier(text, candidate_labels=candidate_labels)
    
    label = result["labels"][0]
    confidence = result["scores"][0]
    
    # Determine alert level and color
    if label == "Phishing Email":
        if confidence > 0.8:
            alert_html = """
                <div style="padding: 15px; background-color: #ffebee; border-left: 6px solid #f44336; margin-bottom: 15px;">
                    <h3 style="color: #d32f2f; margin: 0;">⚠️ High Risk - Likely Phishing Attempt</h3>
                </div>
            """
        else:
            alert_html = """
                <div style="padding: 15px; background-color: #fff3e0; border-left: 6px solid #ff9800; margin-bottom: 15px;">
                    <h3 style="color: #f57c00; margin: 0;">⚠️ Medium Risk - Suspicious Content Detected</h3>
                </div>
            """
    else:
        alert_html = """
            <div style="padding: 15px; background-color: #e8f5e9; border-left: 6px solid #4caf50; margin-bottom: 15px;">
                <h3 style="color: #2e7d32; margin: 0;">✅ Low Risk - Likely Legitimate</h3>
            </div>
        """
    
    # Keyword Analysis
    found_phrases = []
    text_lower = text.lower()
    for phrase in SUSPICIOUS_PHRASES:
        if phrase in text_lower:
            found_phrases.append(phrase)
    
    # Generate Detailed Report
    report = [
        f"### Analysis Results",
        f"**Confidence Score:** {confidence:.1%}",
        f"**Classification:** {label}",
        "\n### Detected Indicators"
    ]
    
    if found_phrases:
        report.append("🚩 **Suspicious elements found:**")
        for phrase in found_phrases:
            report.append(f"- Found '{phrase}'")
    else:
        report.append("✅ No common suspicious phrases detected.")
    
    report.append("\n### Recommendation")
    if confidence > 0.9:
        report.append("🔴 **High confidence in classification - exercise extreme caution!**")
    elif confidence > 0.7:
        report.append("🟡 **Moderate confidence - review carefully and verify sender.**")
    else:
        report.append("🟢 **Low confidence - but always remain vigilant.**")
        
    return alert_html, "\n".join(report), gr.update(visible=True)

def process_input(text_input, image_input):
    """Process either text input or image input."""
    # If text is provided, use it directly
    if text_input.strip():
        return analyze_text(text_input)
    
    # If image is provided, extract text first
    if image_input is not None:
        extracted_text = extract_text_from_image(image_input)
        if extracted_text.strip():
            return analyze_text(extracted_text)
        return (
            """<div style="padding: 15px; background-color: #fff3e0; border-left: 6px solid #ff9800; margin-bottom: 15px;">
                <h3 style="color: #f57c00; margin: 0;">⚠️ OCR Processing Error</h3>
            </div>""",
            "Could not extract text from image. Please ensure the image contains clear, readable text.",
            gr.update(visible=False)
        )
    
    return (
        "",
        "Please provide either text or an image to analyze.",
        gr.update(visible=False)
    )

# Custom theme
custom_theme = gr.themes.Soft().set(
    body_background_fill='*background-1',
    block_background_fill='*background-2',
    block_label_background_fill='*background-3',
    input_background_fill='*background-2',
    button_primary_background_fill='*primary-500',
    button_primary_text_color='white',
)

# Create Gradio interface with tabs
with gr.Blocks(theme=custom_theme) as demo:
    gr.Markdown("""
        # 📧 AI Phishing Detector
        ### Analyze emails and messages for potential phishing attempts
        Upload a screenshot or paste text to check for suspicious content
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("""
                ### How to Use
                1. Either paste message text or upload a screenshot
                2. Click 'Analyze' to check for phishing indicators
                3. Review the detailed analysis results
                
                ℹ️ This tool uses AI to detect:
                - Suspicious language patterns
                - Common phishing phrases
                - Urgency indicators
                - Security threat language
            """)
        
        with gr.Column(scale=2):
            with gr.Tab("✏️ Text Input"):
                text_input = gr.Textbox(
                    lines=8,
                    label="Message Text",
                    placeholder="Paste email or message content here...",
                    elem_id="text_input"
                )
            
            with gr.Tab("📷 Screenshot Upload"):
                image_input = gr.Image(
                    label="Upload Screenshot",
                    type="pil",
                    elem_id="image_input"
                )
            
            analyze_button = gr.Button("🔍 Analyze", variant="primary", size="lg")
    
    with gr.Row():
        with gr.Column():
            alert_html = gr.HTML(label="Alert")
            analysis = gr.Markdown(label="Detailed Analysis")
            
            with gr.Accordion("ℹ️ About this tool", open=False):
                gr.Markdown("""
                    This tool uses advanced AI models to analyze messages for potential phishing attempts:
                    - OCR technology to extract text from screenshots
                    - Zero-shot classification for pattern detection
                    - Keyword analysis for suspicious content
                    
                    **Disclaimer:** This is an educational demo. Always verify suspicious messages with your IT department.
                """)
    
    # Examples
    examples = [
        ["Subject: URGENT - Account Security Alert\n\nDear User,\n\nWe detected unusual activity in your account. Click here immediately to verify your identity and reset your password. If you don't respond within 24 hours, your account will be suspended.\n\nBank Security Team", None],
        ["Subject: Team Meeting Tomorrow\n\nHi everyone,\n\nJust a reminder that we have our weekly team meeting tomorrow at 10 AM in the main conference room. Please bring your project updates.\n\nBest regards,\nSarah", None],
    ]
    gr.Examples(
        examples=examples,
        inputs=[text_input, image_input],
        outputs=[alert_html, analysis],
        fn=process_input,
        cache_examples=True
    )
    
    # Set up event handler
    analyze_button.click(
        fn=process_input,
        inputs=[text_input, image_input],
        outputs=[alert_html, analysis]
    )

# Launch the app
if __name__ == "__main__":
    demo.launch()