VulnMapAI combines the power of nmap’s detailed network scanning and the advanced natural language processing capabilities of GPT-4 to generate comprehensive and intelligible vulnerability reports. It aims to facilitate the identification and understanding of security vulnerabilities.
import asyncio
import json
import os
from yaml import safe_load
from typing import Dict
from config import PROMPTS_FILE
import httpx

from utils.logger import Logger


class Advisor:
    """
    Advisor class that prepares the reports by summarizing discoveries and providing attack advice,
    leveraging GPT-4 models from OpenAI API.
    """

    def __init__(self,
                 discovery_result,
                 title,
                 exploits,
                 api_key=os.getenv('OPENAI_API_KEY'),
                 api_url='https://api.openai.com/v1/chat/completions'):
        
        if not (api_key):
            self.logger.error("OPENAI_API_KEY NOT SET AS AN ENVIRONMENT VARIABLE")
            while not (api_key):
                api_key = input("PLEASE ENTER OPENAI API KEY: ")
        self.full_discovery_result = discovery_result
        self.discovery_result = " ".join(discovery_result.split()[:2000])
        self.title = title
        self.exploits = exploits
        self.api_key = api_key
        self.api_url = api_url
        self.headers = {
            'Authorization': f'Bearer {self.api_key}',
            'Content-Type': 'application/json',
        }
        self.logger = Logger(__name__).get_logger()

    def _query_api(self, json_data: Dict) -> str:
        """
        Private method to query the OpenAI API with provided JSON data.

        Args:
            json_data (dict): JSON data to be sent to the API.

        Returns:
            str: Response from the API.
        """

        async def _send_to_openai_async():
            async with httpx.AsyncClient() as client:
                response = await client.post(
                    url=self.api_url,
                    json=json_data,
                    headers=self.headers,
                    timeout=120,
                )
                if response.status_code == 200:
                    return response.json()["choices"][0]["message"]['content']
                else:
                    response.raise_for_status()

        return asyncio.run(_send_to_openai_async())

    def _load_prompts(self, prompt_type: str) -> Dict[str, str]:
        """
        Load prompts from the prompts.yaml file based on the prompt_type.

        Args:
            prompt_type (str): 'advise_attack_prompt', 'summarize_discoveries_prompt', etc.

        Returns:
            Dict[str, str]: Dictionary containing 'system' and 'user' prompts.
        """
        with open(PROMPTS_FILE, 'r') as prompts_file:
            prompts_data = safe_load(prompts_file)
        return prompts_data.get(prompt_type, {})

    def advise_attack(self) -> str:
        """
        Formulates a query and sends it to the OpenAI API to receive advice on potential attacks.

        Returns:
            str: Advice on potential attacks from the API.
        """
        self.logger.info("Advising attack based on the discovery result.")
        prompts = self._load_prompts('advise_attack_prompt')
        json_data = {
            'model': 'gpt-4',
            'messages': [
                {
                    'role': 'system',
                    'content': prompts['system'],
                },
                {
                    'role': 'user',
                    'content': prompts['user'].format(discovery_result=self.discovery_result),
                },
            ],
        }
        return self._query_api(json_data)

    def summarize_discoveries(self) -> str:
        """
        Formulates a  query and sends it to the OpenAI API to summarize the discoveries made during the scan.

        Returns:
            str: Summarized discoveries from the API.
        """
        self.logger.info("Summarizing the discoveries.")
        prompts = self._load_prompts('summarize_discoveries_prompt')
        json_data = {
            'model': 'gpt-4',
            'messages': [
                {
                    'role': 'system',
                    'content': prompts['system'],
                },
                {
                    'role': 'user',
                    'content': prompts['user'].format(discovery_result=self.discovery_result),
                },
            ],
        }
        return self._query_api(json_data)

    def prepare_report(self) -> Dict:
        """
        Prepares an HTML report based on discoveries, attack advice, and exploits.

        Returns:
            str: A string representation of the HTML report.
        """
        discoveries = self.summarize_discoveries()
        advise = self.advise_attack()

        self.logger.info("Preparing the HTML report.")
        exploits = '<br>'.join(self.exploits)

        self.full_discovery_result = json.loads(self.full_discovery_result)
        if "script" in self.full_discovery_result:
            for k, v in self.full_discovery_result["script"].items():
                self.full_discovery_result[k] = v
            del self.full_discovery_result["script"]

        return {
            "title": self.title,
            "exploits": exploits,
            "discoveries": discoveries,
            "advise": advise,
            "full_discovery_result": self.full_discovery_result
        }

