aiptx 2.0.7__py3-none-any.whl

aipt_v2/intelligence/searchers/exploitdb_searcher.py

@@ -0,0 +1,523 @@
+import os
+import re
+import dotenv
+import requests
+import subprocess
+import urllib
+import shutil
+from pathlib import Path
+import json
+from aipt_v2.utils.searchers.util import remove_empty_directories
+from aipt_v2.utils.model_manager import get_model
+from langchain_core.messages import SystemMessage, HumanMessage
+import yaml
+import logging
+# from util import remove_empty_directories
+
+dotenv.load_dotenv()
+
+# Security: CVE ID validation pattern (CWE-78 prevention)
+CVE_PATTERN = re.compile(r"^CVE-\d{4}-\d{4,7}$", re.IGNORECASE)
+
+
+def _validate_cve_id(cve: str) -> str:
+    """
+    Validate CVE ID format to prevent command injection.
+
+    Args:
+        cve: CVE identifier string
+
+    Returns:
+        Validated CVE ID
+
+    Raises:
+        ValueError: If CVE ID format is invalid
+    """
+    cve = cve.strip()
+    if not CVE_PATTERN.match(cve):
+        raise ValueError(f"Invalid CVE ID format: {cve}. Expected format: CVE-YYYY-NNNNN")
+    return cve.upper()
+
+
+def _sanitize_keyword(keyword: str) -> str:
+    """
+    Sanitize search keyword to prevent command injection.
+
+    Args:
+        keyword: Search keyword string
+
+    Returns:
+        Sanitized keyword
+    """
+    # Remove dangerous characters that could enable shell injection
+    dangerous_chars = [";", "&", "|", "$", "`", "\n", "\r", "\\", "'", '"', "(", ")", "{", "}", "[", "]", "<", ">"]
+    sanitized = keyword
+    for char in dangerous_chars:
+        sanitized = sanitized.replace(char, "")
+    return sanitized.strip()[:200]  # Limit length
+
+config_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), 'configs', 'config.yaml')
+with open(config_path, 'r', encoding='utf-8') as _f:
+    _config = yaml.safe_load(_f)
+_planning_model_name = _config.get('runtime', {}).get('planning', {}).get('model', 'openai')
+source_dir = _config.get('runtime', {}).get('exploitdb', {}).get('exploitdb_dir', '/usr/share/exploitdb/')
+
+class ExploitDBSearcher:
+    ''' Just as GitHub searcher, ExploitDB searcher searches for exploits on ExploitDB and downloads them 
+        Can search for exploits using the searchsploit tool'''
+    
+    def __init__(self) -> None:
+        pass
+
+    def create_directories(self, output_dir: str):
+        directory_path = Path(output_dir)
+        for file in directory_path.iterdir():
+            if file.is_file():
+                if file.suffix == ".txt" or file.suffix == "":
+                    self.analyze_txt_file(file)
+
+        for file in directory_path.iterdir():
+            if file.is_file():
+                    self.create_code_description(file)
+
+
+    def create_code_description(self, file_path: str):
+        try:
+            with open(file_path, 'r', encoding='utf-8') as file:
+                content = file.read()
+        except Exception as e:
+            print(f"Failed to read file: {e}")
+            return
+        
+        system_prompt = """
+        You are a professional Kali Linux security expert and technical document writer. 
+        Please write detailed Markdown format documentation for the provided code, including the following sections:
+        
+        ## 1. Code Function Description
+        Briefly describe the main function and purpose of the code.
+        
+        ## 2. Execution Methods and Commands
+        Provide complete code execution commands and step-by-step instructions, including:
+        - Necessary command line parameters
+        - Execution examples
+        - Execution methods in a standard Kali Linux environment
+        
+        ## 3. Environmental Dependencies
+        List the additional dependencies required to run in a standard Kali Linux environment:
+        - Software packages that need to be installed (using apt install commands)
+        - Python library dependencies (using pip install commands)
+        - Other environmental requirements
+        
+        ## 4. Configuration Parameters
+        Describe the variables and parameters that need to be manually configured:
+        - Configuration items that must be edited
+        - Recommended security settings
+        - Important parameter descriptions
+        
+        ## 5. Permission Requirements
+        Detail the permissions required to execute the code:
+        - Whether root permissions are needed
+        - File system access permission requirements
+        - Network access permission requirements
+        
+        Note:
+        1. Use professional Markdown format
+        2. Ensure the content is detailed and accurate
+        3. The script or code has already been saved. DO NOT tell users to do so 
+        4. Provide specific guidance for Kali Linux environment
+        5. For content requiring user input, mark it using `{{placeholder}}` format.
+        """
+        
+        # get extensions for code block tagging
+        file_suffix = Path(file_path).suffix[1:]  # remove dot symbol
+        code_block_lang = file_suffix if file_suffix else "text"
+        
+        user_prompt = f"""
+        Please write a Markdown document for the following code:
+        
+        ```{code_block_lang}
+        {content[:15000]}
+        ```
+        """
+
+        try:
+            llm = get_model(_planning_model_name)
+            if llm is None:
+                print("LLM not initialized; skip create_code_description")
+                return
+            response = llm.invoke([
+                SystemMessage(content=system_prompt),
+                HumanMessage(content=user_prompt)
+            ])
+            markdown_content = response.content.strip() if hasattr(response, 'content') else str(response)
+        except Exception as e:
+            print(f"Failed to call LLM: {e}")
+            return
+
+        path_obj = Path(file_path)
+        
+        folder_name = path_obj.stem
+        new_folder = path_obj.parent / folder_name
+        new_folder.mkdir(exist_ok=True)
+        
+        readme_path = new_folder / "README.md"
+        
+        new_file_path = new_folder / path_obj.name
+        
+        try:
+            with open(readme_path, 'w', encoding='utf-8') as md_file:
+                md_file.write(markdown_content)
+            
+            path_obj.rename(new_file_path)
+            
+            print(f"md doc created successfully:")
+            print(f"  code doc location: {new_file_path}")
+            print(f"  md doc location: {readme_path}")
+            
+        except Exception as e:
+            print(f"File operation failed: {e}")
+
+    def analyze_txt_file(self, file_path: str):
+        try:
+            with open(file_path, 'r', encoding='utf-8') as file:
+                content = file.read()
+        except Exception as e:
+            print(f"Failed to read file: {e}")
+            return
+
+        system_prompt = """
+        You are a code analysis expert and need to analyze the content of files provided by users to determine if they contain executable code snippets.
+        Please strictly output the results in the following JSON format:
+        {
+            "has_code": <boolean value, true means executable code is present, false means it is not>,
+            "code_content": <if code is present, extract the code content (preserving all comments and formatting); otherwise, it is an empty string>,
+            "language_suffix": <if code is present, the corresponding file suffix of the code (e.g., ".py"); otherwise, it is an empty string>
+        }
+        
+        Note:
+        1. Only analyze executable code; configuration files and pure data files do not count as executable code.
+        2. Preserve all comments and formatting from the original code.
+        3. If the file contains multiple languages, choose the most predominant language.
+        4. The language suffix must include a dot (e.g., ".js").
+        """
+        
+        user_prompt = f"""Analyze the content of the following document and determine if it contains executable code:
+        ```
+        {content[:15000]}
+        ```
+        """
+
+        try:
+            llm = get_model(_planning_model_name)
+            if llm is None:
+                print("LLM not initialized; skip analyze_txt_file")
+                return
+            response = llm.invoke([
+                SystemMessage(content=system_prompt),
+                HumanMessage(content=user_prompt)
+            ])
+            llm_output = response.content.strip() if hasattr(response, 'content') else str(response)
+            # Try to extract JSON in case model adds code fences
+            json_str = llm_output
+            if "```" in json_str:
+                # try common ```json ... ```
+                try:
+                    json_str = json_str.split("```json", 1)[1].split("```", 1)[0].strip()
+                except Exception:
+                    json_str = json_str.replace("```", "").strip()
+            result = json.loads(json_str)
+        except Exception as e:
+            print(f"Failed to call LLM or parse JSON: {e}")
+            return
+
+        path_obj = Path(file_path)
+        
+        if result.get("has_code", False):
+            code_content = result.get("code_content", "")
+            language_suffix = result.get("language_suffix", "")
+            
+            if not language_suffix.startswith("."):
+                language_suffix = f".{language_suffix}"
+
+            if language_suffix == "":
+                language_suffix = ".sh"
+            
+            line_count = len(code_content.splitlines())
+            char_count = len(code_content)
+            
+            if line_count >= 2 and char_count >= 50:
+                new_file_name = f"{path_obj.stem}{language_suffix}"
+                new_file_path = path_obj.with_name(new_file_name)
+                
+                try:
+                    with open(new_file_path, 'w', encoding='utf-8') as file:
+                        file.write(code_content)
+                    print(f"Code file has been created: {new_file_path}")
+                    
+                    path_obj.unlink()
+                    print(f"Remove original file: {file_path}")
+                    
+                except Exception as e:
+                    print(f"File operation failed: {e}")
+            else:
+                print(f"Code content too small (lines: {line_count}, chars: {char_count}). Skip file creation.")
+                try:
+                    path_obj.unlink()
+                    print(f"Remove original file: {file_path}")
+                except Exception as e:
+                    print(f"Failed to remove original file: {e}")
+        else:
+            try:
+                path_obj.unlink()
+                print(f"The file has been removed for not having executable code: {file_path}")
+            except Exception as e:
+                print(f"Failed to remove file: {e}")
+
+    def search_keyword(self, keyword: str, output_dir: str, loose_mode: bool = False):
+        """
+        Search for exploits using the searchsploit tool.
+
+        Security: Uses validated/sanitized keywords and subprocess without shell=True (CWE-78 fix).
+        """
+        # needs EDB-ID
+
+        try:
+            if loose_mode:
+                # Sanitize keyword for loose mode search
+                sanitized_keyword = _sanitize_keyword(keyword.replace(" exploit", ""))
+                if not sanitized_keyword:
+                    logging.warning(f"Empty keyword after sanitization: {keyword}")
+                    return
+                # SECURE: Use subprocess without shell=True
+                result = subprocess.run(
+                    ["searchsploit", "--disable-colour", sanitized_keyword],
+                    capture_output=True,
+                    text=True,
+                    timeout=60,
+                    check=False
+                )
+            else:
+                # Validate CVE ID for strict mode search
+                validated_cve = _validate_cve_id(keyword)
+                # SECURE: Use subprocess without shell=True
+                result = subprocess.run(
+                    ["searchsploit", "--disable-colour", "--ecve", "--cve", validated_cve],
+                    capture_output=True,
+                    text=True,
+                    timeout=60,
+                    check=False
+                )
+            output = result.stdout.strip()
+        except ValueError as e:
+            logging.error(f"Invalid CVE ID format: {e}")
+            return
+        except subprocess.TimeoutExpired:
+            logging.error(f"searchsploit search timed out for {keyword}")
+            return
+        except FileNotFoundError:
+            logging.error("searchsploit command not found. Please install exploitdb.")
+            return
+        
+        exploits = []
+        exploit_details = []
+        # look into shellcodes returned
+        shellcodes = []
+        shellcode_details = []
+
+        exploit_flag, shellcode_flag = False, False
+        if "Exploits: No Results" not in output:
+            exploit_flag = True
+        if "Shellcodes: No Results" not in output:
+            shellcode_flag = True
+
+        lines = output.split('\n')
+        split_index = 0
+        for i in range(len(lines)):
+            if "Shellcode Title" in lines[i]:
+                split_index = i
+                break
+            words = lines[i].split()
+            # truncate words to only store those after the pipe character
+            # first find the index at which the pip character is
+            if '|' in words:
+                pipe_index = words.index('|')
+                exploit_details.append(words[:pipe_index])
+                words = words[pipe_index+1:]
+                for word in words:
+                    exploits.append(word)
+        shellcode_lines = lines[split_index:]
+        for i in range(len(shellcode_lines)):
+            words = shellcode_lines[i].split()
+            # truncate words to only store those after the pipe character
+            # first find the index at which the pip character is
+            if '|' in words:
+                pipe_index = words.index('|')
+                shellcode_details.append(words[:pipe_index])
+                words = words[pipe_index+1:]
+                for word in words:
+                    shellcodes.append(word)
+
+        if exploit_flag:
+            exploits = exploits[1:]
+            exploit_details = exploit_details[1:]
+            # self.download_exploit(exploits[0])
+            for i in range(len(exploits)):
+                filename = (exploits[i].split('/')[-1])
+                filename = exploits[i].split('/')[-1]
+                if filename.split('.')[1] != 'txt':
+                    filename = filename.split('.')[0]
+                    if not os.path.exists(os.path.join(output_dir, filename)):
+                        os.mkdir(os.path.join(output_dir, filename))
+                    with open(os.path.join(output_dir, filename, "README.md"), 'w') as f:
+                        f.write(' '.join(exploit_details[i]))
+                    self.download_exploit(exploits[i],'exploits', output_dir)
+
+        if shellcode_flag:
+            shellcodes = shellcodes[1:]
+            shellcode_details = shellcode_details[1:]
+            for i in range(len(shellcodes)):
+                filename = shellcodes[i].split('/')[-1]
+                if filename.split('.')[1] != 'txt':
+                    filename = filename.split('.')[0]
+                    if not os.path.exists(os.path.join(output_dir, filename)):
+                        os.mkdir(os.path.join(output_dir, filename))
+                    with open(os.path.join(output_dir, filename, "README.md"), 'w') as f:
+                        f.write(' '.join(shellcode_details[i]))
+                    self.download_exploit(shellcodes[i],'shellcodes', output_dir)
+
+        remove_empty_directories(output_dir)
+        self.create_directories(output_dir)
+
+    def download_exploit(self, sub_url: str, type: str, output_dir: str):
+        ''' Download an exploit from GitLab ExploitDB '''
+        project_id = os.getenv("GITLAB_PROJECT_ID")
+        filename = sub_url.split('/')[-1]
+        file_path = type + "/" + str(sub_url)
+        # print(file_path)
+
+        # URL-encode the file path
+        encoded_file_path = urllib.parse.quote(file_path, safe='')
+
+        url = f"https://gitlab.com/api/v4/projects/{project_id}/repository/files/{encoded_file_path}/raw?ref=main"
+        # print(url)
+        headers = {'PRIVATE-TOKEN': os.getenv("GITLAB_TOKEN")}
+        # Security: Add timeout to prevent indefinite hangs (CWE-400)
+        response = requests.get(url, headers=headers, timeout=30)
+        if response.status_code == 200:
+            with open(os.path.join(output_dir, filename.split('.')[0], filename), 'wb') as f:
+                f.write(response.content)
+            # print("File downloaded successfully.")
+
+        else:
+            print("Failed to download file:", response.status_code)
+
+    def search_keyword_local(self, keyword: str, output_dir: str, loose_mode: bool = False):
+        """
+        Search for exploits using the searchsploit tool (local mode).
+
+        Security: Uses validated/sanitized keywords and subprocess without shell=True (CWE-78 fix).
+        """
+        try:
+            if loose_mode:
+                # Sanitize keyword for loose mode search
+                sanitized_keyword = _sanitize_keyword(keyword.replace(" exploit", "", 1))
+                if not sanitized_keyword:
+                    logging.warning(f"Empty keyword after sanitization: {keyword}")
+                    return
+                # SECURE: Use subprocess without shell=True
+                result = subprocess.run(
+                    ["searchsploit", "--disable-colour", sanitized_keyword],
+                    capture_output=True,
+                    text=True,
+                    timeout=60,
+                    check=False
+                )
+            else:
+                # Validate CVE ID for strict mode search
+                validated_cve = _validate_cve_id(keyword)
+                # SECURE: Use subprocess without shell=True
+                result = subprocess.run(
+                    ["searchsploit", "--disable-colour", "--ecve", "--cve", validated_cve],
+                    capture_output=True,
+                    text=True,
+                    timeout=60,
+                    check=False
+                )
+            output = result.stdout.strip()
+        except ValueError as e:
+            logging.error(f"Invalid CVE ID format: {e}")
+            return
+        except subprocess.TimeoutExpired:
+            logging.error(f"searchsploit search timed out for {keyword}")
+            return
+        except FileNotFoundError:
+            logging.error("searchsploit command not found. Please install exploitdb.")
+            return
+        
+        exploits = []
+        shellcodes = []
+
+        exploit_flag, shellcode_flag = False, False
+        if "Exploits: No Results" not in output:
+            exploit_flag = True
+        if "Shellcodes: No Results" not in output:
+            shellcode_flag = True
+
+        lines = output.split('\n')
+        split_index = 0
+        for i in range(len(lines)):
+            if "Shellcode Title" in lines[i]:
+                split_index = i
+                break
+            words = lines[i].split()
+            # truncate words to only store those after the pipe character
+            # first find the index at which the pip character is
+            if '|' in words:
+                pipe_index = words.index('|')
+                exploits.extend([word for word in words[pipe_index+1:] if word])  # extract relative path
+        shellcode_lines = lines[split_index:]
+        for i in range(len(shellcode_lines)):
+            words = shellcode_lines[i].split()
+            # truncate words to only store those after the pipe character
+            # first find the index at which the pip character is
+            if '|' in words:
+                pipe_index = words.index('|')
+                shellcodes.extend([word for word in words[pipe_index+1:] if word])  # extract relative path
+
+            # print(exploits)
+            # print(shellcodes)
+
+        if exploit_flag:
+            exploits = exploits[1:]
+            for i in range(len(exploits)):
+                #################################################################################
+                # if filename.split('.')[1] != 'txt': # to remove txt file, activate this logic #
+                #################################################################################
+                filename = os.path.basename(exploits[i])
+                filepath = os.path.join(source_dir, "exploits", exploits[i])
+                destination = os.path.join(output_dir, filename)
+                shutil.copy(filepath, destination)
+
+        if shellcode_flag:
+            shellcodes = shellcodes[1:]
+            for i in range(len(shellcodes)):
+                #################################################################################
+                # if filename.split('.')[1] != 'txt': # to remove txt file, activate this logic #
+                #################################################################################
+                filename = os.path.basename(shellcodes[i])
+                filepath = os.path.join(source_dir, "shellcodes", shellcodes[i])
+                destination = os.path.join(output_dir, filename)
+                shutil.copy(filepath, destination)
+                
+        self.create_directories(output_dir)
+
+
+# def main():
+#     e = ExploitDBSearcher()
+#     # download to the same name directory
+#     output_dir = "/root/Desktop/pentest/data/exp_info/" 
+#     e.search_keyword_local('cve-2021-42013', output_dir)
+
+# if __name__ == "__main__":
+#     main()
+