Prevent path traversal, refactor

app.py

@@ -1,100 +1,162 @@
-from fastapi import FastAPI, File, UploadFile, HTTPException
-from fastapi.responses import JSONResponse
+from fastapi import FastAPI, File, UploadFile, HTTPException, status
+from fastapi.responses import JSONResponse, FileResponse
+from contextlib import asynccontextmanager
+from pathlib import Path
 import tensorflow as tf
-import tensorflow.keras as keras
 import numpy as np
 import os
 import shutil
-from huggingface_hub import snapshot_download
 import cv2
+import logging
+import uuid
+from huggingface_hub import snapshot_download
+from typing import Optional
+import aiofiles
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+MAL_CLASSES = ['Adialer.C', 'Agent.FYI', 'Allaple.A', 'Allaple.L', 'Alueron.gen!J', 
+               'Autorun.K', 'C2LOP.P', 'C2LOP.gen!g', 'Dialplatform.B', 'Dontovo.A', 
+               'Fakerean', 'Instantaccess', 'Lolyda.AA1', 'Lolyda.AA2', 'Lolyda.AA3', 
+               'Lolyda.AT', 'Malex.gen!J', 'Obfuscator.AD', 'Rbot!gen', 'Skintrim.N', 
+               'Swizzor.gen!E', 'Swizzor.gen!I', 'VB.AT', 'Wintrim.BX', 'Yuner.A']
 
-MAL_CLASSES = ['Adialer.C', 'Agent.FYI', 'Allaple.A', 'Allaple.L', 'Alueron.gen!J', 'Autorun.K', 'C2LOP.P', 'C2LOP.gen!g', 'Dialplatform.B', 'Dontovo.A', 'Fakerean', 'Instantaccess', 'Lolyda.AA1', 'Lolyda.AA2', 'Lolyda.AA3', 'Lolyda.AT', 'Malex.gen!J', 'Obfuscator.AD', 'Rbot!gen', 'Skintrim.N', 'Swizzor.gen!E', 'Swizzor.gen!I', 'VB.AT', 'Wintrim.BX', 'Yuner.A']
 UPLOAD_DIR = "uploads"
 os.makedirs(UPLOAD_DIR, exist_ok=True)
 