advising/llm.py
davidalami/VulnMapAI
# configuration file for storing constants
from pathlib import Path


EXPLOITS_DIR = "/usr/share/metasploit-framework/modules/exploits"
MSF_INDEX_FILE = '/app/resources/msf_index.json'
COMMON_PORTS = "/app/resources/ports.json"
REPORT_DIR = Path('./report')
REPORT_TEMPLATE = '/app/webapp/templates/report_template.html'
PROMPTS_FILE = "/app/advising/prompts.yaml"

config.py
davidalami/VulnMapAI
import json
import pty
import re
import socket
import subprocess
import time
from typing import List, Dict

from pymetasploit3.msfrpc import MsfRpcClient

from config import MSF_INDEX_FILE
from utils.logger import Logger


class ExploitManager:
    """
    Class for testing and executing exploits using Metasploit Framework.
    """

    def __init__(self, host='127.0.0.1', password='msf', port=55552):
        self.logger = Logger(__name__).get_logger()
        self.client = self._start_msf(host, password, port)
        with open(MSF_INDEX_FILE, 'r') as f:
            self.msf_index = json.load(f)
        self.max_try_payloads = 3

    def _start_msf(self, host, password, port) -> MsfRpcClient:
        """
        Starts the Metasploit Framework RPC server.

        Args:
            host (str): The hostname or IP address where the MSF RPC server will bind to.
            password (str): The password to connect to the MSF RPC server.
            port (int): The port number where the MSF RPC server will bind to.

        Returns:
            MsfRpcClient: An instance of the MsfRpcClient, which is connected to the started MSF RPC server.

        Raises:
            TimeoutError: If failed to start MSF and establish a connection within the timeout period.
        """
        self.logger.info("Starting MSF...")
        master, slave = pty.openpty()
        subprocess.Popen(
            ['msfconsole', '-qx', f'load msgrpc ServerHost={host} Pass={password}'],
            stdin=slave, stdout=slave, stderr=slave, universal_newlines=True
        )
        if self._wait_msf_ready(host, port):
            self.logger.info("MSF successfully started.")
            return MsfRpcClient(password, port=port)
        else:
            self.logger.error("MSF failed to start!")
            raise TimeoutError("MSF failed to start!")

    def _wait_msf_ready(self, host: str, port: int, timeout=30) -> bool:
        """
        Waits until the MSF RPC server is ready.

        Args:
            host (str): The host address of the MSF RPC server.
            port (int): Port number of MSF RPC server.
            timeout (int): Maximum time to wait for MSF RPC server to be ready.

        Returns:
            bool: True if the MSF RPC server is ready, False if timed out.
        """
        start_time = time.time()
        while time.time() - start_time < timeout:
            try:
                with socket.create_connection((host, port), timeout=2):
                    self.logger.info(f"MSF started in {time.time() - start_time} seconds")
                    return True
            except (socket.timeout, ConnectionRefusedError):
                time.sleep(1)  # Wait for a second before trying again
        self.logger.error("Timed out waiting for MSF service to start!")
        return False

    def find_exploit_modules(self, service: str, version: str) -> List[str]:
        """
        Searches and returns exploit modules that match the provided service and version.

        Args:
            service (str): The name of the service.
            version (str): The version of the service.

        Returns:
            list: List of matching exploit module paths.
        """
        self.logger.info(f"Searching for exploits of {service} with version {version}")
        service_pattern = re.compile(re.escape(service), re.IGNORECASE)
        version_pattern = re.compile(re.escape(str(version)), re.IGNORECASE)

        matching_modules = set()
        for module in self.msf_index:
            if service_pattern.search(module['name']) or service_pattern.search(module['description']):
                if version_pattern.search(module['name']) or version_pattern.search(module['description']):
                    matching_modules.add(module['path'])

        return list(matching_modules)

    def apply_exploits(self, exploits: List, target: Dict) -> List[str]:
        """
        Applies a list of exploit modules against a target.

        Args:
            exploits (list): List of exploit module names to be applied.
            target (dict): Target information, including host and port.

        Returns:
            list: List of successful exploit names.
        """
        success_exploits = []
        for exploit_name in exploits:
            self.logger.info(f"Trying exploit {exploit_name}")
            exploit = self.client.modules.use('exploit', exploit_name)
            if 'RHOSTS' in exploit.options:
                exploit['RHOSTS'] = target['host']
            elif 'RHOST' in exploit.options:
                exploit['RHOST'] = target['host']
            else:
                # for simplicity, try only exploits with defined host/port
                continue

            if 'RPORT' in exploit.options:
                exploit['RPORT'] = target['port']

            for payload in self.client.call('module.compatible_payloads', [exploit_name])['payloads'][
                           :self.max_try_payloads]:
                self.logger.info(f"Trying payload {payload}")
                exploit.execute(payload=payload)
                if self._wait_for_success():
                    success_exploits.append(exploit_name)
                    break
        return success_exploits

    def _wait_for_success(self, retries: int = 3, delay: int = 5) -> bool:
        """
        Waits for an exploit to succeed.

        Args:
            retries (int): Number of times to check for successful exploitation.
            delay (int): Delay between each check in seconds.

        Returns:
            bool: True if successful, False otherwise.
        """
        for _ in range(retries):
            time.sleep(delay)
            if self.client.sessions.list:
                return True
        return False

    def exploit_target(self, target: Dict, product: str, version: str) -> List[str]:
        """
        Finds suitable exploit modules and applies them to a target.

        Args:
            target (dict): Target information, including host and port.
            product (str): Name of the product.
            version (str): Version of the product.

        Returns:
            list: List of response strings for successful exploits.
        """
        exploits = self.find_exploit_modules(product, version)
        successful_exploits = self.apply_exploits(exploits, target)
        return self._prepare_responses(target, successful_exploits)

    @staticmethod
    def _prepare_responses(target: Dict, successful_exploits: List[str]) -> List[str]:
        """
        Prepares response strings for successful exploits.

        Args:
            target (dict): Target information, including host and port.
            successful_exploits (list): List of successful exploit names.

        Returns:
            list: List of response strings for successful exploits.
        """
        responses = []
        for exploit in successful_exploits:
            msf_one_liner = (
                f'msfconsole -q -x "use {exploit}; set RHOSTS {target["host"]}; '
                f'set RPORT {target["port"]}; run"'
            )
            response = f'Exploit {exploit} succeeded against {target["host"]}:{target["port"]}. ' \
                       f'MSF one-liner: {msf_one_liner}'
            responses.append(response)
        return responses

