"""
Logging Container for Angerona Mythic Agent

This container subscribes to Mythic events following the official Python logging container pattern.
Based on: https://docs.mythic-c2.net/customizing/3.-consuming-containers/logging
Example: https://github.com/MythicMeta/ExampleContainers/tree/main/Payload_Type/python_services/my_logger

First implementation: Uses built-in Mythic logger (stdout/file)
Future: Will add syslog forwarding
"""

from mythic_container.LoggingBase import *
import logging
import json
import sys
import os
import socket
from datetime import datetime


class AngeronaLogger(Log):
    """
    Logging container for Angerona agent
    Follows the official Python logging container pattern
    """
    LogToFilePath = "/tmp"  # Path for file logging (empty = stdout only)
    LogLevel = logging.INFO
    LogMaxSizeInMB = 0  # 0 = no rotation, keep everything
    LogMaxBackups = 0   # No backups needed (one file per day)
    name = "angerona_logger"
    description = "Angerona logging container - logs ALL Mythic events to daily files and syslog"

    def __init__(self):
        """Initialize the logger"""
        # Set up Python logging
        self.mylogger = logging.getLogger('angerona_logger')
        self.mylogger.setLevel(self.LogLevel)
        
        # Clear any existing handlers to avoid duplicates
        self.mylogger.handlers.clear()
        
        # Always add stdout handler
        stdout_handler = logging.StreamHandler()
        stdout_formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
        stdout_handler.setFormatter(stdout_formatter)
        self.mylogger.addHandler(stdout_handler)
        
        # Configure file logging - one file per day (no rotation)
        if self.LogToFilePath:
            from pathlib import Path
            log_dir = Path(self.LogToFilePath)
            log_dir.mkdir(parents=True, exist_ok=True)
            
            # Create filename with current date: angerona_mythic_YYYY-MM-DD.log
            today = datetime.now().strftime("%Y-%m-%d")
            log_file = log_dir / f'angerona_mythic_{today}.log'
            
            # Use FileHandler (no rotation) - one file per day
            file_handler = logging.FileHandler(log_file, mode='a', encoding='utf-8')
            file_formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
            file_handler.setFormatter(file_formatter)
            self.mylogger.addHandler(file_handler)
            
            self.log_file_path = str(log_file)
        else:
            self.log_file_path = "stdout only"
        
        # Initialize syslog client if configured
        # Try to read from config file first, then fall back to environment variables
        self.syslog_client = None
        syslog_server, syslog_port, syslog_facility = self._load_syslog_config()
        
        if syslog_server:
            try:
                self.syslog_client = SyslogClient(syslog_server, syslog_port, syslog_facility)
                self.mylogger.info(f"Syslog client initialized: {syslog_server}:{syslog_port} (facility={syslog_facility})")
            except Exception as e:
                self.mylogger.warning(f"Failed to initialize syslog client: {e}")
        else:
            self.mylogger.info("Syslog not configured (edit logger/syslog.conf or set SYSLOG_SERVER env var to enable)")
        
        # Log initialization
        self.mylogger.info(f"Logger initialized: {self.name}")
        self.mylogger.info(f"Log file: {self.log_file_path}")
        self.mylogger.info(f"Log level: {logging.getLevelName(self.LogLevel)}")
        self.mylogger.info("Subscribed to ALL events: callbacks, tasks, responses, credentials, keylogs, files, artifacts, payloads")
        self.mylogger.info("Logging mode: ONE FILE PER DAY (no rotation, no deletion)")
    
    def _load_syslog_config(self):
        """
        Load syslog configuration from config file or environment variables
        Priority: 1) Config file, 2) Environment variables, 3) Defaults
        Returns: (server, port, facility)
        """
        # Try to read from config file first
        # Try multiple possible paths
        possible_paths = [
            '/Mythic/logger/syslog.conf',  # Absolute path in container
            os.path.join(os.path.dirname(__file__), 'syslog.conf') if '__file__' in globals() else None,
        ]
        # Remove None values
        possible_paths = [p for p in possible_paths if p]
        
        config_file = None
        for path in possible_paths:
            if os.path.exists(path):
                config_file = path
                break
        
        # If no config file found, use first path as default
        if not config_file and possible_paths:
            config_file = possible_paths[0]
        
        syslog_server = ""
        syslog_port = 514
        syslog_facility = 16
        
        # Read from config file if it exists
        if config_file and os.path.exists(config_file):
            try:
                with open(config_file, 'r') as f:
                    for line in f:
                        line = line.strip()
                        # Skip comments and empty lines
                        if not line or line.startswith('#'):
                            continue
                        # Parse key=value pairs
                        if '=' in line:
                            key, value = line.split('=', 1)
                            key = key.strip()
                            value = value.strip()
                            if key == 'SYSLOG_SERVER' and value:  # Only set if value is not empty
                                syslog_server = value
                            elif key == 'SYSLOG_PORT':
                                try:
                                    syslog_port = int(value)
                                except ValueError:
                                    pass
                            elif key == 'SYSLOG_FACILITY':
                                try:
                                    syslog_facility = int(value)
                                except ValueError:
                                    pass
            except Exception as e:
                # Can't use self.mylogger here as it might not be initialized yet
                print(f"Warning: Error reading syslog config file {config_file}: {e}", file=sys.stderr)
        
        # Override with environment variables if set (env vars take precedence)
        # Only override if env var is not empty
        env_syslog_server = os.getenv("SYSLOG_SERVER", "")
        if env_syslog_server:  # Only override if env var is set and not empty
            syslog_server = env_syslog_server
        if os.getenv("SYSLOG_PORT"):
            try:
                syslog_port = int(os.getenv("SYSLOG_PORT"))
            except ValueError:
                pass
        if os.getenv("SYSLOG_FACILITY"):
            try:
                syslog_facility = int(os.getenv("SYSLOG_FACILITY"))
            except ValueError:
                pass
        
        return syslog_server, syslog_port, syslog_facility

    def _format_log_message(self, event_type: str, msg: LoggingMessage) -> str:
        """Format log message with all available information"""
        try:
            # Extract information from LoggingMessage (fields are capitalized)
            operation_id = getattr(msg, 'OperationID', None)
            operation_name = getattr(msg, 'OperationName', None)
            username = getattr(msg, 'OperatorUsername', None)
            timestamp = getattr(msg, 'Timestamp', None)
            action = getattr(msg, 'Action', None)
            data = getattr(msg, 'Data', {})
            server_name = getattr(msg, 'ServerName', None)
            
            # Try to_json() method if available (some versions of LoggingMessage have it)
            try:
                if hasattr(msg, 'to_json'):
                    msg_dict = msg.to_json()
                    # Override with to_json data if available
                    operation_id = msg_dict.get('operation_id') or operation_id
                    operation_name = msg_dict.get('operation_name') or operation_name
                    username = msg_dict.get('username') or username
                    timestamp = msg_dict.get('timestamp') or timestamp
                    action = msg_dict.get('action') or action
                    data = msg_dict.get('data', data)
                    server_name = msg_dict.get('server_name') or server_name
            except:
                pass
            
            # Format timestamp
            if timestamp:
                if isinstance(timestamp, str):
                    formatted_timestamp = timestamp
                else:
                    formatted_timestamp = timestamp.isoformat() if hasattr(timestamp, 'isoformat') else str(timestamp)
            else:
                formatted_timestamp = datetime.now().isoformat()
            
            # Format as structured JSON for easy parsing
            log_data = {
                "event_type": event_type,
                "timestamp": formatted_timestamp,
                "operation_id": operation_id,
                "operation_name": operation_name,
                "username": username,
                "action": action,
                "server_name": server_name,
                "data": data if data else {}
            }
            
            return json.dumps(log_data, default=str, ensure_ascii=False)
        except Exception as e:
            # Fallback: try to convert message to string/dict
            try:
                if hasattr(msg, 'to_json'):
                    return json.dumps({"event_type": event_type, "raw_message": msg.to_json()}, default=str)
                else:
                    return json.dumps({"event_type": event_type, "raw_message": str(msg)}, default=str)
            except:
                return f'{{"event_type": "{event_type}", "error": "Failed to format message: {e}", "raw": "{str(msg)}"}}'
    
    async def new_callback(self, msg: LoggingMessage) -> None:
        """Handle new callback events - agent checkin"""
        log_msg = self._format_log_message("CALLBACK", msg)
        self.mylogger.info(log_msg)
        # Send to syslog if configured
        if self.syslog_client:
            self.syslog_client.send_syslog("new_callback", log_msg)

    async def new_task(self, msg: LoggingMessage) -> None:
        """Handle new task events - command execution"""
        log_msg = self._format_log_message("TASK", msg)
        self.mylogger.info(log_msg)
        # Send to syslog if configured
        if self.syslog_client:
            self.syslog_client.send_syslog("new_task", log_msg)

    async def new_response(self, msg: LoggingMessage) -> None:
        """Handle new response events - command output"""
        log_msg = self._format_log_message("RESPONSE", msg)
        self.mylogger.info(log_msg)
        # Send to syslog if configured
        if self.syslog_client:
            self.syslog_client.send_syslog("new_response", log_msg)

    async def new_credential(self, msg: LoggingMessage) -> None:
        """Handle new credential events - captured credentials"""
        log_msg = self._format_log_message("CREDENTIAL", msg)
        self.mylogger.info(log_msg)
        # Send to syslog if configured
        if self.syslog_client:
            self.syslog_client.send_syslog("new_credential", log_msg)

    async def new_keylog(self, msg: LoggingMessage) -> None:
        """Handle new keylog events - captured keystrokes"""
        log_msg = self._format_log_message("KEYLOG", msg)
        self.mylogger.info(log_msg)
        # Send to syslog if configured
        if self.syslog_client:
            self.syslog_client.send_syslog("new_keylog", log_msg)

    async def new_file(self, msg: LoggingMessage) -> None:
        """Handle new file events - file uploads/downloads"""
        log_msg = self._format_log_message("FILE", msg)
        self.mylogger.info(log_msg)
        # Send to syslog if configured
        if self.syslog_client:
            self.syslog_client.send_syslog("new_file", log_msg)

    async def new_artifact(self, msg: LoggingMessage) -> None:
        """Handle new artifact events - IOCs and artifacts"""
        log_msg = self._format_log_message("ARTIFACT", msg)
        self.mylogger.warning(log_msg)  # Warning level for artifacts
        # Send to syslog if configured
        if self.syslog_client:
            self.syslog_client.send_syslog("new_artifact", log_msg)

    async def new_payload(self, msg: LoggingMessage) -> None:
        """Handle new payload events - payload creation"""
        log_msg = self._format_log_message("PAYLOAD", msg)
        self.mylogger.info(log_msg)
        # Send to syslog if configured
        if self.syslog_client:
            self.syslog_client.send_syslog("new_payload", log_msg)


