Blame for api/badges.py

from flask import Blueprint, render_template, request, redirect, url_for, flash, jsonify, send_file, session, make_response
from auth import jwt_required_with_logging, get_current_user_safe, refresh_jwt_for_user
from web.routes import login_required
from sqlalchemy import text, inspect
from datetime import datetime
from models import Badge, User
from app import db
import io
import json
from flask_jwt_extended import create_access_token, get_jwt

badges_bp = Blueprint('badges', __name__, template_folder='../templates')

# Helper function to check if a user is an admin (consistent with web.routes.is_admin_user_web)
def is_admin_user(user):
    """Check if user is admin using the same logic as is_admin_user_web"""
    return user and user.is_active and (user.is_admin or user.username == 'admin')


# Dynamically generate badge table fields documentation from SQLAlchemy models
from models import Attestation, PolicyResult

def _get_model_fields(model):
    fields = {}
    for column in model.__table__.columns:
        # Always prefer the doc field for description
        if column.doc:
            fields[column.name] = column.doc
        else:
            fields[column.name] = str(column.type)
    return fields

_BADGE_TABLE_FIELDS = {
    'attestations': _get_model_fields(Attestation),
    'policy_results': _get_model_fields(PolicyResult),
}

# In-memory badges store for initial setup (replace with DB later)
badges = []

@badges_bp.route('/badges')
@jwt_required_with_logging
def badges_list():
    user = get_current_user_safe()
    if not user:
        flash('Session expired. Please log in again.', 'warning')
        return redirect(url_for('web.login'))
    
    # Ensure JWT token is refreshed for subsequent requests
    refresh_jwt_for_user(user_id=user.id)
    
    # Check if user is admin using our helper function
    is_admin = is_admin_user(user)
    
    # Set is_admin in session for use in templates, including the modal dialog
    session['is_admin'] = is_admin
    
    # Log user information for debugging
    # admin_status = "Admin" if is_admin else "Regular User"
    # flash(f"User ID: {user.id}, Username: {user.username}, Status: {admin_status}", "info")
    
    # Get the user's own badges and global badges
    badges = Badge.query.filter(
        (Badge.owner_id == user.id) | 
        (Badge.is_global == True)
    ).order_by(Badge.last_modified.desc()).all()
    
    # Pass is_admin to the template directly in addition to the session variable
    return render_template('badges.html', badges=badges, is_admin=is_admin)

