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

import json

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

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