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()

