# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.
import logging
import mimetypes
import os
from typing import Any, Callable, Literal, Optional
import httpx
from pyrit.models import (
PromptRequestPiece,
PromptRequestResponse,
construct_response_from_request,
)
from pyrit.prompt_target import HTTPTarget
from pyrit.prompt_target.common.utils import limit_requests_per_minute
logger = logging.getLogger(__name__)
[docs]
class HTTPXAPITarget(HTTPTarget):
"""
A subclass of HTTPTarget that *only* does "API mode" (no raw HTTP request).
This is a simpler approach for uploading files or sending JSON/form data.
Additionally, if 'file_path' is not provided in the constructor,
we attempt to pull it from the prompt's `converted_value`, assuming
it's a local file path generated by a PromptConverter (like PDFConverter).
"""
[docs]
def __init__(
self,
*,
http_url: str,
method: Literal["GET", "POST", "PUT", "DELETE", "PATCH", "HEAD", "OPTIONS"] = "POST",
file_path: Optional[str] = None,
json_data: Optional[dict] = None,
form_data: Optional[dict] = None,
params: Optional[dict] = None,
headers: Optional[dict] = None,
http2: Optional[bool] = None,
callback_function: Callable | None = None,
max_requests_per_minute: Optional[int] = None,
**httpx_client_kwargs: Any,
) -> None:
"""
Force the parent 'HTTPTarget' to skip raw http_request logic by setting http_request=None.
"""
super().__init__(
http_request="",
prompt_regex_string="",
use_tls=True,
callback_function=callback_function,
max_requests_per_minute=max_requests_per_minute,
**httpx_client_kwargs,
)
self.http_url = http_url
self.method = method
self.file_path = file_path
self.json_data = json_data
self.form_data = form_data
self.params = params
self.headers = headers or {}
self.http2 = http2
# Validate method
if self.method not in {"GET", "POST", "PUT", "DELETE", "PATCH", "HEAD", "OPTIONS"}:
raise ValueError(f"Invalid HTTP method: {self.method}")
# Validate file uploads (only `POST` and `PUT` allow file uploads)
if self.file_path and self.method not in {"POST", "PUT"}:
raise ValueError(f"File uploads are not allowed with HTTP method: {self.method}")
@limit_requests_per_minute
async def send_prompt_async(self, *, prompt_request: PromptRequestResponse) -> PromptRequestResponse:
"""
Override the parent's method to skip raw http_request usage,
and do a standard "API mode" approach.
- If file_path is set or we can deduce it from the prompt piece, we upload a file.
- Otherwise, we send normal requests with JSON or form_data (if provided).
"""
self._validate_request(prompt_request=prompt_request)
request_piece: PromptRequestPiece = prompt_request.request_pieces[0]
# If user didn't set file_path, see if the PDF path is in converted_value
if not self.file_path:
possible_path = request_piece.converted_value
if isinstance(possible_path, str) and os.path.exists(possible_path):
logger.info(f"HTTPXApiTarget: auto-using file_path from {possible_path}")
self.file_path = possible_path
if not self.http_url:
raise ValueError("No `http_url` provided for HTTPXApiTarget.")
http2_version = self.http2 if self.http2 is not None else False
async with httpx.AsyncClient(http2=http2_version, **self.httpx_client_kwargs) as client:
try:
if self.file_path and os.path.exists(self.file_path):
# Handle file upload (only for POST & PUT)
filename = os.path.basename(self.file_path)
mime_type = mimetypes.guess_type(filename)[0] or "application/octet-stream"
with open(self.file_path, "rb") as fp:
file_bytes = fp.read()
files = {"file": (filename, file_bytes, mime_type)}
logger.info(f"HTTPXApiTarget: uploading file={filename} via {self.method} to {self.http_url}")
response = await client.request(
method=self.method,
url=self.http_url,
headers=self.headers,
files=files,
follow_redirects=True,
)
else:
# No file upload, handle based on HTTP method
logger.info(f"HTTPXApiTarget: sending {self.method} to {self.http_url} with possible JSON/form.")
response = await client.request(
method=self.method,
url=self.http_url,
headers=self.headers,
params=self.params if self.method in {"GET", "HEAD"} else None,
json=self.json_data if self.method in {"POST", "PUT", "PATCH"} else None,
data=self.form_data if self.method in {"POST", "PUT", "PATCH"} else None,
follow_redirects=True,
)
except httpx.TimeoutException:
logger.error(f"Timeout error for URL: {self.http_url}")
raise
except httpx.RequestError as e:
logger.error(f"Request failed: {e}")
raise
except FileNotFoundError as e:
logger.error(f"File not found: {self.file_path}. Exception: {e}")
raise
response_content = response.content
# If a callback function was set, let them parse the response
if self.callback_function:
response_content = self.callback_function(response=response)
# Reuse parent's response object construction
response_entry = construct_response_from_request(
request=request_piece, response_text_pieces=[str(response_content)]
)
return response_entry
