lollms-webui/tests/pentests/multipart_dos/test_no_dos.py

59 lines
2.2 KiB
Python

import requests
from requests_toolbelt import MultipartEncoder
import io
import zipfile
def test_valid_file_upload():
base_url = "http://localhost:9600/upload_app"
client_id = "3qxKnpFF8aJU8KsZAAAH" # Replace with a valid client id
url = f"{base_url}?client_id={client_id}" # Add client_id as a query parameter
# Create a test zip file in memory
zip_buffer = io.BytesIO()
with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zipf:
zipf.writestr('index.html', '<html><body>Test</body></html>')
zipf.writestr('description.yaml', 'name: TestApp\n')
zipf.writestr('icon.png', b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x06\x00\x00\x00\x1f\x15\xc4\x89\x00\x00\x00\nIDATx\x9cc\x00\x01\x00\x00\x05\x00\x01\r\n-\xb4\x00\x00\x00\x00IEND\xaeB`\x82')
zip_buffer.seek(0)
try:
# Prepare the multipart/form-data request
m = MultipartEncoder(
fields={'file': ('test_upload.zip', zip_buffer, 'application/zip')}
)
headers = {
'Content-Type': m.content_type
}
# Send the POST request
response = requests.post(url, data=m, headers=headers)
# Check the response
print(f"Status Code: {response.status_code}")
print(f"Response Content: {response.text}")
# If it's a 422 error, try to parse and print the JSON error message
if response.status_code == 422:
try:
error_detail = response.json()
print("Error details:")
print(error_detail)
except:
print("Could not parse error details as JSON")
# Assert the expected behavior
assert response.status_code == 200, f"Expected status code 200, but got {response.status_code}"
assert "App 'TestApp' uploaded successfully" in response.text, "File upload confirmation message not found in response"
print("Test passed successfully!")
except requests.RequestException as e:
print(f"Request failed: {e}")
except AssertionError as e:
print(f"Test failed: {e}")
if __name__ == "__main__":
test_valid_file_upload()