Update backend_structured/routes.py

backend_structured/routes.py

@@ -1,13 +1,53 @@
-from flask import Blueprint
+from flask import (
+    Flask, Blueprint, request, jsonify, current_app, send_from_directory,
+    send_file, render_template, abort, g, Response, make_response
+)
 api_bp = Blueprint("api", __name__)
+
 import sys
 import os
-sys.path.insert(0, os.path.abspath('backend'))
+import re
+import time
+import json
+import math
+import uuid
+import html
+import io
+import hashlib
+import sqlite3
+import socket
+import ssl
+import base64
+import bcrypt
+import jwt
+import requests
+import urllib3
+import ipaddress
+import queue
+import threading
+import statistics
+import itertools
+import traceback
+import concurrent.futures
+from datetime import datetime, timezone, timedelta
+from typing import Any, Callable, Literal
+from collections import defaultdict
+from functools import wraps
+from urllib.parse import urljoin, urlparse
+from dotenv import load_dotenv
+
+load_dotenv()
+
+BASE_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
+BACKEND_DIR = os.path.join(BASE_DIR, 'backend')
+if BASE_DIR not in sys.path:
+    sys.path.insert(0, BASE_DIR)
+if BACKEND_DIR not in sys.path:
+    sys.path.insert(0, BACKEND_DIR)
 
 from bs4 import BeautifulSoup
 from celery import Celery
 from celery.schedules import crontab
