From 8f5bc0c2188529444c7870dee578bae9be2bfa3a Mon Sep 17 00:00:00 2001 From: Antonino Lorenzo <94693967+antoninoLorenzo@users.noreply.github.com> Date: Wed, 17 Jul 2024 12:16:46 +0200 Subject: [PATCH] Added tool guidelines generation + nmap guidelines --- scripts/gen_tool_guidelines.py | 51 +++++++++++++++++++--------------- tools_settings/nmap.json | 46 ++++++++++++++++++++++++------ 2 files changed, 66 insertions(+), 31 deletions(-) diff --git a/scripts/gen_tool_guidelines.py b/scripts/gen_tool_guidelines.py index 104c127..1c30817 100644 --- a/scripts/gen_tool_guidelines.py +++ b/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) diff --git a/tools_settings/nmap.json b/tools_settings/nmap.json index cbca05d..0a0f9ed 100644 --- a/tools_settings/nmap.json +++ b/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." ] } \ No newline at end of file