exploit/metasploit.py
davidalami/VulnMapAI
import argparse

from orchestration.pentest import PenetrationTester
from webapp.app import app


def main(targets, top_ports):
    pentester = PenetrationTester(targets, top_ports)
    pentester.scan_and_report()
    app.run(debug=False, host='0.0.0.0', port=1337)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Run penetration tests on the given target subnets.')

    # Targets Argument
    parser.add_argument('targets', nargs='+', help='List of target subnets.')

    # Top Ports Argument
    parser.add_argument('-p', '--top_ports', type=int, default=300, help='Number of most common ports to scan.')

    # Parse the arguments
    args = parser.parse_args()

    # Validate top_ports
    if args.top_ports < 1:
        parser.error("top_ports must be a positive integer.")

    # Maximum limit of top_ports
    if args.top_ports > 8367:
        parser.error("top_ports must be less than or equal to 8367.")

    # Run the main function with parsed arguments
    main(args.targets, args.top_ports)

main.py
davidalami/VulnMapAI
import json

from advising.llm import Advisor
from config import REPORT_DIR
from exploit.metasploit import ExploitManager
from recon.async_scanner import NetworkScanner
from utils.logger import Logger


class PenetrationTester:
    """
    A class representing a Penetration Tester which combines the functionality of a NetworkScanner and an ExploitManager.

    Attributes:
    target_subnets (list): A list of subnets to be scanned.
    scanner (NetworkScanner): An instance of NetworkScanner to perform network scans on the target subnets.
    exploit_manager (ExploitManager): An instance of ExploitManager to find and manage exploits for discovered services.
    """
    def __init__(self, target_subnets, top_ports):
        self.target_subnets = target_subnets
        self.scanner = NetworkScanner(target_subnets, top_ports)
        self.exploit_manager = ExploitManager()
        self.logger = Logger(__name__).get_logger()

    def scan_and_report(self):
        """
        Performs scanning on target subnets, exploits the targets and prepares a report for each discovered service.
        """

        self.logger.info("Cleaning results from previous scan...")

        # Ensure the directory exists before proceeding
        if REPORT_DIR.exists() and REPORT_DIR.is_dir():
            # Iterate over all the HTML files in the directory
            for json_file in REPORT_DIR.glob('*.json'):
                # Delete the file
                json_file.unlink()

        self.logger.info(f"Old report files were removed from {str(REPORT_DIR.absolute())}")

        self.logger.info("Starting scan and report process.")
        for result in self.scanner.execute_scan():
            target = {'host': result['host'], 'port': result['port']}
            title = ' '.join([result['product'] or result['name'], result['version']]).strip()
            self.logger.info(f"Trying exploits for target: {result['host']}:{result['port']}, {title}")
            exploits = self.exploit_manager.exploit_target(target, result['product'] or result['name'],
                                                           result['version'])

            advisor = Advisor(
                discovery_result=json.dumps(result),
                title=title,
                exploits=exploits
            )

            report_file = REPORT_DIR / f"{title}.json"
            self.logger.info(f"Preparing report and saving to {report_file}")
            with report_file.open("w") as f:
                f.write(json.dumps(advisor.prepare_report()))

