mirror of https://github.com/aliasrobotics/cai.git
71 lines
2.8 KiB
Python
71 lines
2.8 KiB
Python
import os
|
|
import json
|
|
import re
|
|
import numpy as np
|
|
from datasets import Dataset
|
|
|
|
|
|
def load_step_data(step_folder):
|
|
all_samples = []
|
|
for filename in os.listdir(step_folder):
|
|
if filename.endswith(".json"):
|
|
filepath = os.path.join(step_folder, filename)
|
|
with open(filepath, "r") as f:
|
|
data = json.load(f)
|
|
|
|
initial_prompt = data.get("initial_prompt", "")
|
|
SYSTEM_PROMPT = data.get("system_prompt", "") # MODIFIED
|
|
|
|
for step in data.get("steps", []):
|
|
if step["step_number"] > 1:
|
|
user_content = f"Initial Description of Penetration Testing Task: {initial_prompt}"
|
|
for prev_step in data.get("steps", [])[
|
|
: step["step_number"] - 1
|
|
]:
|
|
user_content += f"\n===Step {prev_step['step_number']}==="
|
|
user_content += f"\nThought: {prev_step['thought']}"
|
|
user_content += f"\nCommand: {prev_step['command']}"
|
|
user_content += f"\nObservation: {prev_step['result']}"
|
|
user_content += "\nWhat are the next steps and actions?"
|
|
else:
|
|
user_content = f"Initial Description of Penetration Testing Task: {initial_prompt}\nWhat are the next instruction and commands?"
|
|
|
|
content = f"""You will guide the penetration testing process step by step. At each stage, you will receive the context from the previous steps. Based on that information, you must determine the next logical action to advance the assessment. Continue this process until the penetration test is fully completed.
|
|
|
|
### Response Format (Always follow this format exactly):
|
|
=== Step i ===
|
|
Thought: [Explain your reasoning and goals here.]
|
|
Command: [Write the full command you would run.]
|
|
|
|
### Previous Steps Context
|
|
{user_content}"""
|
|
|
|
assistant_content = f"===Step {step['step_number']}===\nThought: {step['thought']}\nCommand: {step['command']}"
|
|
|
|
sample = {
|
|
"prompt": [
|
|
{"role": "system", "content": SYSTEM_PROMPT},
|
|
{"role": "user", "content": content},
|
|
],
|
|
"answer": assistant_content,
|
|
}
|
|
all_samples.append(sample)
|
|
|
|
return Dataset.from_list(all_samples)
|
|
|
|
|
|
def main():
|
|
dataset = load_step_data("data/walkthroughs")
|
|
print(dataset)
|
|
print(dataset.features)
|
|
print(dataset.column_names)
|
|
|
|
answers = dataset["answer"]
|
|
prompts = dataset["prompt"]
|
|
print(
|
|
answers[:2]
|
|
) # increasing the index is going to increase the the number of steps visualized
|
|
|
|
|
|
main()
|