Added tool guidelines generation + nmap guidelines

scripts/gen_tool_guidelines.py

@@ -21,12 +21,11 @@
 import json
 import argparse
 import textwrap
-from textwrap import dedent
+from pprint import pprint
 
 from tqdm import tqdm
 import google.generativeai as genai
 from google.generativeai.types import HarmCategory, HarmBlockThreshold
-from dotenv import load_dotenv
 
 
 PROMPTS = {
@@ -104,6 +103,7 @@ def parse_arguments(argv: list):
 
     parser.add_argument(
         "--output-path",
+        default='./tool_guidelines.json',
         help="Specifies the path for tool guidelines"
     )
 
@@ -113,7 +113,8 @@ def parse_arguments(argv: list):
     return {
         'tool_name': arguments.tool_name,
         'docs_path': arguments.docs_path,
-        'api_key': arguments.api_key
+        'api_key': arguments.api_key,
+        'output_path': arguments.output_path
     }
 
 
@@ -127,7 +128,6 @@ def load_tool_docs(path: str, tool_name: str) -> str:
     """
     if not os.path.exists(path):
         raise RuntimeError(f"Path don't exists: {path}")
-    print('Loading tool documentation')
 
     out = f'# {tool_name}\n\n'
     with open(path, 'r', encoding='utf-8') as fp:
@@ -147,7 +147,6 @@ def load_tool_docs(path: str, tool_name: str) -> str:
 
 def generate_summary(full_docs: str) -> str:
     """Generates the summary of the tool documentation"""
-    print('Generating summary')
     llm = genai.GenerativeModel('gemini-1.5-flash')
     prompt = PROMPTS['summarization'].format(tool_documentation=full_docs)
     response = llm.generate_content(prompt, safety_settings=SAFETY_SETTINGS)
@@ -157,7 +156,6 @@ def generate_summary(full_docs: str) -> str:
 def generate_scenarios(docs_summary: str,
                        num_scenarios: int = 3) -> list[str]:
     """Generates usage scenarios based on the summarized tool documentation"""
-    print('Generating Scenarios')
     llm = genai.GenerativeModel('gemini-1.5-flash')
     prompt = PROMPTS['scenario'].format(tool_documentation=docs_summary)
 
@@ -172,21 +170,28 @@ def generate_tool_documentation(path: str, tool_name) -> dict:
     """Generates the tool documentation dictionary"""
     tool_docs = load_tool_docs(path, tool_name)
 
-    tool_docs_summary = generate_summary(tool_docs)
-    scenarios = generate_scenarios(tool_docs_summary)
+    with tqdm(total=3, desc="Overall Progress") as pbar:
+        pbar.set_description("Generating summary")
+        tool_docs_summary = generate_summary(tool_docs)
+        pbar.update(1)
+
+        pbar.set_description("Generating scenarios")
+        scenarios = generate_scenarios(tool_docs_summary)
+        pbar.update(1)
+
+        pbar.set_description("Generating tool documentation")
+        llm = genai.GenerativeModel(
+            'gemini-1.5-flash',
+            generation_config={"response_mime_type": "application/json"}
+        )
+        prompt = PROMPTS['output'].format(
+            tool_documentation=tool_docs_summary,
+            tool_scenarios='\n'.join(scenarios)
+        )
 
-    llm = genai.GenerativeModel(
-        'gemini-1.5-flash',
-        generation_config={"response_mime_type": "application/json"}
-    )
-    prompt = PROMPTS['output'].format(
-        tool_documentation=tool_docs_summary,
-        tool_scenarios='\n'.join(scenarios)
-    )
-
-    print('Generating tool documentation')
-    response = llm.generate_content(prompt, safety_settings=SAFETY_SETTINGS)
-    tool_guidelines = json.loads(response.text)
+        response = llm.generate_content(prompt, safety_settings=SAFETY_SETTINGS)
+        tool_guidelines = json.loads(response.text)
+        pbar.update(1)
 
     return tool_guidelines
 
@@ -196,6 +201,8 @@ def generate_tool_documentation(path: str, tool_name) -> dict:
 
     genai.configure(api_key=args['api_key'])
     output = generate_tool_documentation(args['docs_path'], args['tool_name'])
-    print(output)
-
+    print('Result: \n')
+    pprint(output)
 
+    with open(args['output_path'], 'w', encoding='utf-8') as fp_out:
+        json.dump(output, fp_out,indent=4)

tools_settings/nmap.json

@@ -1,13 +1,41 @@
 {
-    "name": "nmap",
-    "tool_description": "Nmap is a tool for network exploration used during Reconnaissance in Penetration Testing.",
+    "name": "Nmap",
+    "tool_description": "Nmap is a versatile tool used for network discovery and security auditing.",
     "args_description": [
-        "The most important nmap command argument is the iPv4 address of the target machine.\n",
-        "Ensure that the IP address is in the format xxx.xxx.xxx.xxx with points separating the octets.\n",
-        "Some of the other arguments are as follows:\n",
-        "-p     : Specifies the port(s) to scan.\n",
-        "-sV    : Enables version detection on open ports\n",
-        "-sC    : Executes a set of default scripts that are deemed useful for most scanning situations.\n",
-        "-T<0-5>: Sets the timing template (higher is faster)\n"
+        "**-sL:** List Scan - Simply lists targets without sending any packets. ",
+        "**-sn:** Ping Scan - Disables port scan, only performs host discovery.",
+        "**-Pn:** Treats all hosts as online, skips host discovery. ",
+        "**-PS/PA/PU/PY:** TCP SYN/ACK, UDP, or SCTP discovery to given ports.",
+        "**-PE/PP/PM:** ICMP echo, timestamp, and netmask request discovery probes.",
+        "**-PO:** IP Protocol Ping.",
+        "**-sS:** TCP SYN Scan - Default and most popular scan.",
+        "**-sT:** TCP Connect Scan - Used when raw packet privileges are not available.",
+        "**-sU:** UDP Scan.",
+        "**-sY:** SCTP INIT Scan - Similar to TCP SYN scan.",
+        "**-sN/sF/sX:** TCP Null, FIN, and Xmas scans - Exploit loopholes in the TCP RFC.",
+        "**-sA:** TCP ACK Scan - Used to map firewall rulesets.",
+        "**-sW:** TCP Window Scan - Similar to ACK scan but differentiates open and closed ports.",
+        "**-sM:** TCP Maimon Scan - Similar to NULL/FIN/Xmas scans, but uses FIN/ACK probe.",
+        "**--scanflags:** Customize TCP scan flags.",
+        "**-sZ:** SCTP COOKIE ECHO Scan - Less obvious than INIT scan.",
+        "**-sI:** Idle Scan - Truly blind TCP port scan.",
+        "**-sO:** IP Protocol Scan - Determines supported IP protocols.",
+        "**-b:** FTP Bounce Scan - Uses FTP to scan other hosts.",
+        "**-sV:** Enables service and version detection.",
+        "**-O:** Enables OS detection.",
+        "**-sC:** Performs a script scan using the default set of scripts.",
+        "**--script:** Runs a script scan using a comma-separated list of filenames, categories, and directories.",
+        "**-oN:** Normal output.",
+        "**-oX:** XML output.",
+        "**-oS:** ScRiPt KiDDi3 0utPUt.",
+        "**-oG:** Grepable output (deprecated).",
+        "**-oA:** Output to all formats.",
+        "**-6:** Enable IPv6 scanning.",
+        "**-A:** Enables OS detection, version detection, script scanning, and traceroute.",
+        "**--datadir:** Specify custom Nmap data file location.",
+        "**--send-eth/--send-ip:** Send using raw ethernet frames or IP packets.",
+        "**--privileged/--unprivileged:** Assume the user is fully privileged or lacks raw socket privileges.",
+        "**-V:** Print version number.",
+        "**-h:** Print this help summary page."
     ]
 }