orchestration/pentest.py
davidalami/VulnMapAI
"""
This module provides functionality to build an index of exploits found in the specified directory.
It will search for Ruby files, parse the exploit's name and description, and save the data to a JSON index file.
"""

import os
import re
import json

from config import EXPLOITS_DIR, MSF_INDEX_FILE


def build_index(exploits_dir=EXPLOITS_DIR) -> None:
    """
    Parses and indexes the Name and Description from Ruby (.rb) exploit files located in the specified directory.
    Saves the indexed data as a JSON file (`MSF_INDEX_FILE`).
    The indexing will go through each Ruby file, extract the exploit's name and description, and add this information 
    to the index, along with the relative path to the exploit.
    
    Args:
        exploits_dir (str): The path to the directory containing the Ruby exploit files. Defaults to `EXPLOITS_DIR` from config.
    Returns:
        None
    """
    index = []
    name_pattern = re.compile(r"'Name'\s*=>\s*'([^']*)'", re.IGNORECASE)
    description_pattern = re.compile(r"'Description'\s*=>\s*%q\{([^\}]*)\}", re.IGNORECASE | re.DOTALL)

    for foldername, _, filenames in os.walk(exploits_dir):
        for filename in filenames:
            if filename.endswith('.rb'):
                filepath = os.path.join(foldername, filename)
                try:
                    with open(filepath, 'r', encoding='utf-8', errors='ignore') as file:
                        content = file.read()
                        name_match = name_pattern.search(content)
                        description_match = description_pattern.search(content)
                        if name_match and description_match:
                            index.append({
                                'name': name_match.group(1),
                                'description': description_match.group(1),
                                'path': os.path.relpath(filepath, exploits_dir).replace(os.path.sep, '/')[:-3] #remove extension
                            })
                except Exception as e:
                    print(f"Error reading {filepath}: {e}")

    with open(MSF_INDEX_FILE, 'w') as f:
        json.dump(index, f)


if __name__=='__main__':
    build_index()


rebuild_msf_index.py
davidalami/VulnMapAI
import json
import multiprocessing
import time
from typing import List, Dict
from utils.logger import Logger
import nmap

from config import COMMON_PORTS


