Integrate VirusTotal for URL and file scanning. Add file analysis route.

app/config.py

@@ -5,6 +5,7 @@ load_dotenv(dotenv_path=os.path.join(os.path.dirname(__file__), "..", ".env"))
 
 GEMINI_API_KEY: str = os.environ["GEMINI_API_KEY"]
 HUGGING_FACE_TOKEN: str = os.environ["HUGGING_FACE_TOKEN"]
+VIRUS_TOTAL_API_KEY: str = os.environ.get("VIRUS_TOTAL", "")
 
 HF_IMAGE_MODEL = "dima806/deepfake_vs_real_image_detection"
 HF_AUDIO_MODEL = "mo-thecreator/Deepfake-audio-detection"

app/main.py

@@ -8,7 +8,7 @@ from fastapi.staticfiles import StaticFiles
 from fastapi.responses import FileResponse
 import os
 
-from app.routers import text, image, video, audio
+from app.routers import text, image, video, audio, file
 
 app = FastAPI(
     title="Anti-Phishing AI Backend",
@@ -27,6 +27,7 @@ app.include_router(text.router, prefix="/analyze", tags=["Text"])
 app.include_router(image.router, prefix="/analyze", tags=["Image"])
 app.include_router(video.router, prefix="/analyze", tags=["Video"])
 app.include_router(audio.router, prefix="/analyze", tags=["Audio"])
+app.include_router(file.router, prefix="/analyze", tags=["File"])
 
 FRONTEND_DIR = os.path.join(os.path.dirname(__file__), "..", "frontend")
 

app/routers/file.py

@@ -0,0 +1,47 @@
+from fastapi import APIRouter, HTTPException, UploadFile, File
+from app.models import AnalysisResult
+from app.tools.virustotal_tools import scan_file_virustotal
+
+router = APIRouter()
+
+def _risk_level(score: float) -> str:
+    if score < 0.3:
+        return "LOW"
+    elif score < 0.6:
+        return "MEDIUM"
+    elif score < 0.85:
+        return "HIGH"
+    return "CRITICAL"
+
+@router.post("/file", response_model=AnalysisResult)
+async def analyze_file(file: UploadFile = File(...)):
+    try:
+        content = await file.read()
+        vt_result = scan_file_virustotal(content, file.filename)
+        
+        # Calculate a simple risk score based on VT malicious count
+        # 0 malicious = 0 score, 5+ malicious = 1.0 score
+        malicious = vt_result.get("malicious", 0)
+        risk_score = min(malicious / 5.0, 1.0)
+        
+        threat_types = []
+        if malicious > 0:
+            threat_types.append("malicious_file")
+        if vt_result.get("suspicious", 0) > 0:
+            threat_types.append("suspicious_file")
+            
+        explanation = f"VirusTotal analysis for {file.filename}: {malicious} malicious detections."
+        if "status" in vt_result:
+            explanation = vt_result["message"]
+
+        return AnalysisResult(
+            risk_score=risk_score,
+            risk_level=_risk_level(risk_score),
+            threat_types=threat_types,
+            explanation=explanation,
+            tool_outputs={"virustotal_file": vt_result}
+        )
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+    finally:
+        await file.close()

app/tools/text_tools.py

@@ -17,8 +17,10 @@ def extract_urls(text: str) -> list[str]:
     return _SUSPICIOUS_TLD.findall(text)
 
 
-def score_url(url: str) -> dict:
+def score_url(url: str, use_vt: bool = True) -> dict:
     from urllib.parse import urlparse
+    from app.tools.virustotal_tools import scan_url_virustotal
+    
     parsed = urlparse(url)
     domain = parsed.netloc.lower()
     flags = []
@@ -39,12 +41,22 @@ def score_url(url: str) -> dict:
         flags.append("ip_address_url")
         is_suspicious = True
 
-    return {"url": url, "suspicious": is_suspicious, "flags": flags}
+    vt_result = {}
+    if use_vt:
+        vt_data = scan_url_virustotal(url)
+        if "malicious" in vt_data:
+            vt_result = vt_data
+            if vt_data["malicious"] > 0:
+                flags.append(f"virustotal_malicious:{vt_data['malicious']}")
+                is_suspicious = True
+
+    return {"url": url, "suspicious": is_suspicious, "flags": flags, "virustotal": vt_result}
 
 
 def analyze_urls_in_text(text: str) -> dict:
     urls = extract_urls(text)
-    scored = [score_url(u) for u in urls]
+    # Only use VT for top 3 URLs to avoid long wait times and rate limits
+    scored = [score_url(u, use_vt=(i < 3)) for i, u in enumerate(urls)]
     suspicious_count = sum(1 for s in scored if s["suspicious"])
     return {
         "urls_found": len(urls),

app/tools/virustotal_tools.py

@@ -0,0 +1,95 @@
+import requests
+import base64
+import hashlib
+import time
+from app.config import VIRUS_TOTAL_API_KEY
+
+def get_url_id(url: str) -> str:
+    """VT v3 uses base64 without padding for URL IDs."""
+    return base64.urlsafe_b64encode(url.encode()).decode().strip("=")
+
+def scan_url_virustotal(url: str) -> dict:
+    if not VIRUS_TOTAL_API_KEY:
+        return {"error": "VirusTotal API key not configured"}
+
+    url_id = get_url_id(url)
+    endpoint = f"https://www.virustotal.com/api/v3/urls/{url_id}"
+    headers = {
+        "x-apikey": VIRUS_TOTAL_API_KEY
+    }
+
+    try:
+        response = requests.get(endpoint, headers=headers)
+        if response.status_code == 200:
+            data = response.json()
+            stats = data.get("data", {}).get("attributes", {}).get("last_analysis_stats", {})
+            return {
+                "malicious": stats.get("malicious", 0),
+                "suspicious": stats.get("suspicious", 0),
+                "harmless": stats.get("harmless", 0),
+                "undetected": stats.get("undetected", 0),
+                "link": f"https://www.virustotal.com/gui/url/{url_id}"
+            }
+        elif response.status_code == 404:
+            return {"status": "not_found", "message": "URL not previously scanned by VirusTotal"}
+        else:
+            return {"error": f"VirusTotal API returned {response.status_code}", "details": response.text}
+    except Exception as e:
+        return {"error": str(e)}
+
+def scan_file_virustotal(file_content: bytes, filename: str) -> dict:
+    """
+    Scans a file using VirusTotal v3 API. 
+    First checks by hash, then uploads if not found.
+    """
+    if not VIRUS_TOTAL_API_KEY:
+        return {"error": "VirusTotal API key not configured"}
+
+    sha256_hash = hashlib.sha256(file_content).hexdigest()
+    headers = {"x-apikey": VIRUS_TOTAL_API_KEY}
+    
+    # 1. Check if hash already exists
+    endpoint = f"https://www.virustotal.com/api/v3/files/{sha256_hash}"
+    try:
+        response = requests.get(endpoint, headers=headers)
+        if response.status_code == 200:
+            return _parse_vt_file_response(response.json())
+        
+        # 2. If not found, upload the file
+        if response.status_code == 404:
+            upload_url = "https://www.virustotal.com/api/v3/files"
+            # For files > 32MB, we'd need a special upload URL, but let's assume standard for now
+            files = {"file": (filename, file_content)}
+            upload_response = requests.post(upload_url, headers=headers, files=files)
+            
+            if upload_response.status_code == 200:
+                analysis_id = upload_response.json().get("data", {}).get("id")
+                # In a real app, we might poll or return the ID. 
+                # For this tool, let's wait a few seconds and try to get the report by hash
+                # usually hash report is available shortly after upload if it's small
+                time.sleep(2) 
+                # Re-check by hash
+                response = requests.get(endpoint, headers=headers)
+                if response.status_code == 200:
+                    return _parse_vt_file_response(response.json())
+                return {"status": "queued", "analysis_id": analysis_id, "message": "File uploaded, analysis in progress."}
+            else:
+                return {"error": f"Upload failed: {upload_response.status_code}", "details": upload_response.text}
+                
+        return {"error": f"VT API error: {response.status_code}", "details": response.text}
+    except Exception as e:
+        return {"error": str(e)}
+
+def _parse_vt_file_response(data: dict) -> dict:
+    attributes = data.get("data", {}).get("attributes", {})
+    stats = attributes.get("last_analysis_stats", {})
+    return {
+        "malicious": stats.get("malicious", 0),
+        "suspicious": stats.get("suspicious", 0),
+        "harmless": stats.get("harmless", 0),
+        "undetected": stats.get("undetected", 0),
+        "file_type": attributes.get("type_description", "unknown"),
+        "size": attributes.get("size", 0),
+        "sha256": attributes.get("sha256", ""),
+        "link": f"https://www.virustotal.com/gui/file/{attributes.get('sha256')}"
+    }

tests/test_file_vt.py

@@ -0,0 +1,31 @@
+import os
+import sys
+import hashlib
+
+# Add the project root to sys.path
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+
+from app.tools.virustotal_tools import scan_file_virustotal
+
+def test_file_scan_hash():
+    # EICAR test file content (standard for testing antivirus)
+    eicar_content = b'X5O!P%@AP[4\\PZX54(P^)7CC)7}$EICAR-STANDARD-ANTIVIRUS-TEST-FILE!$H+H*'
+    print(f"Testing VT for EICAR file hash: {hashlib.sha256(eicar_content).hexdigest()}")
+    result = scan_file_virustotal(eicar_content, "eicar.com")
+    print("VT Result (EICAR):")
+    import json
+    print(json.dumps(result, indent=2))
+
+def test_file_scan_upload():
+    # A unique small file to trigger upload
+    import time
+    unique_content = f"This is a unique test file created at {time.time()}".encode()
+    print(f"\nTesting VT for unique file upload (hash: {hashlib.sha256(unique_content).hexdigest()})")
+    result = scan_file_virustotal(unique_content, "test.txt")
+    print("VT Result (Unique Upload):")
+    import json
+    print(json.dumps(result, indent=2))
+
+if __name__ == "__main__":
+    test_file_scan_hash()
+    test_file_scan_upload()

tests/test_vt.py

@@ -0,0 +1,25 @@
+import os
+import sys
+
+# Add the project root to sys.path
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+
+from app.tools.virustotal_tools import scan_url_virustotal
+from app.tools.text_tools import analyze_urls_in_text
+
+def test_vt():
+    url = "http://malware.testing.google.test/testing/malware/"
+    print(f"Testing VT for URL: {url}")
+    result = scan_url_virustotal(url)
+    print("VT Result:", result)
+
+def test_text_analysis():
+    text = "Please check this link: http://malware.testing.google.test/testing/malware/"
+    print(f"\nTesting text analysis for: {text}")
+    result = analyze_urls_in_text(text)
+    import json
+    print("Analysis Result:", json.dumps(result, indent=2))
+
+if __name__ == "__main__":
+    test_vt()
+    test_text_analysis()