@badges_bp.route('/badges/new', methods=['GET', 'POST'])
@jwt_required_with_logging
def badges_new():
    """Create a new badge. Handles both initial form display (GET) and form submission (POST)."""
    # Common authentication check for both GET and POST
    user = get_current_user_safe()
    if not user:
        flash('Session expired. Please log in again.', 'warning')
        return redirect(url_for('web.login'))
        
    # Check if user is admin and set in session
    is_admin = is_admin_user(user)
    session['is_admin'] = is_admin
    
    # Initialize variables used for template rendering
    test_results = None
    test_query = None
    badge_data = None
    
    # ===== POST request handling =====
    if request.method == 'POST':
        # Parse form data based on content type
        if request.is_json:
            data = request.get_json()
            name = data.get('name', '')
            description = data.get('description', '')
            is_global = bool(data.get('is_global', False)) and is_admin  # Only admins can set global
            products = data.get('products', '')
            criteria = data.get('criteria', '')
            table = data.get('table', 'attestations')
            positive = bool(data.get('positive', False))
            visible = data.get('visible', True)
            action_type = data.get('actionType', 'save')
        else:
            name = request.form['name']
            description = request.form.get('description', '')
            is_global = bool(request.form.get('is_global')) and is_admin  # Only admins can set global
            products = request.form.get('products', '')
            criteria = request.form.get('criteria', '')
            table = request.form.get('table', 'attestations')
            positive = bool(request.form.get('positive'))
            visible = 'visible' in request.form
            action_type = request.form.get('actionType', 'save')

        # Set the owner to the current user
        owner_id = user.id

        # Create badge data structure for template rendering
        badge_data = {
            'name': name,
            'description': description,
            'is_global': is_global,
            'owner_id': owner_id,
            'products': products,
            'criteria': criteria,
            'table': table,
            'positive': positive,
            'visible': visible
        }
        
        # ==== Handle "Test Query" action ====
        if action_type == 'test_query':
            import re
            if not re.match(r'^[a-zA-Z0-9_ =\'/<>",.%()!\-]+$', criteria):
                test_results = 'Invalid criteria: only safe SQL characters allowed.'
                test_query = None
            else:
                # Build and execute the query
                user_id = user.id
                products_clean = (products or '').strip().lower()
                
                if products_clean == 'all':
                    product_clause = "product_key LIKE '%'"
                else:
                    product_list = [f"'{p.strip()}'" for p in products.split(',') if p.strip()]
                    if product_list:
                        product_clause = f"product_key IN ({', '.join(product_list)})"
                    else:
                        product_clause = "product_key LIKE '%'"
                        
                base_query = f"SELECT * FROM {table} WHERE user_id={user_id} AND {product_clause}"
                if criteria.strip():
                    base_query += f" AND {criteria.strip()}"
                base_query += " LIMIT 10"
                test_query = base_query
                
                try:
                    result = db.session.execute(text(base_query))
                    rows = result.fetchall()
                    
                    def row_to_dict(row):
                        if hasattr(row, '_asdict'):
                            return row._asdict()
                        try:
                            return dict(row)
                        except Exception:
                            return str(row)
                            
                    test_results = '\n'.join([str(row_to_dict(row)) for row in rows]) if rows else 'No results.'
                except Exception as e:
                    test_results = f'Query error: {e}'
            
            return render_template('badge_form.html', badge=badge_data, test_results=test_results, 
                                  test_query=test_query, is_admin=is_admin)
        
        # ==== Handle "Save" action ====
        elif action_type == 'save':
            from sqlalchemy.exc import IntegrityError
            
            # Check if a badge with this name already exists
            existing_badge = Badge.query.filter_by(name=name).first()
            if existing_badge:
                error_msg = "A badge with this name already exists. Please choose a different name."
                if request.headers.get('Accept') == 'application/json' or request.is_json:
                    return jsonify({'error': error_msg}), 400
                else:
                    flash(error_msg, 'danger')
                    return render_template('badge_form.html', badge=badge_data, error=error_msg, is_admin=is_admin)
            
            # Create and save the new badge
            badge = Badge(
                name=name,
                description=description,
                owner_id=owner_id,
                is_global=is_global,
                last_modified=datetime.utcnow(),
                products=products,
                criteria=criteria,
                table=table,
                positive=positive,
                visible=visible
            )
            db.session.add(badge)
            
            try:
                db.session.commit()
                flash('Badge created successfully.', 'success')
                # Refresh JWT after badge creation
                return refresh_jwt_for_user(user_id=user.id, redirect_url=url_for('badges.badges_list'))
            except IntegrityError as e:
                db.session.rollback()
                error_msg = "Badge name already exists. Please choose a different name."
                if request.headers.get('Accept') == 'application/json' or request.is_json:
                    return jsonify({'error': error_msg}), 400
                else:
                    flash(error_msg, 'danger')
                    return render_template('badge_form.html', badge=badge_data, error=error_msg, is_admin=is_admin)
            except Exception as e:
                db.session.rollback()
                error_msg = f"Failed to create badge: {e}"
                if request.headers.get('Accept') == 'application/json' or request.is_json:
                    return jsonify({'error': error_msg}), 500
                else:
                    flash(error_msg, 'danger')
                    return render_template('badge_form.html', badge=badge_data, error=error_msg, is_admin=is_admin)
    
    # ===== GET request handling =====
    # Just display the empty form
    return render_template('badge_form.html', badge=None, is_admin=is_admin)

