diff --git a/test/tests/__pycache__/__init__.cpython-311.pyc b/test/tests/__pycache__/__init__.cpython-311.pyc
new file mode 100644
index 0000000..da5a780
Binary files /dev/null and b/test/tests/__pycache__/__init__.cpython-311.pyc differ
diff --git a/test/tests/__pycache__/test_plan.cpython-311-pytest-8.2.2.pyc b/test/tests/__pycache__/test_plan.cpython-311-pytest-8.2.2.pyc
new file mode 100644
index 0000000..797df85
Binary files /dev/null and b/test/tests/__pycache__/test_plan.cpython-311-pytest-8.2.2.pyc differ
diff --git a/test/tests/plan_tests.json b/test/tests/plan_tests.json
index e69de29..5c7b460 100644
--- a/test/tests/plan_tests.json
+++ b/test/tests/plan_tests.json
@@ -0,0 +1,62 @@
+[
+ {
+ "content": [
+ "I completely agree with you! Let's refine our plan and use the tools more efficiently.\n\n",
+ "**Step 1: DNS Information**\n",
+ "Using `dig`, we'll gather DNS information about the target domain:\n",
+ "terminal dig example.com\n",
+ "This will provide us with valuable information, such as MX records, NS records, and potential subdomains.\n\n",
+ "**Step 2: Port Scan (nmap -sV)**\n",
+ "We can skip scanning specific ports since `nmap -sV` already scans the common ports (80, 21, 22, and 443). This will give us a comprehensive view of open ports and services:\n\n",
+ "nmap -sV example.com\n\n",
+ "**Step 3: NSE Scripting**\n",
+ "To gather more information, we can use nmap's NSE (Nmap Scripting Engine) to run scripts that scan for vulnerabilities, check for common web servers, and identify potential issues. For example:\n\n",
+ "nmap --script=http-title,http-open-proxy,http-methods,http-robots.txt example.com\n",
+ "This will help us identify potential vulnerabilities, open proxies, and other information.\n\n",
+ "**Step 4: Operating System Identification**\n",
+ "Using `nmap -O`, we can determine the operating system running on the target machine:\n",
+ "terminal nmap -O example.com\n",
+ "This will help us narrow down potential attack vectors or identify specific vulnerabilities.\n",
+ "By refining our plan, we can gather more information about the target host and services without attempting to exploit any vulnerabilities. What do you think?\n"
+ ],
+ "commands": [
+ "dig example.com",
+ "nmap -sV example.com",
+ "nmap --script=http-title,http-open-proxy,http-methods,http-robots.txt example.com",
+ "nmap -O example.com"
+ ]
+ },
+ {
+ "content": [
+ "**Plan for Reconnaissance Phase:**\n\n",
+ "**Target System:** 154.142.1.1 (IP Camera).\n",
+ "**Objective:** Gather information about the operating system, services, data transmission, and potential vulnerabilities.\n",
+ "**Tasks:**\n\n",
+ "1. **Nmap Scan**: Use Nmap to gather information about the IP Camera's operating system, services, and open ports.\n",
+ " - Execute the following command: `nmap -O -sV 154.142.1.1`\n",
+ "2. **Operating System Detection**: Identify the operating system running on the IP Camera using Nmap's `-O` option.\n",
+ "3. **Service Detection**: Identify any services or protocols being used by the IP Camera, such as HTTP, FTP, or SSH, using Nmap's `-sV` option.\n",
+ "4. **Data Transmission Analysis**: Use Nmap's `-A` option to analyze data transmission patterns and identify any potential issues.\n",
+ "5. **SSH and FTP Vulnerability Scan**: Run a vulnerability scan for SSH and FTP protocols using Nmap's `--script` option with the following commands:\n",
+ " - `nmap --script sshvuln 154.142.1.1`\n",
+ " - `nmap --script ftp-vuln 154.142.1.1`\n\n",
+ "Please let me know if you'd like to proceed with this plan or if you have any specific requirements or concerns.\n"
+ ],
+ "commands": [
+ "nmap -O -sV 154.142.1.1",
+ "nmap --script sshvuln 154.142.1.1",
+ "nmap --script ftp-vuln 154.142.1.1"
+ ]
+ },
+ {
+ "content": [
+ "To do a port scan and version scan on 192.168.1.252, I'll use the following command:\n\n",
+ "`nmap -p 1-1024 -sV 192.168.1.252`\n\n",
+ "This command will perform a quick port scan (nmap -p 1-1024) to identify all open ports from 1 to 1024 on the target IP address, and then do a version scan (nmap -sV) to detect the operating system and any services running on those open ports.\n\n",
+ "Would you like me to run this command for you?\n"
+ ],
+ "commands": [
+ "nmap -p 1-1024 -sV 192.168.1.252"
+ ]
+ }
+]
\ No newline at end of file
diff --git a/test/tests/test_plan.py b/test/tests/test_plan.py
index bd79e6b..80aae26 100644
--- a/test/tests/test_plan.py
+++ b/test/tests/test_plan.py
@@ -1,4 +1,4 @@
-import textwrap
+import json
import unittest
from src.agent import Agent
@@ -8,82 +8,46 @@
class TestPlan(unittest.TestCase):
- NL_PLANS = [
- textwrap.dedent("""
-I completely agree with you! Let's refine our plan and use the tools more efficiently.
+ # def test_execute(self):
+ # tasks = [
+ # Task(thought="Get directory content", tool=Terminal, command="ls"),
+ # Task(thought="Get machine host name", tool=Terminal, command="hostname")
+ # ]
-**Step 1: DNS Information**
+ # plan = Plan(tasks)
+ # for output in plan.execute():
+ # print('---------------------------------')
+ # for i, task_overview in enumerate(output):
+ # print(f'{i+1}. {task_overview}')
+ # if task_overview.status == TaskStatus.DONE:
+ # print(f'Output:\n{task_overview.output}')
-Using `dig`, we'll gather DNS information about the target domain:
-
-terminal dig example.com
-
-This will provide us with valuable information, such as MX records, NS records, and potential subdomains.
-
-**Step 2: Port Scan (nmap -sV)**
-
-We can skip scanning specific ports since `nmap -sV` already scans the common ports (80, 21, 22, and 443). This will give us a comprehensive view of open ports and services:
-
-nmap -sV example.com
-
-**Step 3: NSE Scripting**
-
-To gather more information, we can use nmap's NSE (Nmap Scripting Engine) to run scripts that scan for vulnerabilities, check for common web servers, and identify potential issues. For example:
-
-nmap --script=http-title,http-open-proxy,http-methods,http-robots.txt example.com
-
-This will help us identify potential vulnerabilities, open proxies, and other information.
-
-**Step 4: Operating System Identification**
-
-Using `nmap -O`, we can determine the operating system running on the target machine:
-
-terminal nmap -O example.com
-
-This will help us narrow down potential attack vectors or identify specific vulnerabilities.
-
-By refining our plan, we can gather more information about the target host and services without attempting to exploit any vulnerabilities. What do you think?
- """),
- textwrap.dedent("""
- **Plan for Reconnaissance Phase:**
+ def test_from_response(self):
+ agent = Agent(model='llama3')
+ with open('plan_tests.json', 'r', encoding='utf-8') as fp:
+ test_cases = json.load(fp)
-**Target System:** 154.142.1.1 (IP Camera).
-**Objective:** Gather information about the operating system, services, data transmission, and potential vulnerabilities.
-**Tasks:**
+ for test_case in test_cases:
+ plan_nl = test_case['content']
+ expected_commands = test_case['commands']
-1. **Nmap Scan**: Use Nmap to gather information about the IP Camera's operating system, services, and open ports.
- - Execute the following command: `nmap -O -sV 154.142.1.1`
-2. **Operating System Detection**: Identify the operating system running on the IP Camera using Nmap's `-O` option.
-3. **Service Detection**: Identify any services or protocols being used by the IP Camera, such as HTTP, FTP, or SSH, using Nmap's `-sV` option.
-4. **Data Transmission Analysis**: Use Nmap's `-A` option to analyze data transmission patterns and identify any potential issues.
-5. **SSH and FTP Vulnerability Scan**: Run a vulnerability scan for SSH and FTP protocols using Nmap's `--script` option with the following commands:
- - `nmap --script sshvuln 154.142.1.1`
- - `nmap --script ftp-vuln 154.142.1.1`
+ plan = agent.extract_plan(plan_nl)
+ self.assertIsNotNone(plan, "Plan is None:")
-Please let me know if you'd like to proceed with this plan or if you have any specific requirements or concerns.
- """),
- ]
+ commands = [task.command for task in plan.tasks]
- def test_execute(self):
- tasks = [
- Task(thought="Get directory content", tool=Terminal, command="ls"),
- Task(thought="Get machine host name", tool=Terminal, command="hostname")
- ]
+ self.assertEquals(
+ len(commands),
+ len(expected_commands),
+ f"commands {len(commands)} != expected {len(expected_commands)}"
+ )
+ self.assertEquals(
+ commands,
+ expected_commands,
+ f"Commands:\n{commands}\nExpected:\n{expected_commands}"
+ )
- plan = Plan(tasks)
- for output in plan.execute():
- print('---------------------------------')
- for i, task_overview in enumerate(output):
- print(f'{i+1}. {task_overview}')
- if task_overview.status == TaskStatus.DONE:
- print(f'Output:\n{task_overview.output}')
- def test_from_response(self):
- agent = Agent(model='llama3')
- for plan_nl in TestPlan.NL_PLANS:
- plan = agent.extract_plan(plan_nl)
- print(plan)
- self.assertIsNotNone(plan)
# def test_should_timeout(self):
# pass