111 lines
4.2 KiB
Python
111 lines
4.2 KiB
Python
import concurrent.futures
|
|
import io
|
|
import random
|
|
import string
|
|
|
|
import requests
|
|
|
|
# Endpoint
|
|
endpoint_url = "http://IP:PORT/api/files/" # Change IP and PORT as needed
|
|
|
|
# Number of files to send
|
|
num_files = 3
|
|
num_threads = 5 # Change the number of threads as needed
|
|
|
|
# Test file path
|
|
test_file_path = "/path/to/virus.exe"
|
|
|
|
|
|
# Function to generate random file content
|
|
def generate_random_content():
|
|
content_length = random.randint(1024, 20480) # Random content size between 1KB and 20KB
|
|
return ''.join(random.choices(string.ascii_letters + string.digits, k=content_length))
|
|
|
|
|
|
# Function to send files in a thread and return the key
|
|
def send_files_thread(thread_id, use_test_file=False):
|
|
keys = []
|
|
for i in range(num_files):
|
|
if use_test_file and i == 0:
|
|
# Send the test file
|
|
with open(test_file_path, 'rb') as test_file:
|
|
test_file_content = io.BytesIO(test_file.read())
|
|
files = {'file': ('shell_reverse.exe', test_file_content)}
|
|
else:
|
|
# Generate random file content
|
|
content = generate_random_content()
|
|
file_content = io.BytesIO(content.encode())
|
|
files = {'file': ('file{}.txt'.format(i), file_content)}
|
|
|
|
# Define any additional data you want to send with the request
|
|
data = {'ttl': 60}
|
|
|
|
try:
|
|
# Make the POST request
|
|
response = requests.post(endpoint_url, files=files, data=data)
|
|
|
|
# Check if the request was successful
|
|
if response.status_code == 201:
|
|
items = response.json()
|
|
for item in items:
|
|
key = item.get('key')
|
|
if key:
|
|
keys.append(key)
|
|
print(f"Thread {thread_id}: File {i} uploaded successfully with key {key}.")
|
|
else:
|
|
print(f"Thread {thread_id}: File {i} uploaded, but key not found in the response.")
|
|
else:
|
|
print(f"Thread {thread_id}: File {i} upload failed with status code {response.status_code}.")
|
|
except Exception as e:
|
|
print(f"Thread {thread_id}: An error occurred while sending file {i}: {str(e)}")
|
|
finally:
|
|
# Close the file content
|
|
if use_test_file:
|
|
test_file_content.close()
|
|
else:
|
|
file_content.close()
|
|
|
|
return keys
|
|
|
|
|
|
# Function to receive files in a thread using the provided key
|
|
def receive_files_thread(thread_id, file_key):
|
|
try:
|
|
# Make the GET request to retrieve the file using the provided key
|
|
response = requests.get(endpoint_url + file_key)
|
|
|
|
# Check if the request was successful
|
|
if response.status_code == 200:
|
|
# Save or process the received file content here
|
|
print(f"Thread {thread_id}: File with key {file_key} received successfully.")
|
|
else:
|
|
print(
|
|
f"Thread {thread_id}: File with key {file_key} retrieval failed with status code {response.status_code}.")
|
|
except Exception as e:
|
|
print(f"Thread {thread_id}: An error occurred while receiving the file with key {file_key}: {str(e)}")
|
|
|
|
|
|
# Create and send/receive files using multiple threads
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=num_threads) as executor:
|
|
# Select a random subset of threads to use the test file
|
|
use_test_file_threads = random.sample(range(num_threads), min(num_threads, num_files))
|
|
|
|
# Start multiple threads for sending files
|
|
send_futures = []
|
|
for thread_id in range(num_threads):
|
|
use_test_file = thread_id in use_test_file_threads
|
|
future = executor.submit(send_files_thread, thread_id, use_test_file)
|
|
send_futures.append(future)
|
|
|
|
# Retrieve the keys from the responses
|
|
keys = []
|
|
for future in concurrent.futures.as_completed(send_futures):
|
|
keys.extend(future.result())
|
|
|
|
# Start multiple threads for receiving files using the obtained keys
|
|
receive_futures = [executor.submit(receive_files_thread, thread_id, file_key) for thread_id, file_key in
|
|
enumerate(keys) if file_key]
|
|
|
|
# Wait for all receive threads to complete
|
|
concurrent.futures.wait(receive_futures)
|