from mythic_container.MythicCommandBase import *
from mythic_container.MythicRPC import *
import re
import string, json
import os
from datetime import datetime

import logging

logging.basicConfig(level=logging.INFO)

class LsArguments(TaskArguments):
    def __init__(self, command_line, **kwargs):
        super().__init__(command_line, **kwargs)
        self.args = [
            CommandParameter(
                name="path", 
                type=ParameterType.String, 
                description="Path del fichero o de la carpeta del sistema",
                parameter_group_info=[ParameterGroupInfo(
                    required=False
                )]
            ),
        ]

    async def parse_arguments(self):
        if len(self.command_line) > 0:
            # Check if it's JSON from File Browser
            if self.command_line[0] == '{':
                try:
                    temp_json = json.loads(self.command_line)
                    if 'host' in temp_json:
                        # File Browser format: combine path + file
                        full_path = temp_json.get('full_path', '')
                        if full_path:
                            self.add_arg("path", full_path)
                        else:
                            path = temp_json.get('path', '')
                            file = temp_json.get('file', '')
                            if path and file:
                                # Normalize path separators
                                if not path.endswith('\\') and not path.endswith('/'):
                                    path = path + "\\"
                                self.add_arg("path", path + file)
                            elif path:
                                self.add_arg("path", path)
                            else:
                                self.add_arg("path", ".\\*")
                        self.add_arg("host", temp_json.get('host', ''))
                        self.add_arg("file_browser", "true")
                        return
                except:
                    pass
            # Regular command line usage
            self.add_arg("path", self.command_line)
            self.add_arg("host", "")
            self.add_arg("file_browser", "false")
        else:
            # No arguments, list current directory
            self.add_arg("path", ".\\*")
            self.add_arg("host", "")
            self.add_arg("file_browser", "false")

    async def parse_dictionary(self, dictionary):
        logging.info(f"Parse Dictionary - {dictionary}")
        if "host" in dictionary:
            # File Browser format
            logging.info(f"Command came from File Browser UI - {dictionary}")
            full_path = dictionary.get('full_path', '')
            if full_path:
                self.add_arg("path", full_path)
            else:
                path = dictionary.get("path", '')
                file = dictionary.get("file", '')
                if path and file:
                    if not path.endswith('\\') and not path.endswith('/'):
                        path = path + "\\"
                    self.add_arg("path", path + file)
                elif path:
                    self.add_arg("path", path)
                else:
                    self.add_arg("path", ".\\*")
            self.add_arg("host", dictionary.get("host", ""))
            self.add_arg("file_browser", "true")
        else:
            # Command line usage
            logging.info(f"Command came from CMDLINE - {dictionary}")
            arg_path = dictionary.get("path")
            if arg_path:
                self.add_arg("path", arg_path)
            else:
                self.add_arg("path", ".\\*")
            self.add_arg("host", "")
            self.add_arg("file_browser", "false")
        
        self.load_args_from_dictionary(dictionary)

