| 1 | import asyncio |
| 2 | |
| 3 | import httpx |
| 4 | from fastapi import HTTPException |
| 5 | from fastapi import UploadFile |
| 6 | from loguru import logger |
| 7 | |
| 8 | from app.threat_intel.schema.virustotal import FileAnalysisResponse |
| 9 | from app.threat_intel.schema.virustotal import FileReportResponse |
| 10 | from app.threat_intel.schema.virustotal import FileSubmissionRequest |
| 11 | from app.threat_intel.schema.virustotal import FileSubmissionResponse |
| 12 | |
| 13 | |
| 14 | async def submit_file_to_virustotal( |
| 15 | api_key: str, |
| 16 | file: UploadFile, |
| 17 | request: FileSubmissionRequest, |
| 18 | ) -> FileSubmissionResponse: |
| 19 | """ |
| 20 | Submit a file to VirusTotal for analysis. |
| 21 | |
| 22 | Args: |
| 23 | api_key (str): The VirusTotal API key |
| 24 | file (UploadFile): The file to be analyzed |
| 25 | request (FileSubmissionRequest): Additional parameters for submission |
| 26 | |
| 27 | Returns: |
| 28 | FileSubmissionResponse: Response containing submission ID |
| 29 | |
| 30 | Raises: |
| 31 | HTTPException: If the submission fails |
| 32 | """ |
| 33 | url = "https://www.virustotal.com/api/v3/files" |
| 34 | |
| 35 | # Headers - exactly match the working example |
| 36 | headers = {"accept": "application/json", "x-apikey": api_key} |
| 37 | |
| 38 | # Prepare the file for upload |
| 39 | file_content = await file.read() |
| 40 | |
| 41 | # Reset file pointer for potential reuse |
| 42 | await file.seek(0) |
| 43 | |
| 44 | # Prepare the files dictionary - exactly match the working pattern |
| 45 | files = {"file": (file.filename, file_content, "application/octet-stream")} |
| 46 | |
| 47 | # Prepare form data if password is provided |
| 48 | data = {} |
| 49 | if request.password: |
| 50 | data["password"] = request.password |
| 51 | |
| 52 | logger.info(f"Submitting file {file.filename} to VirusTotal (size: {len(file_content)} bytes)") |
| 53 | |
| 54 | try: |
| 55 | async with httpx.AsyncClient(timeout=300.0) as client: # 5 minute timeout for file uploads |
| 56 | response = await client.post(url, headers=headers, files=files, data=data if data else None) |
| 57 | |
| 58 | # Log the actual request headers for debugging |
| 59 | logger.info(f"Request headers sent: {response.request.headers}") |
| 60 | logger.info(f"Response status: {response.status_code}") |
| 61 | |
| 62 | response.raise_for_status() |
| 63 | response_data = response.json() |
| 64 | |
| 65 | return FileSubmissionResponse( |
| 66 | data=response_data["data"], |
| 67 | success=True, |
| 68 | message=f"File {file.filename} submitted successfully for analysis", |
| 69 | ) |
| 70 | |
| 71 | except httpx.HTTPStatusError as e: |
| 72 | logger.error(f"HTTP error submitting file to VirusTotal: {e.response.status_code} - {e.response.text}") |
| 73 | |
| 74 | # Parse the error response to provide better error messages |
| 75 | try: |
| 76 | error_data = e.response.json() |
| 77 | error_message = error_data.get("error", {}).get("message", str(e.response.text)) |
| 78 | except (ValueError, KeyError): |
| 79 | error_message = str(e.response.text) |
| 80 | |
| 81 | if e.response.status_code == 400: |
| 82 | # Handle specific 400 errors |
| 83 | if "Invalid zip file" in error_message: |
| 84 | raise HTTPException( |
| 85 | status_code=400, |
| 86 | detail="File format not supported or corrupted. VirusTotal accepts executables, documents, archives, and other common file types.", |
| 87 | ) |
| 88 | elif "File too large" in error_message: |
| 89 | raise HTTPException( |
| 90 | status_code=413, |
| 91 | detail="File too large. Maximum file size is 32MB for free API keys, 650MB for premium.", |
| 92 | ) |
| 93 | elif "missing" in error_message.lower(): |
| 94 | raise HTTPException( |
| 95 | status_code=400, |
| 96 | detail="File upload failed. Please ensure the file is properly formatted and try again.", |
| 97 | ) |
| 98 | else: |
| 99 | raise HTTPException(status_code=400, detail=f"Bad request: {error_message}") |
| 100 | elif e.response.status_code == 429: |
| 101 | raise HTTPException(status_code=429, detail="Rate limit exceeded. Please try again later.") |
| 102 | elif e.response.status_code == 413: |
| 103 | raise HTTPException(status_code=413, detail="File too large. Maximum file size is 32MB for free API keys.") |
| 104 | else: |
| 105 | raise HTTPException(status_code=e.response.status_code, detail=f"Failed to submit file: {error_message}") |
| 106 | except httpx.RequestError as e: |
| 107 | logger.error(f"Request error submitting file to VirusTotal: {e}") |
| 108 | raise HTTPException(status_code=500, detail=f"Network error occurred: {str(e)}") |
| 109 | except Exception as e: |
| 110 | logger.error(f"Unexpected error submitting file to VirusTotal: {e}") |
| 111 | raise HTTPException(status_code=500, detail=f"Unexpected error occurred: {str(e)}") |
| 112 | |
| 113 | |
| 114 | async def get_file_analysis_status( |
| 115 | api_key: str, |
| 116 | analysis_id: str, |
| 117 | ) -> FileAnalysisResponse: |
| 118 | """ |
| 119 | Get the status of a file analysis. |
| 120 | |
| 121 | Args: |
| 122 | api_key (str): The VirusTotal API key |
| 123 | analysis_id (str): The analysis ID returned from file submission |
| 124 | |
| 125 | Returns: |
| 126 | FileAnalysisResponse: Current analysis status |
| 127 | |
| 128 | Raises: |
| 129 | HTTPException: If the request fails |
| 130 | """ |
| 131 | url = f"https://www.virustotal.com/api/v3/analyses/{analysis_id}" |
| 132 | headers = {"x-apikey": api_key} |
| 133 | |
| 134 | logger.info(f"Checking analysis status for ID: {analysis_id}") |
| 135 | |
| 136 | try: |
| 137 | async with httpx.AsyncClient(timeout=30.0) as client: |
| 138 | response = await client.get(url, headers=headers) |
| 139 | response.raise_for_status() |
| 140 | response_data = response.json() |
| 141 | |
| 142 | return FileAnalysisResponse(data=response_data["data"], success=True, message="Analysis status retrieved successfully") |
| 143 | |
| 144 | except httpx.HTTPStatusError as e: |
| 145 | logger.error(f"HTTP error getting analysis status: {e.response.status_code} - {e.response.text}") |
| 146 | raise HTTPException(status_code=e.response.status_code, detail=f"Failed to get analysis status: {e.response.text}") |
| 147 | except httpx.RequestError as e: |
| 148 | logger.error(f"Request error getting analysis status: {e}") |
| 149 | raise HTTPException(status_code=500, detail=f"Network error occurred: {str(e)}") |
| 150 | |
| 151 | |
| 152 | async def get_file_report( |
| 153 | api_key: str, |
| 154 | file_id: str, |
| 155 | ) -> FileReportResponse: |
| 156 | """ |
| 157 | Get the detailed analysis report for a file. |
| 158 | |
| 159 | Args: |
| 160 | api_key (str): The VirusTotal API key |
| 161 | file_id (str): The file ID (hash) to get report for |
| 162 | |
| 163 | Returns: |
| 164 | FileReportResponse: Detailed analysis report |
| 165 | |
| 166 | Raises: |
| 167 | HTTPException: If the request fails |
| 168 | """ |
| 169 | url = f"https://www.virustotal.com/api/v3/files/{file_id}" |
| 170 | headers = {"x-apikey": api_key} |
| 171 | |
| 172 | logger.info(f"Getting file report for ID: {file_id}") |
| 173 | |
| 174 | try: |
| 175 | async with httpx.AsyncClient(timeout=30.0) as client: |
| 176 | response = await client.get(url, headers=headers) |
| 177 | response.raise_for_status() |
| 178 | response_data = response.json() |
| 179 | |
| 180 | return FileReportResponse(data=response_data["data"], success=True, message="File report retrieved successfully") |
| 181 | |
| 182 | except httpx.HTTPStatusError as e: |
| 183 | logger.error(f"HTTP error getting file report: {e.response.status_code} - {e.response.text}") |
| 184 | if e.response.status_code == 404: |
| 185 | raise HTTPException(status_code=404, detail="File not found or not yet analyzed") |
| 186 | raise HTTPException(status_code=e.response.status_code, detail=f"Failed to get file report: {e.response.text}") |
| 187 | except httpx.RequestError as e: |
| 188 | logger.error(f"Request error getting file report: {e}") |
| 189 | raise HTTPException(status_code=500, detail=f"Network error occurred: {str(e)}") |
| 190 | |
| 191 | |
| 192 | async def submit_and_wait_for_analysis( |
| 193 | api_key: str, |
| 194 | file: UploadFile, |
| 195 | request: FileSubmissionRequest, |
| 196 | max_wait_time: int = 300, # 5 minutes |
| 197 | poll_interval: int = 10, # 10 seconds |
| 198 | ) -> FileReportResponse: |
| 199 | """ |
| 200 | Submit a file and wait for analysis to complete, then return the report. |
| 201 | |
| 202 | Args: |
| 203 | api_key (str): The VirusTotal API key |
| 204 | file (UploadFile): The file to be analyzed |
| 205 | request (FileSubmissionRequest): Additional parameters for submission |
| 206 | max_wait_time (int): Maximum time to wait in seconds |
| 207 | poll_interval (int): Time between status checks in seconds |
| 208 | |
| 209 | Returns: |
| 210 | FileReportResponse: Complete analysis report |
| 211 | |
| 212 | Raises: |
| 213 | HTTPException: If submission fails or analysis times out |
| 214 | """ |
| 215 | # Submit the file |
| 216 | submission_response = await submit_file_to_virustotal(api_key, file, request) |
| 217 | analysis_id = submission_response.data.id |
| 218 | |
| 219 | logger.info(f"File submitted, analysis ID: {analysis_id}. Waiting for completion...") |
| 220 | |
| 221 | # Wait for analysis to complete |
| 222 | waited_time = 0 |
| 223 | while waited_time < max_wait_time: |
| 224 | status_response = await get_file_analysis_status(api_key, analysis_id) |
| 225 | |
| 226 | if status_response.data.attributes.status == "completed": |
| 227 | # Analysis is complete, get the file hash from the analysis ID |
| 228 | # The analysis ID format is usually base64 encoded, but we can extract file hash from response |
| 229 | logger.info("Analysis completed, retrieving detailed report...") |
| 230 | |
| 231 | # For now, we'll return the analysis status response |
| 232 | # In a real implementation, you might want to extract the file hash and get the full report |
| 233 | return FileReportResponse(data=status_response.data, success=True, message="File analysis completed successfully") |
| 234 | |
| 235 | logger.info(f"Analysis status: {status_response.data.attributes.status}. Waiting {poll_interval} seconds...") |
| 236 | await asyncio.sleep(poll_interval) |
| 237 | waited_time += poll_interval |
| 238 | |
| 239 | # If we get here, the analysis timed out |
| 240 | raise HTTPException( |
| 241 | status_code=408, |
| 242 | detail=f"Analysis did not complete within {max_wait_time} seconds. You can check the status later using analysis ID: {analysis_id}", |
| 243 | ) |