@badges_bp.route('/badges/edit/<int:badge_id>', methods=['GET', 'POST'])
@jwt_required_with_logging
def badges_edit(badge_id):
    user = get_current_user_safe()
    if not user:
        flash('Session expired. Please log in again.', 'warning')
        return redirect(url_for('web.login'))
        
    # Check if user is admin using our helper function
    is_admin = is_admin_user(user)
    session['is_admin'] = is_admin
        
    badge = Badge.query.get_or_404(badge_id)
    
    # Check if user is authorized to access this badge
    # Users can only edit their own badges
    # But users can view and test global badges even if they didn't create them
    if badge.owner_id != user.id and not badge.is_global:
        flash('You are not authorized to edit this badge. You can only edit badges you created.', 'danger')
        return redirect(url_for('badges.badges_list'))
    
    # Non-admin users can't edit global badges, only test them
    if badge.is_global and not is_admin:
        if request.method == 'POST' and request.form.get('actionType') != 'test_query':
            flash('As a non-administrator, you cannot modify global badges.', 'warning')
            return redirect(url_for('badges.badges_list'))
    
    test_results = None
    test_query = None
    
    if request.method == 'POST':
        # Support both JSON and form submissions
        if request.is_json:
            data = request.get_json()
            positive = bool(data.get('positive', False))
            visible = data.get('visible', True)
            # Only admins can set global badges
            is_global = bool(data.get('is_global', False)) and is_admin
            badge_data = {
                'id': badge.id,
                'name': data.get('name', ''),
                'description': data.get('description', ''),
                'is_global': is_global,
                'products': data.get('products', ''),
                'criteria': data.get('criteria', ''),
                'table': data.get('table', 'attestations'),
                'positive': positive,
                'visible': visible
            }
            action_type = data.get('actionType', 'save')
        else:
            positive = bool(request.form.get('positive'))
            # Only admins can set global badges
            is_global = bool(request.form.get('is_global')) and is_admin
            badge_data = {
                'id': badge.id,
                'name': request.form['name'],
                'description': request.form.get('description', ''),
                'is_global': is_global,
                'products': request.form.get('products', ''),
                'criteria': request.form.get('criteria', ''),
                'table': request.form.get('table', 'attestations'),
                'positive': positive,
                'visible': 'visible' in request.form
            }
            action_type = request.form.get('actionType', 'save')
        
        if action_type == 'test_query':
            import re
            user = get_current_user_safe()
            criteria = request.form.get('criteria', '')
            products = request.form.get('products', '')
            table = request.form.get('table', 'attestations')
            
            if not re.match(r'^[a-zA-Z0-9_ =\'/<>",.%()!\-]+$', criteria):
                test_results = 'Invalid criteria: only safe SQL characters allowed.'
                test_query = None
            else:
                # Ensure user is authenticated, otherwise redirect to login
                if not user:
                    flash('Session expired. Please log in again.', 'warning')
                    return redirect(url_for('web.login'))
                
                user_id = user.id  # Now we're sure user is not None
                products_clean = (products or '').strip().lower()
                if products_clean == 'all':
                    product_clause = "product_key LIKE '%'"
                else:
                    product_list = [f"'{p.strip()}'" for p in products.split(',') if p.strip()]
                    if product_list:
                        product_clause = f"product_key IN ({', '.join(product_list)})"
                    else:
                        product_clause = "product_key LIKE '%'"
                base_query = f"SELECT * FROM {table} WHERE user_id={user_id} AND {product_clause}"
                if criteria.strip():
                    base_query += f" AND {criteria.strip()}"
                base_query += " LIMIT 10"
                test_query = base_query
                try:
                    result = db.session.execute(text(base_query))
                    rows = result.fetchall()
                    def row_to_dict(row):
                        if hasattr(row, '_asdict'):
                            return row._asdict()
                        try:
                            return dict(row)
                        except Exception:
                            return str(row)
                    test_results = '\n'.join([str(row_to_dict(row)) for row in rows]) if rows else 'No results.'
                except Exception as e:
                    test_results = f'Query error: {e}'
            
            return render_template('badge_form.html', badge=badge_data, test_results=test_results, test_query=test_query, badge_id=badge_id, is_admin=is_admin)
        
        elif action_type == 'save':
            # Non-admin users shouldn't be able to save changes to global badges
            if badge.is_global and not is_admin:
                error_msg = "As a non-administrator, you cannot modify global badges."
                if request.headers.get('Accept') == 'application/json' or request.is_json:
                    return jsonify({'error': error_msg}), 403
                else:
                    flash(error_msg, 'warning')
                    return redirect(url_for('badges.badges_list'))
            
            # Check if a different badge with the same name already exists (excluding current badge)
            existing_badge = Badge.query.filter(Badge.name == badge_data['name'], Badge.id != badge_id).first()
            if existing_badge:
                error_msg = "A badge with this name already exists. Please choose a different name."
                if request.headers.get('Accept') == 'application/json' or request.is_json:
                    return jsonify({'error': error_msg}), 400
                else:
                    flash(error_msg, 'danger')
                    return render_template('badge_form.html', badge=badge_data, error=error_msg, badge_id=badge_id, is_admin=is_admin)
            
            if request.is_json:
                badge.name = badge_data['name']
                badge.description = badge_data['description']
                # Only admins can set global flag
                if is_admin:
                    badge.is_global = badge_data['is_global']
                badge.products = badge_data['products']
                badge.criteria = badge_data['criteria']
                badge.table = badge_data['table']
                badge.positive = badge_data['positive']
                badge.visible = badge_data.get('visible', True)
            else:
                badge.name = request.form['name']
                badge.description = request.form.get('description', '')
                # Only admins can set global flag
                if is_admin:
                    badge.is_global = bool(request.form.get('is_global'))
                badge.products = request.form.get('products', '')
                badge.criteria = request.form.get('criteria', '')
                badge.table = request.form.get('table', 'attestations')
                badge.positive = positive
                badge.visible = 'visible' in request.form
            
            # Badge owner never changes - it's always the user who created it
            badge.last_modified = datetime.utcnow()
            db.session.commit()

            flash('Badge updated successfully.', 'success')
            return redirect(url_for('badges.badges_list'))
            
    return render_template('badge_form.html', badge=badge, badge_id=badge_id, is_admin=is_admin)