class LsCommand(CommandBase):
    cmd = "ls"
    needs_admin = False
    help_cmd = "ls [path]"
    description = "List directory contents with detailed file information including type (File/Directory), size, modification date, and name. Supports wildcards (e.g., 'ls C:\\Windows\\*.exe'). If no path is specified, lists the current working directory. Compatible with Mythic File Browser for interactive file management."
    version = 1
    supported_ui_features = ["file_browser:list"]
    author = "Kaseya OFSTeam"
    attackmapping = ["T1106", "T1083"]
    argument_class = LsArguments
    browser_script = BrowserScript(
        script_name="ls", author="Kaseya OFSTeam", for_new_ui=True
    )
    attributes = CommandAttributes(
        builtin=False,
        supported_os=[ SupportedOS.Windows ],
        suggested_command=False
    )

    # async def create_tasking(self, task: MythicTask) -> MythicTask:
    #     task.display_params = task.args.get_arg("command")
    #     return task
    
    async def create_go_tasking(self, taskData: PTTaskMessageAllData) -> PTTaskCreateTaskingMessageResponse:
        response = PTTaskCreateTaskingMessageResponse(
            TaskID=taskData.Task.ID,
            Success=True,
        )
        
        path = taskData.args.get_arg("path")
        logging.info(f"create_go_tasking - path : {path}")
        response.DisplayParams = path
        
        # 
        if uncmatch := re.match(r"^\\\\(?P<host>[^\\]+)\\(?P<path>.*)$", path):
            taskData.args.add_arg("host", uncmatch.group("host"))
            taskData.args.set_arg("path", uncmatch.group("path"))
        else:
            # Set the host argument to an empty string if it does not exist
            taskData.args.add_arg("host", "")
        if host := taskData.args.get_arg("host"):
            host = host.upper()

            # Resolve 'localhost' and '127.0.0.1' aliases
            if host == "127.0.0.1" or host.lower() == "localhost":
                host = taskData.Callback.Host

            taskData.args.set_arg("host", host)

        return response

    async def process_response(self, task: PTTaskMessageAllData, response: any) -> PTTaskProcessResponseMessageResponse:
        resp = PTTaskProcessResponseMessageResponse(TaskID=task.Task.ID, Success=True)
        
        # Log that process_response was called
        logging.info(f"[FILE_BROWSER] process_response called for task {task.Task.ID}")
        logging.info(f"[FILE_BROWSER] response type: {type(response)}")
        if response:
            logging.info(f"[FILE_BROWSER] response keys: {list(response.keys()) if isinstance(response, dict) else 'not a dict'}")
        
        # Check if this was from File Browser
        is_file_browser = task.args.get_arg("file_browser") == "true"
        logging.info(f"[FILE_BROWSER] is_file_browser={is_file_browser}")
        
        # Check if translator already parsed file_browser format
        # The translator automatically detects and parses File Browser format,
        # so we only need to enhance it with host info if needed
        if isinstance(response, dict) and "file_browser" in response:
            logging.info(f"[FILE_BROWSER] Translator already parsed file_browser, just adding host if missing")
            file_browser_data = response.get("file_browser")
            if file_browser_data and not file_browser_data.get("host"):
                host = task.args.get_arg("host") or ""
                if not host and task.Callback:
                    host = task.Callback.Host
                file_browser_data["host"] = host
            resp.ProcessResponse = response
            return resp
        
        # Only parse if translator didn't already do it (fallback for non-File Browser commands)
        if is_file_browser and response:
            try:
                # Agent output format:
                # Line 1: directory path
                # Lines 2+: D/F\tSIZE\tMM/DD/YYYY HH:MM:SS\tNAME
                
                output_text = response.get("user_output", "")
                if not output_text:
                    return resp
                
                lines = [l.strip() for l in output_text.strip().split('\n') if l.strip()]
                if len(lines) < 1:
                    logging.warning("[FILE_BROWSER] No lines in output")
                    return resp
                
                logging.info(f"[FILE_BROWSER] Parsing {len(lines)} lines from output")
                logging.info(f"[FILE_BROWSER] First few lines: {lines[:3]}")
                
                # First line is the directory path being listed
                dir_path = lines[0].strip()
                
                # Normalize path
                if dir_path.endswith('\\*'):
                    dir_path = dir_path[:-2]
                if dir_path.endswith('*'):
                    dir_path = dir_path[:-1]
                dir_path = dir_path.rstrip('\\/')
                
                # Determine parent path
                if dir_path.endswith(':\\') or dir_path == '':
                    parent_path = ""
                else:
                    parent_path = os.path.dirname(dir_path) if os.path.dirname(dir_path) else ""
                
                # Get hostname
                host = task.args.get_arg("host") or ""
                if not host:
                    host = task.Callback.Host if task.Callback else ""
                
                # Determine if the listed item is a file or directory
                # Since we're listing a directory, it's typically not a file
                is_file = False
                name = os.path.basename(dir_path) if dir_path else ""
                if not name:
                    # Root directory case (e.g., C:\)
                    if len(dir_path) == 2 and dir_path[1] == ':':
                        name = dir_path[0] + ':'
                        parent_path = ""
                    else:
                        name = dir_path
                
                # Parse file/directory entries
                # Format: D/F\tSIZE\tMM/DD/YYYY HH:MM:SS\tNAME
                files = []
                for i, line in enumerate(lines[1:], 1):  # Start from line 1 (skip path line)
                    parts = line.split('\t')
                    logging.debug(f"[FILE_BROWSER] Line {i}: {line[:100]}... (parts: {len(parts)})")
                    if len(parts) >= 4:
                        file_type = parts[0].strip()  # D or F
                        size_str = parts[1].strip()
                        datetime_str = parts[2].strip()  # MM/DD/YYYY HH:MM:SS (already combined)
                        file_name = parts[3].strip() if len(parts) > 3 else ""
                        
                        # Skip empty lines or invalid entries
                        if not file_type or not file_name:
                            continue
                        
                        # Skip . and ..
                        if file_name in ['.', '..']:
                            continue
                        
                        try:
                            # Parse date: MM/DD/YYYY HH:MM:SS
                            date_obj = datetime.strptime(datetime_str, "%m/%d/%Y %H:%M:%S")
                            modify_time_ms = int(date_obj.timestamp() * 1000)
                            # Use modify_time as access_time (agent doesn't provide separate access time)
                            access_time_ms = modify_time_ms
                            
                            # Filter out invalid dates (before 1980 or after year 2100)
                            # This helps filter out files with corrupted timestamps (e.g., 1970 epoch issues)
                            if modify_time_ms < 315532800000 or modify_time_ms > 4102444800000:  # Before 1980-01-01 or after 2100-01-01
                                logging.warning(f"[FILE_BROWSER] Skipping file with invalid date: {file_name} (timestamp: {modify_time_ms}, date: {datetime_str})")
                                continue
                        except Exception as date_err:
                            logging.warning(f"Could not parse date '{datetime_str}': {date_err}")
                            # Skip files with unparseable dates
                            continue
                        
                        try:
                            # Parse size - remove any text like "b" or "by" at the end
                            size_str_clean = re.sub(r'[^\d]', '', size_str)  # Remove all non-digits
                            if size_str_clean:
                                size = int(size_str_clean) if file_type == 'F' else 0
                            else:
                                size = 0
                        except Exception as size_err:
                            logging.warning(f"Could not parse size '{size_str}': {size_err}")
                            size = 0
                        
                        # Get file attributes as permissions (Windows specific)
                        permissions = {
                            "attributes": "DIRECTORY" if file_type == 'D' else "FILE"
                        }
                        
                        # Skip files with invalid names or empty names
                        if file_name and file_name.strip() and file_name not in ['.', '..']:
                            files.append({
                                "is_file": file_type == 'F',
                                "permissions": permissions,
                                "name": file_name,
                                "access_time": access_time_ms,
                                "modify_time": modify_time_ms,
                                "size": size
                            })
                            logging.debug(f"[FILE_BROWSER] Added: {file_name} ({'file' if file_type == 'F' else 'dir'}, size={size})")
                    else:
                        logging.warning(f"[FILE_BROWSER] Skipping line {i} - insufficient parts: {len(parts)}")
                
                # Remove duplicates by name to ensure clean response
                seen_files = {}
                unique_files = []
                for f in files:
                    file_name = f.get("name", "")
                    # Only add if we haven't seen this file name before
                    if file_name and file_name not in seen_files:
                        seen_files[file_name] = True
                        unique_files.append(f)
                
                # Use current timestamp for modify_time to help Mythic identify fresh responses
                import time
                current_timestamp_ms = int(time.time() * 1000)
                
                # Build file_browser response
                file_browser_data = {
                    "host": host,
                    "is_file": is_file,
                    "permissions": {"attributes": "DIRECTORY"} if not is_file else {"attributes": "FILE"},
                    "name": name,
                    "parent_path": parent_path,
                    "success": True,
                    "access_time": current_timestamp_ms,  # Use current time to help with cache invalidation
                    "modify_time": current_timestamp_ms,  # Use current time to help with cache invalidation
                    "size": 0,
                    "update_deleted": True,  # Tell Mythic this is the complete list - mark missing files as deleted
                    "files": unique_files  # Use deduplicated files list - complete and fresh
                }
                
                logging.info(f"[FILE_BROWSER] Sending fresh listing for '{name}' with {len(unique_files)} unique files (timestamp: {current_timestamp_ms}, update_deleted=True)")
                
                # Add file_browser to response
                # Merge with existing response (like rm.py does)
                resp.ProcessResponse = response if isinstance(response, dict) else {}
                resp.ProcessResponse["file_browser"] = file_browser_data
                logging.info(f"[FILE_BROWSER] Added {len(files)} files/directories to file_browser response")
                logging.info(f"[FILE_BROWSER] ProcessResponse keys: {list(resp.ProcessResponse.keys())}")
                
            except Exception as e:
                logging.error(f"Error parsing file browser response: {str(e)}")
                import traceback
                logging.error(traceback.format_exc())
        
        return resp