-app = FastAPI()
+# Environment configuration
+MODEL_REPO = os.getenv("MODEL_REPO", "GranularFireplace/malware")
+MODEL_FILE = os.getenv("MODEL_FILE", "model_v2_with_weight.keras")
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Manage model loading and unloading during app lifecycle"""
+    try:
+        logger.info("Downloading model from Hugging Face Hub...")
+        download_dir = snapshot_download(MODEL_REPO)
+        app.state.model = tf.keras.models.load_model(os.path.join(download_dir, MODEL_FILE))
+        logger.info("Model loaded successfully")
+    except Exception as e:
+        logger.error(f"Error loading model: {str(e)}")
+        raise
+    yield
+    # Cleanup resources if needed
+    app.state.model = None
+
+app = FastAPI(lifespan=lifespan)
 
 @app.get("/")
-def greet_json():
+async def greet_json():
     return {"Hello": "World!"}
 
-@app.post("/upload")
-def upload_file(file: UploadFile = File(...)):
+@app.post("/upload", status_code=status.HTTP_201_CREATED)
+async def upload_file(file: UploadFile = File(...)):
+    """Handle file uploads with async operations and path sanitization"""
     try:
-        # Save the uploaded file
-        file_path = os.path.join(UPLOAD_DIR, file.filename)
-        with open(file_path, "wb") as buffer:
-            shutil.copyfileobj(file.file, buffer)
-        return {"filename": file.filename, "message": "File uploaded successfully"}
+        # Sanitize filename to prevent path traversal
+        filename = Path(file.filename).name
+        file_path = os.path.join(UPLOAD_DIR, filename)
+        
+        async with aiofiles.open(file_path, "wb") as buffer:
+            content = await file.read()
+            await buffer.write(content)
+            
+        logger.info(f"File uploaded successfully: {filename}")
+        return {"filename": filename, "message": "File uploaded successfully"}
+    
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error uploading file: {str(e)}")
+        logger.error(f"Error uploading file: {str(e)}")
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=f"Error uploading file: {str(e)}"
+        )
 
 @app.get("/files")
-def list_files():
+async def list_files():
+    """List all uploaded files"""
     try:
         files = os.listdir(UPLOAD_DIR)
         return {"files": files}
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error listing files: {str(e)}")
+        logger.error(f"Error listing files: {str(e)}")
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=f"Error listing files: {str(e)}"
+        )
 
 @app.get("/download/{file_name}")
-def download_file(file_name: str):
-    file_path = os.path.join(UPLOAD_DIR, file_name)
-    if os.path.exists(file_path):
-        return JSONResponse(content={"message": "Download endpoint available. Use a client like cURL or a browser to download.", "file_path": file_path})
-    else:
-        raise HTTPException(status_code=404, detail="File not found")
-
-@app.get("/analyse/{file_name}")
-def analyse(file_name: str):
-    download_dir = snapshot_download("GranularFireplace/malware")
-    print(download_dir)
-    print(os.path.join(download_dir, 'model_v2_with_weight.keras'))
-    img = keras.preprocessing.image.load_img(
-        os.path.join(UPLOAD_DIR, file_name), target_size=(64, 64)
-    )
-    
-    img_array = keras.preprocessing.image.img_to_array(tf.image.rgb_to_grayscale(img))
-    
-    img_array = tf.expand_dims(img_array, 0) # Create a batch
+async def download_file(file_name: str):
+    """Serve files for download with proper security checks"""
+    # Sanitize input filename
+    sanitized_name = Path(file_name).name
+    file_path = os.path.join(UPLOAD_DIR, sanitized_name)
     
-    model = tf.keras.models.load_model(os.path.join(download_dir, 'model_v2_with_weight.keras'))
-    predicted_label = model.predict(img_array)
+    if not os.path.exists(file_path):
+        logger.warning(f"File not found: {sanitized_name}")
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND,
+            detail="File not found"
+        )
     
-    return {"result": MAL_CLASSES[np.argmax(predicted_label)]}
-
-def convert_binary_to_grayscale_image(binary_file_path, output_image_path, height=None, width=None):
-    with open(binary_file_path, "rb") as bin_file:
-        binary_data = bin_file.read()
-
-    grayscale_image = np.frombuffer(binary_data, dtype=np.uint8)
+    return FileResponse(file_path, filename=sanitized_name)
 
-    total_pixels = len(grayscale_image)
-    if height is None and width is None:
-        raise ValueError("Either height or width must be specified.")
-    elif height is None:
-        height = total_pixels // width
-    elif width is None:
-        width = total_pixels // height
-
-    if height * width != total_pixels:
-        raise ValueError(
-            f"The binary file size ({total_pixels}) is not compatible with the specified dimensions ({height}x{width})."
+def predict_malware(img_array: np.ndarray) -> str:
+    """Make prediction using the preloaded model"""
+    try:
+        prediction = app.state.model.predict(img_array)
+        return MAL_CLASSES[np.argmax(prediction)]
+    except Exception as e:
+        logger.error(f"Prediction error: {str(e)}")
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail="Error processing prediction"
         )
 
-    grayscale_image = grayscale_image.reshape((height, width))
-
-    # Save
-    cv2.imwrite(output_image_path, grayscale_image)
+async def process_image(file_path: str, target_size: tuple = (64, 64)) -> np.ndarray:
+    """Process image file for model input"""
+    try:
+        img = tf.keras.utils.load_img(file_path, target_size=target_size, color_mode="grayscale")
+        img_array = tf.keras.utils.img_to_array(img)
+        return tf.expand_dims(img_array, axis=0)
+    except Exception as e:
+        logger.error(f"Image processing error: {str(e)}")
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail="Invalid image file format"
+        )
 
-    print(f"Grayscale image saved to {output_image_path} with dimensions ({height}, {width})")
-    return grayscale_image
+@app.get("/analyse/{file_name}")
+async def analyse(file_name: str):
+    """Analyze image files"""
+    sanitized_name = Path(file_name).name
+    file_path = os.path.join(UPLOAD_DIR, sanitized_name)
+    
+    if not os.path.exists(file_path):
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND,
+            detail="File not found"
+        )
+    
+    try:
+        img_array = await process_image(file_path)
+        result = predict_malware(img_array)
+        return {"result": result}
+    except HTTPException as he:
+        raise he
+    except Exception as e:
+        logger.error(f"Analysis error: {str(e)}")
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=f"Analysis failed: {str(e)}"
+        )
 
-def get_image_width(file_path):
-    file_size_kb = os.path.getsize(file_path) / 1024  # File size in kilobytes
+def get_image_width(file_path: str) -> int:
+    """Determine image width based on file size"""
+    file_size_kb = os.path.getsize(file_path) / 1024
 
-    # Define the file size range and corresponding image width
-    file_size_ranges = [
+    size_ranges = [
         (0, 10, 32),
         (10, 30, 64),
         (30, 60, 128),
@@ -105,31 +167,61 @@ def get_image_width(file_path):
         (1000, float('inf'), 1024)
     ]
 
-    # Determine the image width based on the file size
-    for lower_bound, upper_bound, width in file_size_ranges:
-        if lower_bound <= file_size_kb < upper_bound:
+    for lower, upper, width in size_ranges:
+        if lower <= file_size_kb < upper:
             return width
 
-    # If no range matches (unlikely due to the final range), raise an error
-    raise ValueError("File size is out of expected range.")
+    raise ValueError("File size out of expected range")
 
-@app.get("/analysebin/{file_name}")
-def analyse_bin(file_name: str):
-    download_dir = snapshot_download("GranularFireplace/malware")
-    print(download_dir)
-    print(os.path.join(download_dir, 'model_v2_with_weight.keras'))
+def convert_binary_to_image(binary_path: str, output_path: str, width: int):
+    """Convert binary file to grayscale image with error handling"""
+    try:
+        with open(binary_path, "rb") as f:
+            binary_data = f.read()
+
+        grayscale = np.frombuffer(binary_data, dtype=np.uint8)
+        height = len(grayscale) // width
+        grayscale = grayscale[:height*width].reshape((height, width))
+        
+        cv2.imwrite(output_path, grayscale)
+        logger.debug(f"Image converted: {output_path}")
+        
+    except Exception as e:
+        logger.error(f"Binary conversion error: {str(e)}")
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail="Invalid binary file format"
+        )
 
-    convert_binary_to_grayscale_image(os.path.join(UPLOAD_DIR, file_name), "image.png", width=get_image_width(os.path.join(UPLOAD_DIR, file_name)))
-    
-    img = keras.preprocessing.image.load_img(
-        "image.png", target_size=(64, 64)
-    )
-    
-    img_array = keras.preprocessing.image.img_to_array(tf.image.rgb_to_grayscale(img))
-    
-    img_array = tf.expand_dims(img_array, 0) # Create a batch
+@app.get("/analysebin/{file_name}")
+async def analyse_bin(file_name: str):
+    """Analyze binary files by converting to images"""
+    sanitized_name = Path(file_name).name
+    file_path = os.path.join(UPLOAD_DIR, sanitized_name)
     
-    model = tf.keras.models.load_model(os.path.join(download_dir, 'model_v2_with_weight.keras'))
-    predicted_label = model.predict(img_array)
+    if not os.path.exists(file_path):
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND,
+            detail="File not found"
+        )
+
+    temp_image = os.path.join(UPLOAD_DIR, f"temp_{uuid.uuid4()}.png")
     
-    return {"result": MAL_CLASSES[np.argmax(predicted_label)]}
+    try:
+        width = get_image_width(file_path)
+        convert_binary_to_image(file_path, temp_image, width)
+        img_array = await process_image(temp_image)
+        result = predict_malware(img_array)
+        return {"result": result}
+        
+    except HTTPException as he:
+        raise he
+    except Exception as e:
+        logger.error(f"Binary analysis error: {str(e)}")
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=f"Binary analysis failed: {str(e)}"
+        )
+    finally:
+        if os.path.exists(temp_image):
+            os.remove(temp_image)