main.py

"""
VeriScan Global Scoring Protocol:
---------------------------------
All algorithms are standardized to a "Suspicion Score" range: [0.0 - 1.0]
- 0.0: High Authenticity / Clean
- 1.0: High Risk / Manipulated
- Scaling: UI components multiply by 100 for percentage display.

Algorithm Output Ranges (Pre-Normalization):
- ELA: 0-100 (Internal) -> Normalized to 0.0-1.0
- MOVE-COPY: 0.0-1.0 (Lowe's Ratio based)
- NOISE-VAR: Raw Sigma (0-50+) -> Normalized to 0.0-1.0
- CNN-DETECT: Probability (0.0-1.0)
- SRNet: Stego Probability (0.0-1.0)

Auther: Elal Gilboa
"""

from fastapi import FastAPI, UploadFile, File, HTTPException, Form
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from PIL import Image
import numpy as np
import io
import base64
import logging
from typing import Optional, Dict, Any

from verification_manager.image_authentication_engine.compression_artifact_analyzer import ELAAuthenticator
from verification_manager.image_authentication_engine.noise_consistency_analyzer import NoiseConsistencyAnalyzer
from verification_manager.image_authentication_engine.copy_move_forg_detect_alg import MoveCopyForgeryDetector
from verification_manager.cnn_detection_engine.cnn_detector import CNNDetector
from verification_manager.injection_defense_engine.image_injection_detector import ImageInjectionDetector
from verification_manager.injection_defense_engine.image_injection_neutralizer import ImageInjectionNeutralizer
import verification_manager.injection_defense_engine.injection_detector_constants as Keys

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Initialize detectors globally to load weights only once
cnn_detector = CNNDetector()  # Loads .pth file once at startup
injection_detector = ImageInjectionDetector()  # Steganography detection
injection_neutralizer = ImageInjectionNeutralizer()  # Image sanitization

# Initialize FastAPI app
app = FastAPI(title="VeriScan API", version="1.0.0")

# Add CORS middleware to allow frontend requests
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"]
                   , )


# Pydantic models for request/response validation
class AnalysisRequest(BaseModel):
    method: str


class AnalysisResponse(BaseModel):
    method: str
    score: float
    description: str
    heatmap: Optional[str] = None  # Base64 encoded image


class InjectionDetectionResponse(BaseModel):
    score: float
    description: str
    heatmap: Optional[str] = None
    risk_level: str
    is_injection_detected: bool
    detailed_results: Optional[Dict[str, Any]] = None
    dominant_detection_method: str


class NeutralizationRequest(BaseModel):
    actions: list[str]  # ['blur', 'compress', or both]


class NeutralizationResponse(BaseModel):
    status: str
    message: str
    sanitized_image: str  # Base64 encoded


@app.get("/")
def root():
    """Health check endpoint."""
    return {"status": "VeriScan API is running"}


def _sanitize_for_json(obj):
    """
    Recursively sanitize an object to make it JSON-serializable.

    Converts:
    - NumPy arrays to Python lists
    - NumPy scalar types to Python native types
    - PIL Images to Base64 strings
    - Nested dictionaries and lists
    Args:
        obj: Object to sanitize
    Returns:
        JSON-serializable version of the object
    """
    if isinstance(obj, np.ndarray):
        return obj.tolist()
    elif isinstance(obj, np.integer):
        return int(obj)
    elif isinstance(obj, np.floating):
        return float(obj)
    elif isinstance(obj, np.bool_):
        return bool(obj)
    elif isinstance(obj, Image.Image):
        # Convert PIL Image to Base64
        return _image_to_base64(obj)
    elif isinstance(obj, dict):
        return {key: _sanitize_for_json(value) for key, value in obj.items()}
    elif isinstance(obj, (list, tuple)):
        return [_sanitize_for_json(item) for item in obj]
    else:
        return obj


