fix_phase2.py

#!/usr/bin/env python3
"""Apply Phase 2 High Priority Fixes"""

with open('bespo_notifications/bespo_notifications/bot/commands.py', 'r') as f:
    content = f.read()

# Fix 2.1: Add validate_doc_id function after imports
old_imports = '''from .helpers import validate_user_status, get_or_register_user, format_file_size, truncate'''
new_imports = '''from .helpers import validate_user_status, get_or_register_user, format_file_size, truncate


def validate_doc_id(doc_id: str) -> bool:
    """
    Validate document ID format to prevent injection attacks.
    
    Valid formats:
    - ABC-123
    - ABC_123
    - ABC123
    
    Invalid:
    - <script> tags
    - SQL keywords
    - Path traversal
    """
    if not doc_id:
        return False
    # Allow alphanumeric with hyphens and underscores
    if not all(c.isalnum() or c in '-_' for c in doc_id):
        return False
    # Reasonable length limit
    if len(doc_id) > 100:
        return False
    return True'''

content = content.replace(old_imports, new_imports)

# Fix 2.1a: Add validation to cmd_status after "if not doc_id_prefix"
old_status = '''    if not doc_id_prefix:
        await update.message.reply_text(
            "Usage: /status `<document_id>`", parse_mode="Markdown"
        )
        return

    doc_svc = DocumentService()'''

new_status = '''    if not doc_id_prefix:
        await update.message.reply_text(
            "Usage: /status `<document_id>`", parse_mode="Markdown"
        )
        return

    # BUG FIX 2.1: Validate document ID to prevent injection
    if not validate_doc_id(doc_id_prefix):
        await update.message.reply_text(
            "❌ Invalid document ID format. Use only alphanumeric characters, hyphens, and underscores."
        )
        return

    doc_svc = DocumentService()'''

content = content.replace(old_status, new_status)

# Fix 2.1b: Add validation to cmd_history
old_history = '''    if not doc_id_prefix:
        await update.message.reply_text(
            "Usage: /history <code><document_id></code>", parse_mode="HTML"
        )
        return

    doc_svc = DocumentService()'''

new_history = '''    if not doc_id_prefix:
        await update.message.reply_text(
            "Usage: /history <code><document_id></code>", parse_mode="HTML"
        )
        return

    # BUG FIX 2.1: Validate document ID to prevent injection
    if not validate_doc_id(doc_id_prefix):
        await update.message.reply_text(
            "❌ Invalid document ID format. Use only alphanumeric characters, hyphens, and underscores."
        )
        return

    doc_svc = DocumentService()'''

content = content.replace(old_history, new_history)

# Fix 2.1c: Add validation to cmd_cancel_doc
old_cancel = '''    if len(parts) < 2:
        await update.message.reply_text(
            "Usage: /cancel <code><document_id></code>", parse_mode="HTML"
        )
        return

    doc_id = parts[1]
    doc_svc = DocumentService()'''

new_cancel = '''    if len(parts) < 2:
        await update.message.reply_text(
            "Usage: /cancel <code><document_id></code>", parse_mode="HTML"
        )
        return

    doc_id = parts[1]
    
    # BUG FIX 2.1: Validate document ID to prevent injection
    if not validate_doc_id(doc_id):
        await update.message.reply_text(
            "❌ Invalid document ID format. Use only alphanumeric characters, hyphens, and underscores."
        )
        return
    
    doc_svc = DocumentService()'''

content = content.replace(old_cancel, new_cancel)

with open('bespo_notifications/bespo_notifications/bot/commands.py', 'w') as f:
    f.write(content)

print("Phase 2 Fix 2.1 applied to commands.py")