@badges_bp.route('/badges/delete/<int:badge_id>', methods=['POST'])
@jwt_required_with_logging
def badges_delete(badge_id):
    user = get_current_user_safe()
    if not user:
        flash('Session expired. Please log in again.', 'warning')
        return redirect(url_for('web.login'))
        
    badge = Badge.query.get_or_404(badge_id)
    
    # Users can only delete badges they own
    if badge.owner_id != user.id:
        flash('You are not authorized to delete this badge. You can only delete badges you created.', 'danger')
        return redirect(url_for('badges.badges_list'))
    
    db.session.delete(badge)
    db.session.commit()
    flash('Badge deleted.', 'info')
    return redirect(url_for('badges.badges_list'))


# --- Badge JSON API for modal editing ---
@badges_bp.route('/badges/<int:badge_id>/json')
@jwt_required_with_logging
def badge_json(badge_id):
    # Make sure to refresh the JWT token if needed
    user = get_current_user_safe()
    if not user:
        # Try to use session user if available
        from flask import session
        if 'user_id' in session:
            user_id = session.get('user_id')
            user = User.query.get(user_id)
            # If session has a user but get_current_user_safe failed, refresh JWT
            if user:
                refresh_jwt_for_user(user_id=user.id)
        if not user:
            return jsonify({'error': 'Unauthorized'}), 401
        
    badge = Badge.query.get(badge_id)
    if not badge:
        return jsonify({'error': 'Badge not found'}), 404
        
    # Check if user is authorized to view badge details
    # Users can view their own badges and global badges
    if not (badge.owner_id == user.id or badge.is_global):
        return jsonify({'error': 'Unauthorized access to badge details'}), 403
    
    # Check if user is admin
    is_admin = is_admin_user(user)
    
    # Include can_edit flag in response
    # Users can only edit their own badges, and non-admins can't edit global badges
    can_edit = badge.owner_id == user.id and (not badge.is_global or is_admin)
    
    return jsonify({
        'id': badge.id,
        'name': badge.name,
        'description': badge.description,
        'is_global': badge.is_global,
        'owner_id': badge.owner_id,
        'table': badge.table,
        'products': badge.products,
        'criteria': badge.criteria,
        'positive': badge.positive,
        'visible': badge.visible,
        'can_edit': can_edit
    })

# --- Badge settings import/export logic for reuse ---
def export_badge_settings_for_user(user):
    """Return badge settings for user as dict."""
    # Get the user's own badges and global badges
    badges = Badge.query.filter(
        (Badge.owner_id == user.id) | 
        (Badge.is_global == True)
    ).all()
        
    return {badge.name: badge.to_dict() for badge in badges}