@app.post("/analyze", response_model=AnalysisResponse)
async def analyze_image(file: UploadFile = File(...), method: str = Form("ELA")):
    """
    Analyze an uploaded image using the specified forensic algorithm.

    Args:
        file: Uploaded image file (JPG, JPEG, or PNG).
        method: Analysis method (ELA, MOVE-COPY, NOISE-VAR, CNN-DETECT).

    Returns:
        AnalysisResponse: Contains score, description, and optional Base64 heatmap.

    Raises:
        HTTPException: If file validation fails or analysis encounters an error.
    """
    try:
        # Validate file extension
        valid_extensions = {'.jpg', '.jpeg', '.png', '.pgm'}
        file_ext = '.' + file.filename.split('.')[-1].lower()

        if file_ext not in valid_extensions:
            raise HTTPException(status_code=400, detail=f"Invalid file type. Supported: {', '.join(valid_extensions)}")

        # Read file bytes and convert to PIL Image (RGB)
        file_bytes = await file.read()
        image = Image.open(io.BytesIO(file_bytes)).convert('RGB')
        image_array = np.array(image, dtype=np.uint8)

        logger.info(f"Processing image with method: {method}")

        # Route to appropriate analyzer
        if method.upper() == "ELA":
            result = _run_ela_analysis(image_array)
        elif method.upper() == "MOVE-COPY":
            result = _run_move_copy_analysis(image_array)
        elif method.upper() == "NOISE-VAR":
            result = _run_noise_variance_analysis(image_array)
        elif method.upper() == "CNN-DETECT":
            result = _run_cnn_detection(image_array)
        else:
            raise HTTPException(
                status_code=400,
                detail=f"Unknown analysis method: {method}")

        logger.info(f"Analysis completed. Score: {result['score']}")

        # Build response
        response_data = {
            "method": method,
            Keys.SCORE: result.get('score', 0.0),
            Keys.DESCRIPTION: result.get('description', 'Analysis completed'),
        }

        # Add heatmap if it exists (not for CNN-DETECT)
        if Keys.HEATMAP in result and result[Keys.HEATMAP] is not None:
            heatmap_b64 = _image_to_base64(result[Keys.HEATMAP])
            response_data[Keys.HEATMAP] = heatmap_b64

        return AnalysisResponse(**response_data)

    except HTTPException as e:
        raise e
    except Exception as e:
        logger.error(f"Error during analysis: {str(e)}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"Analysis failed: {str(e)}")


@app.post("/detect-injection", response_model=InjectionDetectionResponse)
async def run_steganographic_injection_analysis(file: UploadFile = File(...)):
    """
    Detect steganographic prompt injections in uploaded image by running multiple detection algorithms and aggregating
    results. Algorithm returns score value in the 0.0-1.0 range, which is converted to percentage for frontend display.
    The response includes a risk level and detailed results from each sub-analyzer to provide transparency and
    actionable insights for users.
    Args:
        file: Uploaded image file (JPG, JPEG, or PNG).
    Returns:
        InjectionDetectionResponse: Detection results with risk assessment contains:
        - score: Aggregated trust score (0-100%)
        - description: Summary of findings and risk level
        - heatmap: Base64-encoded heatmap image (if applicable)
        - risk_level: Classified risk level (low, medium, high)
        - is_injection_detected: Boolean indicating if injection is likely
        - detailed_results: Comprehensive diagnostics from each sub-analyzer for transparency
        - dominant_detection_method: The method that contributed most to the final score
    Raises:
        HTTPException: If file validation or detection fails.
    """
    try:
        # Validate file extension
        valid_extensions = {'.jpg', '.jpeg', '.png', '.pgm'}
        file_ext = '.' + file.filename.split('.')[-1].lower()

        if file_ext not in valid_extensions:
            raise HTTPException(
                status_code=400,
                detail=f"Invalid file type. Supported: {', '.join(valid_extensions)}")

        # Read and convert image
        file_bytes = await file.read()
        image = Image.open(io.BytesIO(file_bytes)).convert('RGB')
        image_array = np.array(image, dtype=np.uint8)

        logger.info("Running advanced steganographic injection detection...")

        # Run injection detection
        detection_results = injection_detector.analyze(image_array)
        detailed_data = detection_results.get(Keys.DETAILED, {})
        if Keys.NEURAL_DETECTOR in detailed_data:
            neural_detailed = detailed_data[Keys.NEURAL_DETECTOR].get(Keys.DETAILED, {})
            if Keys.FULLY_PROFILE in neural_detailed:
                neural_detailed[Keys.FULLY_PROFILE] = {
                    algo: float(prob * 100.0)  # Convert to percentage for frontend display
                    for algo, prob in neural_detailed[Keys.FULLY_PROFILE].items()
                }

        # Convert heatmap to Base64 if it exists
        heatmap_b64 = None
        if detection_results.get(Keys.HEATMAP) is not None:
            heatmap_b64 = _image_to_base64(Image.fromarray(detection_results[Keys.HEATMAP]))

        # Return structured response using centralized Keys to prevent KeyErrors
        return InjectionDetectionResponse(
            score=detection_results[Keys.SCORE] * 100.0,  # Convert to percentage for frontend display
            description=detection_results[Keys.DESCRIPTION],
            heatmap=heatmap_b64,
            risk_level=detection_results[Keys.RISK_LEVEL],
            is_injection_detected=detection_results[Keys.IS_DETECTED],
            detailed_results=_sanitize_for_json(detection_results.get(Keys.DETAILED, {})),
            dominant_detection_method=detection_results.get(Keys.DOMINANT_METHOD, 'Unknown'))

    except HTTPException as e:
        raise e
    except Exception as e:
        logger.error(f"Injection detection error: {str(e)}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"Detection failed: {str(e)}")


@app.post("/neutralize-injection", response_model=NeutralizationResponse)
async def neutralize_injection(file: UploadFile = File(...), actions: str = Form("compress")):
    """
    Sanitize image by neutralizing steganographic injections.
    Args:
        file: Uploaded image file.
        actions: Comma-separated actions ('blur', 'compress', or both).
    Returns:
        NeutralizationResponse: Sanitized image as Base64 and status message.
    Raises:
        HTTPException: If file or neutralization fails.
    """
    try:
        # Validate file
        valid_extensions = {'.jpg', '.jpeg', '.png', '.pgm'}
        file_ext = '.' + file.filename.split('.')[-1].lower()

        if file_ext not in valid_extensions:
            raise HTTPException(
                status_code=400,
                detail=f"Invalid file type. Supported: {', '.join(valid_extensions)}"
            )

        # Parse actions
        action_list = [a.strip().lower() for a in actions.split(',') if a.strip()]
        valid_actions = {'blur', 'compress'}

        if not all(a in valid_actions for a in action_list):
            raise HTTPException(
                status_code=400,
                detail=f"Invalid actions. Supported: {valid_actions}"
            )

        # Read and convert image
        file_bytes = await file.read()
        image = Image.open(io.BytesIO(file_bytes)).convert('RGB')
        image_array = np.array(image, dtype=np.uint8)

        logger.info(f"Neutralizing image with actions: {action_list}")

        # Apply neutralization
        sanitized_array = injection_neutralizer.neutralize_injection(
            image_array,
            action_list
        )

        # Convert to PIL and encode as Base64
        sanitized_image = Image.fromarray(sanitized_array)
        sanitized_b64 = _image_to_base64(sanitized_image)

        message = f"Image sanitized using: {', '.join(action_list)}"
        if not action_list:
            message = "No sanitization actions applied (image unchanged)"

        return NeutralizationResponse(
            status="success",
            message=message,
            sanitized_image=sanitized_b64
        )

    except HTTPException as e:
        raise e
    except Exception as e:
        logger.error(f"Neutralization error: {str(e)}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"Neutralization failed: {str(e)}")


def _run_ela_analysis(image_array: np.ndarray) -> dict:
    """Run ELA algorithm on image array.
    Args:
        image_array (np.ndarray): Input image as a NumPy array.
    Returns:
        dict: Contains 'score', 'description', and 'heatmap' (PIL Image).
    NOTE: ELA returns values in a 0-100 range so they are returned as is
    """
    # by
    ela_authenticator = ELAAuthenticator(quality=95)
    results = ela_authenticator.analyze(image_array)
    heatmap_array = results.get(Keys.HEATMAP)

    return {
        Keys.SCORE: float(results.get(Keys.SCORE, 0.0) * 2.0),
        Keys.DESCRIPTION: results.get(Keys.DESCRIPTION, ''),
        Keys.HEATMAP: Image.fromarray(heatmap_array) if heatmap_array is not None else None
    }


def _run_move_copy_analysis(image_array: np.ndarray) -> dict:
    """Run Move-Copy Forgery Detection algorithm.
    Algorithm returns a score in the 0.0-1.0 range based on Lowe's Ratio, so we convert it to percentage for
    consistency.
    Args:
        image_array (np.ndarray): Input image as a NumPy array.
    Returns:
        dict: Contains 'score', 'description', and 'heatmap' (PIL Image).
        """
    detector = MoveCopyForgeryDetector()
    result = detector.analyze(image_array)

    return {
        Keys.SCORE: float(result.get(Keys.SCORE, 0.0) * 100.0),  # Convert to percentage
        Keys.DESCRIPTION: result.get(Keys.DESCRIPTION, ''),
        Keys.HEATMAP: result.get(Keys.HEATMAP)
    }


def _run_noise_variance_analysis(image_array: np.ndarray) -> dict:
    """Run Noise Consistency Variance analysis.
    Algorithm returns a normalized score in the 0.0-1.0 range based on global noise estimation, so we convert it to
    percentage for consistency.
    Args:
        image_array (np.ndarray): Input image as a NumPy array.
    Returns:
        dict: Contains 'score', 'description', and 'heatmap' (PIL Image)"""
    analyzer = NoiseConsistencyAnalyzer()
    result = analyzer.analyze(image_array)
    score = float(result.get(Keys.SCORE, 0.0) * 100.0)  # Convert to percentage

    return {
        Keys.SCORE: score,
        Keys.DESCRIPTION: result.get(Keys.DESCRIPTION, ''),
        Keys.HEATMAP: result.get(Keys.HEATMAP)
    }


def _run_cnn_detection(image_array: np.ndarray) -> dict:
    """Run CNN-based detection.
    Algorithm returns a probability in the 0.0-1.0 range, so we convert it to percentage for consistency.
    Args:
        image_array (np.ndarray): Input image as a NumPy array.
    Returns:
        dict: Contains 'score' (percentage), 'description', and optional 'heatmap'
    """
    result = cnn_detector.analyze(image_array)

    # Convert probability (0-1) to percentage (0-100)
    prob: float = result.get(Keys.AI_PROBABILITY, 0.0)
    score: float = float(prob * 100)

    # Determine the status based on the threshold
    is_ai = result.get(Keys.IS_SUSPICIOUS, False)

    if is_ai:
        description = f"High probability of AI generation detected ({score:.1f}%). Findings consistent with {result['weights_used']} model."
    else:
        description = f"Image appears to be authentic. AI probability is low ({score:.1f}%)."

    return {
        Keys.SCORE: score,
        Keys.DESCRIPTION: description,
        Keys.HEATMAP: result.get(Keys.HEATMAP)
    }


def _image_to_base64(image: Image.Image) -> str:
    """
    Convert PIL Image to Base64 string.

    Args:
        image (Image.Image): PIL Image object.

    Returns:
        str: Data URL with Base64-encoded PNG image.
    """
    buffered = io.BytesIO()
    image.save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue()).decode()
    return f"data:image/png;base64,{img_str}"