class NetworkScanner:
    """
    Сlass for performing network scans using nmap.
    """

    def __init__(self, target_subnets: List[str], top_ports=300):
        self.target_subnets = target_subnets
        self.top_ports = top_ports
        self.ports = NetworkScanner._load_common_ports(COMMON_PORTS)
        self.results_queue = multiprocessing.Queue()
        self.logger = Logger(__name__).get_logger()

    @staticmethod
    def _load_common_ports(common_ports_file: str) -> List[str]:
        with open(common_ports_file) as f:
            return json.load(f)

    def extract_port_info(self, scanned_host: str, protocol_type: str, port_number: int, scan_data: Dict) -> Dict:
        """
        Extracts information for a scanned port.

        Args:
            scanned_host (str): IP address of the scanned host.
            protocol_type (str): Protocol type, e.g., 'tcp'.
            port_number (int): Port number being scanned.
            scan_data (dict): Scan data returned from nmap scan.

        Returns:
            dict: A dictionary containing port information.
        """
        self.logger.info(f"Extracting port info for {scanned_host}, {protocol_type}, {port_number}")
        protocol_data = scan_data[protocol_type][port_number]
        port_info = {'host': scanned_host, 'protocol': protocol_type, 'port': port_number}
        if protocol_data['state'] != 'open':
            return None
        for key in ['name', 'product', 'version', 'script']:
            if key in protocol_data:
                port_info[key] = protocol_data[key]

        return port_info

    def scan_range(self, subnet_address: str, start_port: int, end_port: int):
        """
        Scans a range of ports on a subnet.

        Args:
            subnet_address (str): The address of the subnet to be scanned.
            start_port (int): The starting port number.
            end_port (int): The ending port number.
        """
        self.logger.info(f"Scanning range {start_port}-{end_port} on subnet {subnet_address}")
        port_scanner = nmap.PortScanner()
        ports_enumeration = ",".join(self.ports[start_port:end_port + 1])
        port_scanner.scan(hosts=subnet_address, arguments='-sV -T3 -sT --script=discovery -p' + ports_enumeration)

        for scanned_host in port_scanner.all_hosts():
            for protocol_type in port_scanner[scanned_host].all_protocols():
                if protocol_type != "tcp":
                    continue

                for port_number in port_scanner[scanned_host][protocol_type].keys():
                    port_info = self.extract_port_info(scanned_host, protocol_type, port_number,
                                                       port_scanner[scanned_host])
                    if port_info:
                        self.results_queue.put(port_info)

    def _create_processes(self, subnet_address: str, cpu_count: int, port_range: int, last_port: int) -> List:
        """
        Creates and starts a list of processes to scan a subnet.

        Args:
            subnet_address (str): The address of the subnet to be scanned.
            cpu_count (int): Number of CPU cores available.
            port_range (int): Range of ports to be scanned per process.
            last_port (int): The last port in the range to be scanned.

        Returns:
            list: List of created and started multiprocessing.Process objects.
        """
        processes = []
        for i in range(cpu_count):
            start_port = i * port_range
            end_port = (i + 1) * port_range if i < cpu_count - 1 else last_port
            p = multiprocessing.Process(target=self.scan_range, args=(subnet_address, start_port, end_port))
            processes.append(p)
            p.start()
        return processes

    def _yield_results(self, processes: List) -> Dict:
        """
        Yields results from the results_queue as they become available.

        Args:
            processes (list): List of multiprocessing.Process objects.

        Yields:
            dict: Port information dictionary from the results_queue.
        """
        while processes:
            while not self.results_queue.empty():
                yield self.results_queue.get()

            for p in processes[:]:
                if not p.is_alive():
                    processes.remove(p)
                    while not self.results_queue.empty():
                        yield self.results_queue.get()

            time.sleep(1)

    def execute_scan(self) -> Dict:
        """
        Executes the scanning, evently distributing the ports to be scanned among CPU cores.

        Yields:
            dict: Port information dictionary from the results_queue.
        """
        cpu_count = multiprocessing.cpu_count()
        port_range = self.top_ports // cpu_count

        self.logger.info(f"Starting the scanning process with {cpu_count} processes for {self.top_ports} most common ports")

        for subnet_address in self.target_subnets:
            processes = self._create_processes(subnet_address, cpu_count, port_range, self.top_ports)
            yield from self._yield_results(processes)

recon/async_scanner.py
davidalami/VulnMapAI
import pytest
import socket
import subprocess
import time


def is_port_open(host, port, timeout=1):
    """Check if a given port is open."""
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.settimeout(timeout)
        try:
            s.connect((host, port))
            return True
        except socket.error:
            return False