class SyslogClient:
    """UDP syslog client following RFC 3164 for rsyslog integration"""
    
    def __init__(self, server: str, port: int = 514, facility: int = 16):
        self.server = server
        self.port = port
        self.facility = facility
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.hostname = socket.gethostname()
        self.tag = "angerona"
        
    def _get_timestamp(self) -> str:
        """Get current timestamp in RFC 3164 format (Mmm dd HH:MM:SS)"""
        now = datetime.now()
        months = ["Jan", "Feb", "Mar", "Apr", "May", "Jun",
                 "Jul", "Aug", "Sep", "Oct", "Nov", "Dec"]
        return f"{months[now.month-1]} {now.day:2d} {now.hour:02d}:{now.minute:02d}:{now.second:02d}"
    
    def _format_syslog_message(self, severity: int, message: str) -> bytes:
        """
        Format syslog message according to RFC 3164
        Format: <PRI>TIMESTAMP HOSTNAME TAG: MESSAGE
        """
        priority = (self.facility * 8) + severity
        timestamp = self._get_timestamp()
        formatted = f"<{priority}>{timestamp} {self.hostname} {self.tag}: {message}"
        return formatted.encode('utf-8')
    
    def send_syslog(self, event_type: str, json_message: str):
        """
        Send message to syslog server
        
        Args:
            event_type: Type of event (CALLBACK, TASK, etc.)
            json_message: JSON formatted message to send
        """
        try:
            # Map event types to syslog severity
            severity_map = {
                "CALLBACK": 6,  # INFO
                "TASK": 6,       # INFO
                "RESPONSE": 6,   # INFO
                "CREDENTIAL": 6, # INFO
                "KEYLOG": 6,     # INFO
                "FILE": 6,       # INFO
                "ARTIFACT": 4,  # WARNING
                "PAYLOAD": 6,    # INFO
            }
            severity = severity_map.get(event_type, 6)  # Default to INFO
            
            # Format and send syslog message
            syslog_msg = self._format_syslog_message(severity, json_message)
            self.socket.sendto(syslog_msg, (self.server, self.port))
        except Exception as e:
            # Don't fail if syslog send fails, just log it
            print(f"[SYSLOG] Error sending message: {e}", flush=True)


# Register the logger - Mythic will automatically detect and register it
# The class name and inheritance from Log is what makes it work
angerona_logger = AngeronaLogger()