-from collections import defaultdict
 from scanners.base_scanner import (
     active_scan_logs, add_log, get_scan_logs, parse_domain,
     cleanup_scan_logs, schedule_log_cleanup, emit_scan_progress
@@ -16,80 +56,45 @@ from scanners import get_pipeline, get_phases, build_scanner, apply_scan_options
 from utils.fuzzer_engine import ContextAwareFuzzer
 from cryptography import x509
 from cryptography.hazmat.backends import default_backend
-from datetime import datetime, timezone
-from datetime import datetime, timezone, timedelta
-from datetime import datetime, timezone, timezone
-from dotenv import load_dotenv
-load_dotenv()
 
 import stripe
-from flask import Blueprint, request, jsonify, current_app, send_from_directory
 from werkzeug.utils import secure_filename
-from flask import Blueprint, send_file, jsonify, request
-from flask import Flask
-from flask import jsonify
-from flask import render_template
-from flask import request, abort, g, Response, make_response
 from flask_cors import CORS
 from flask_limiter import Limiter
 from flask_limiter.util import get_remote_address
 from flask_socketio import SocketIO, emit, join_room, leave_room
 from flask_sqlalchemy import SQLAlchemy
-from functools import wraps
-from markupsafe import escape  # always available with Flask
+from markupsafe import escape
+
 from reportlab.lib import colors
 from reportlab.lib.pagesizes import letter
 from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
 from reportlab.pdfgen import canvas
-from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Table, TableStyle, PageBreak, Image, Flowable
+from reportlab.platypus import (
+    SimpleDocTemplate, Paragraph, Spacer, Table, TableStyle, PageBreak, Image, Flowable
+)
 from reportlab.graphics.shapes import Drawing
 from reportlab.graphics.charts.barcharts import VerticalBarChart
-from sqlalchemy import event
-from sqlalchemy import func
-from sqlalchemy import inspect, text
-from sqlalchemy import text
+
+from sqlalchemy import event, func, inspect, text
 from sqlalchemy.engine import Engine
-from typing import Any
-from typing import Any, Callable
-from typing import Callable
-from typing import Literal
-from urllib.parse import urljoin, urlparse
-from urllib.parse import urlparse
-import base64
-import bcrypt
-import concurrent.futures
-from backend.utils.email_service import (
-    send_welcome_email, 
-    send_scan_started, 
-    send_scan_completed, 
-    send_scan_failed,
-    send_critical_alert
-)
 
-import hashlib
-import html
-import io
-import itertools
-import json
-import jwt
-import math
-import os
-import re
-import re, time, ipaddress, os, hashlib, threading
-import requests
-import socket
-import sqlite3
-import ssl
-import statistics
-import threading
-import time
-import traceback
-import urllib.error
-import urllib.parse
-import urllib.request
-import urllib3
-import uuid
-import ipaddress
+try:
+    from backend.utils.email_service import (
+        send_welcome_email, 
+        send_scan_started, 
+        send_scan_completed, 
+        send_scan_failed,
+        send_critical_alert
+    )
+except ImportError:
+    from utils.email_service import (
+        send_welcome_email, 
+        send_scan_started, 
+        send_scan_completed, 
+        send_scan_failed,
+        send_critical_alert
+    )
 
 
 
@@ -814,39 +819,85 @@ def upload_organization_logo(current_user):
     if not org:
         return jsonify({'message': 'Organization not found'}), 404
         
-    if 'logo' not in request.files:
-        return jsonify({'message': 'No file part'}), 400
-        
-    file = request.files['logo']
-    if file.filename == '':
-        return jsonify({'message': 'No selected file'}), 400
-        
-    if file:
-        from utils.firebase_storage import is_available, upload_bytes
-        
-        filename = secure_filename(f"org_{org.id}_{file.filename}")
+    file_bytes = None
+    content_type = 'image/png'
+    ext = 'png'
+    
+    # 1. Direct File Upload (Multipart Form Data)
+    if 'logo' in request.files and request.files['logo'].filename != '':
+        file = request.files['logo']
+        filename_raw = secure_filename(file.filename) or "logo.png"
         file_bytes = file.read()
         content_type = file.content_type or 'image/png'
+        if '.' in filename_raw:
+            ext = filename_raw.rsplit('.', 1)[1].lower()
+    else:
+        # 2. Web Image URL Download Fallback (JSON or Form Data)
+        data = request.get_json(silent=True) or request.form or {}
+        logo_url_input = data.get('logo_url') or data.get('image_url')
         
-        if is_available():
-            # Upload to Firebase Cloud Storage
-            blob_path = f"logos/{org.id}/{filename}"
-            logo_url = upload_bytes(file_bytes, content_type, blob_path)
-            if not logo_url:
-                return jsonify({'message': 'Firebase upload failed'}), 500
-        else:
-            # Fallback: save to local disk
-            upload_folder = current_app.config['UPLOAD_FOLDER']
-            os.makedirs(upload_folder, exist_ok=True)
-            file_path = os.path.join(upload_folder, filename)
-            with open(file_path, 'wb') as f:
-                f.write(file_bytes)
-            logo_url = f"/uploads/logos/{filename}"
-        
-        org.report_logo_url = logo_url
-        db.session.commit()
-        
-        return jsonify({'message': 'Logo uploaded successfully', 'report_logo_url': logo_url}), 200
+        if logo_url_input and (logo_url_input.startswith('http://') or logo_url_input.startswith('https://')):
+            try:
+                resp = requests.get(
+                    logo_url_input,
+                    timeout=10,
+                    headers={'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36'}
+                )
+                if resp.status_code == 200 and resp.content:
+                    file_bytes = resp.content
+                    content_type = resp.headers.get('Content-Type', 'image/png').split(';')[0].strip()
+                    if 'jpeg' in content_type: ext = 'jpg'
+                    elif 'webp' in content_type: ext = 'webp'
+                    elif 'svg' in content_type: ext = 'svg'
+                    elif 'gif' in content_type: ext = 'gif'
+                    else: ext = 'png'
+            except Exception as e:
+                current_app.logger.warn(f"Failed to fetch logo_url: {e}")
+
+    if not file_bytes:
+        return jsonify({'message': 'No valid logo file or URL provided'}), 400
+
+    from utils.firebase_storage import is_available, upload_bytes
+    
+    filename = secure_filename(f"org_{org.id}_{int(time.time())}.{ext}")
+    logo_url = None
+    
+    if is_available():
+        # Upload to Firebase Cloud Storage
+        blob_path = f"logos/{org.id}/{filename}"
+        logo_url = upload_bytes(file_bytes, content_type, blob_path)
+        if not logo_url:
+            current_app.logger.warning(f"[Firebase] Upload failed for {blob_path}. Falling back to local disk storage.")
+
+    if not logo_url:
+        # Fallback: save to local disk
+        upload_folder = current_app.config.get('UPLOAD_FOLDER', os.path.join(os.getcwd(), 'uploads', 'logos'))
+        os.makedirs(upload_folder, exist_ok=True)
+        file_path = os.path.join(upload_folder, filename)
+        with open(file_path, 'wb') as f:
+            f.write(file_bytes)
+        logo_url = f"/uploads/logos/{filename}"
+    
+    org.report_logo_url = logo_url
+    db.session.commit()
+    
+    return jsonify({'message': 'Logo uploaded successfully', 'report_logo_url': logo_url}), 200
+
+
+@auth_bp.route('/uploads/logos/<filename>', methods=['GET'])
+@auth_bp.route('/uploads/<path:filename>', methods=['GET'])
+def serve_uploaded_logo(filename):
+    upload_folder = current_app.config.get('UPLOAD_FOLDER', os.path.join(os.getcwd(), 'uploads', 'logos'))
+    # Direct check in logos directory
+    if os.path.exists(os.path.join(upload_folder, filename)):
+        return send_from_directory(upload_folder, filename)
+    
+    # Parent uploads directory check
+    base_uploads = os.path.abspath(os.path.join(upload_folder, '..'))
+    if os.path.exists(os.path.join(base_uploads, filename)):
+        return send_from_directory(base_uploads, filename)
+
+    return jsonify({'message': 'File not found'}), 404
 
 @auth_bp.route('/organizations/webhook', methods=['GET', 'PUT'])
 @token_required