def test_server_start():
    """
    Test the startup of the server by running a Python script in the background.

    The test performs the following steps:
    1. Initiates the server by running 'main.py' script with '127.0.0.1' as an argument.
    2. Waits for an initial duration to account for any startup delay.
    3. Checks for 60 seconds to see if the server starts and becomes accessible on '127.0.0.1:1337'.
    4. If the server starts within the duration, the test is considered successful.
    5. If the server doesn't start within 60 seconds, or if the Python script encounters an issue,
       the test is marked as failed.

    Returns:
        None. If the server doesn't start as expected, an assertion error is raised.
    """
    # Run the python script in the background
    process = subprocess.Popen(["python", "main.py", "127.0.0.1"])

    # Give an initial delay, assuming the server might not start instantly
    time.sleep(2)

    # Check if the server starts within 60 seconds
    for _ in range(60):
        if is_port_open("127.0.0.1", 1337):
            # If server started, end the test as successful
            process.terminate()
            return

        time.sleep(1)

    # If the server didn't start within 60 seconds, or if the python process had a problem
    process.terminate()
    assert False, "Server did not start within 60 seconds or had an error"


tests/integration_test.py
davidalami/VulnMapAI
import logging

class Logger:

    def __init__(self, name, log_level=logging.INFO):
        self.logger = logging.getLogger(name)
        self.logger.setLevel(log_level)

        # Create a file handler for writing log files
        log_file_path = 'runtime.log'
        file_handler = logging.FileHandler(log_file_path)

        # Create a console handler for printing logs to the console
        console_handler = logging.StreamHandler()

        # Create a formatter and set it for both handlers
        formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
        file_handler.setFormatter(formatter)
        console_handler.setFormatter(formatter)

        # Add both handlers to the logger
        self.logger.addHandler(file_handler)
        self.logger.addHandler(console_handler)

    def get_logger(self):
        return self.logger

utils/logger.py
davidalami/VulnMapAI
import os
import json
import openai

from flask import (
    Flask,
    render_template,
    send_from_directory,
    request, jsonify
)

from config import REPORT_DIR


app = Flask(__name__)

app.config['REPORT_FOLDER'] = str(REPORT_DIR.absolute())


@app.route('/')
def index():
    """
    Serve the main index page which displays a list of all HTML reports.

    The function scans the report folder as specified in the app's configuration
    for all HTML files and then renders them in the `index.html` template.

    Returns:
        Rendered HTML template with the list of HTML report files.
    """
    files = [f for f in os.listdir(app.config['REPORT_FOLDER']) if f.endswith('.json')]
    return render_template("index.html", files=files)


@app.route('/report/<filename>')
def serve_report(filename):
    """
    Serves a specific report file from the configured report directory.

    Returns:
        The requested report file.
    """
    kwargs = json.load(open(os.path.join(app.config['REPORT_FOLDER'], filename)))
    return render_template("report_template.html", **kwargs)

@app.route('/ask-chatgpt', methods=['POST'])
def ask_chatgpt():
    """
    Handle POST requests to obtain responses from ChatGPT.

    This function is responsible for processing POST requests containing user queries
    and sending these queries to the ChatGPT API for responses. It then returns the
    generated response in a JSON format.

    Returns:
        JSON response containing the ChatGPT response to the user's query.
    """
    # Extract the user's query from the POST request data
    user_query = request.form['query']

    # Check if the 'query' parameter is provided in the request
    if not user_query:
        return jsonify({"error": "Please provide a 'query' parameter in the POST request."}), 400

    try:
        # Make a request to the ChatGPT API
        response = openai.Completion.create(
            engine="text-davinci-003",  # You can choose a different engine if needed
            prompt=user_query,
            max_tokens=50,  # Adjust the max tokens as per your requirements
            n = 1  # Number of responses to generate
        )

        # Extract the response text from the API response
        chatgpt_response = response.choices[0].text

        # Return the ChatGPT response as JSON
        return jsonify({"response": chatgpt_response})

    except Exception as e:
        # Handle any exceptions that may occur during the API request
        return jsonify({"error": str(e)}), 500


if __name__ == '__main__':
    app.run()
webapp/app.py
davidalami/VulnMapAI