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/', 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/', 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//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)