lollms-webui/tests/pentests/code_injection/test_code_injection.py

from fastapi.testclient import TestClient
from main import app  # Replace with the actual name of your FastAPI app
import json

client = TestClient(app)

def test_open_code_in_vs_code_valid():
    response = client.post(
        "/open_code_in_vs_code",
        data=json.dumps({
            "discussion_id": 1,
            "message_id": 1,
            "code": "print('Hello, World!')"
        }),
        headers={"content-type": "application/json"},
    )
    assert response.status_code == 200
    assert response.json()["status"] == True

def test_open_code_in_vs_code_invalid():
    response = client.post(
        "/open_code_in_vs_code",
        data=json.dumps({
            "discussion_id": "1; copy file.exe /some/path/",
            "message_id": "1",
            "code": "print('Hello, World!')"
        }),
        headers={"content-type": "application/json"},
    )
    assert response.status_code == 422  # Unprocessable Entity

def test_open_code_in_vs_code_attack():
    response = client.post(
        "/open_code_in_vs_code",
        data=json.dumps({
            "discussion_id": 1,
            "message_id": 1,
            "code": "print('This is a harmless test.')"  # Dangerous code
        }),
        headers={"content-type": "application/json"},
    )
    assert response.status_code == 200
    assert response.json()["status"] == False  # The code should not be executed