# --- Badge settings API endpoints (must be after badges_bp is defined) ---
@badges_bp.route('/badges/settings/download', methods=['GET'])
@jwt_required_with_logging
def download_badge_settings():
    """Download exactly the badges shown on the screen for the current user as JSON."""
    user = get_current_user_safe()
    if not user:
        return jsonify({'error': 'Unauthorized'}), 401
    
    # Get the user's own badges and global badges
    badges = Badge.query.filter(
        (Badge.owner_id == user.id) | 
        (Badge.is_global == True)
    ).order_by(Badge.last_modified.desc()).all()
        
    data = {badge.name: badge.to_dict() for badge in badges}
    json_bytes = json.dumps(data, indent=2).encode('utf-8')
    return send_file(
        io.BytesIO(json_bytes),
        mimetype='application/json',
        as_attachment=True,
        download_name='badges_settings.json'
    )

@badges_bp.route('/badges/settings/upload', methods=['POST'])
@jwt_required_with_logging
def upload_badge_settings():
    """Upload badge settings from JSON. Updates existing or adds new. Returns info on accepted/rejected entries."""
    user = get_current_user_safe()
    if not user:
        return jsonify({'error': 'Unauthorized'}), 401
    try:
        if request.is_json:
            data = request.get_json()
        else:
            # Accept file upload as form-data
            file = request.files.get('file')
            if not file:
                return jsonify({'error': 'No file uploaded'}), 400
            data = json.load(file)
    except Exception as e:
        return jsonify({'error': f'Invalid JSON: {e}'}), 400

    required_fields = {'name', 'description', 'products', 'criteria', 'table', 'positive'}
    accepted, rejected = [], []
    for badge_name, badge_data in data.items():
        # Validate fields
        if not isinstance(badge_data, dict) or not required_fields.issubset(badge_data):
            rejected.append({'name': badge_name, 'reason': 'Missing required fields'})
            continue
        # Find existing badge by name
        badge = Badge.query.filter_by(name=badge_name).first()
        # Always set owner_id to current user
        owner_id = user.id
        visible = badge_data.get('visible', True)
        
        # Check if badge exists and belongs to this user
        if badge:
            if badge.owner_id != user.id:
                rejected.append({'name': badge_name, 'reason': 'Cannot update badge owned by another user'})
                continue
            
            # Update fields of user's own badge
            badge.description = badge_data['description']
            badge.products = badge_data['products']
            badge.criteria = badge_data['criteria']
            badge.table = badge_data['table']
            badge.positive = badge_data['positive']
            badge.visible = visible
            badge.last_modified = datetime.utcnow()
            db.session.commit()
            accepted.append({'name': badge_name, 'action': 'updated'})
        else:
            # Add new badge
            try:
                # Double check for name uniqueness across all badges
                if Badge.query.filter_by(name=badge_name).first():
                    rejected.append({'name': badge_name, 'reason': 'A badge with this name already exists'})
                    continue
                    
                # Create new badge owned by current user
                new_badge = Badge(
                    name=badge_name,  # Use the key from the JSON, not badge_data['name']
                    description=badge_data['description'],
                    owner_id=owner_id,
                    is_global=False,  # New badges through upload are never global
                    last_modified=datetime.utcnow(),
                    products=badge_data['products'],
                    criteria=badge_data['criteria'],
                    table=badge_data['table'],
                    positive=badge_data['positive'],
                    visible=visible
                )
                db.session.add(new_badge)
                db.session.commit()
                accepted.append({'name': badge_name, 'action': 'added'})
            except Exception as e:
                db.session.rollback()
                rejected.append({'name': badge_name, 'reason': f'Error adding: {e}'})
    return jsonify({'accepted': accepted, 'rejected': rejected})


@badges_bp.route('/badges/tables_doc')
@jwt_required_with_logging
def table_docs():
    """Render documentation for badge query tables and their fields."""
    inspector = inspect(db.engine)
    docs = []
    for table_name, field_map in _BADGE_TABLE_FIELDS.items():
        try:
            columns = inspector.get_columns(table_name)
        except Exception:
            columns = []
        for col in columns:
            docs.append({
                'table': table_name,
                'field': col.get('name'),
                'type': str(col.get('type')),
                'explanation': field_map.get(col.get('name'), '')
            })
    return render_template('badge_tables_doc.html', docs=docs)