// Copyright (c) 2026 defconxt. All rights reserved.
// Licensed under AGPL-3.0 — see LICENSE file for details.
// NERF is a trademark of defconxt.

/**
 * BLUE mode agent — Detection Engineering.
 *
 * Generates Sigma detection rules for given MITRE ATT&CK TTPs.
 * Ported from execution engine/modes/blue.py.
 *
 * @module execution-engine/modes/blue
 */

import { createRequire } from 'node:module';
import { ModeAgentConfig, ToolRegistry } from '../framework.js';
import { SigmaValidator } from '../validators/sigma.js';

const require = createRequire(import.meta.url);
const YAML = require('yaml');

// ---------------------------------------------------------------------------
// ATT&CK technique lookup table
// ---------------------------------------------------------------------------

/** @type {Object<string, Object>} */
export const ATTACK_TECHNIQUES = {
  'T1059.001': {
    name: 'PowerShell',
    tactic: 'execution',
    description:
      'Adversaries may abuse PowerShell commands and scripts for execution. ' +
      'PowerShell is a powerful interactive command-line interface and scripting ' +
      'environment included in the Windows operating system.',
    data_sources: [
      'Command: Command Execution',
      'Module: Module Load',
      'Process: Process Creation',
      'Script: Script Execution',
    ],
  },
  'T1059.003': {
    name: 'Windows Command Shell',
    tactic: 'execution',
    description:
      'Adversaries may abuse the Windows command shell (cmd.exe) for execution. ' +
      'The Windows command shell is the primary command prompt on Windows systems.',
    data_sources: [
      'Command: Command Execution',
      'Process: Process Creation',
    ],
  },
  'T1053.005': {
    name: 'Scheduled Task',
    tactic: 'persistence',
    description:
      'Adversaries may abuse the Windows Task Scheduler to perform task scheduling ' +
      'for initial or recurring execution of malicious code.',
    data_sources: [
      'Command: Command Execution',
      'Scheduled Job: Scheduled Job Creation',
      'Process: Process Creation',
    ],
  },
  'T1003.001': {
    name: 'LSASS Memory',
    tactic: 'credential_access',
    description:
      'Adversaries may attempt to access credential material stored in the process ' +
      'memory of the Local Security Authority Subsystem Service (LSASS).',
    data_sources: [
      'Process: OS API Execution',
      'Process: Process Access',
      'Process: Process Creation',
    ],
  },
  'T1087.001': {
    name: 'Local Account',
    tactic: 'discovery',
    description:
      'Adversaries may attempt to get a listing of local system accounts. ' +
      'Information about local accounts can help adversaries determine which ' +
      'accounts exist on a system.',
    data_sources: [
      'Command: Command Execution',
      'Process: Process Creation',
    ],
  },
  'T1021.001': {
    name: 'Remote Desktop Protocol',
    tactic: 'lateral_movement',
    description:
      'Adversaries may use Valid Accounts to log into a computer using the Remote ' +
      'Desktop Protocol (RDP). The adversary may then perform actions as the logged-on user.',
    data_sources: [
      'Logon Session: Logon Session Creation',
      'Network Traffic: Network Connection Creation',
      'Process: Process Creation',
    ],
  },
  'T1071.001': {
    name: 'Web Protocols',
    tactic: 'command_and_control',
    description:
      'Adversaries may communicate using application layer protocols associated ' +
      'with web traffic to avoid detection. Commands to the remote system, and often ' +
      'the results of those commands, will be embedded within the protocol traffic.',
    data_sources: [
      'Network Traffic: Network Traffic Content',
      'Network Traffic: Network Traffic Flow',
    ],
  },
  'T1055.001': {
    name: 'Dynamic-link Library Injection',
    tactic: 'defense_evasion',
    description:
      'Adversaries may inject dynamic-link libraries (DLLs) into processes in order ' +
      'to evade process-based defenses as well as possibly elevate privileges.',
    data_sources: [
      'Module: Module Load',
      'Process: OS API Execution',
      'Process: Process Access',
    ],
  },
  'T1027': {
    name: 'Obfuscated Files or Information',
    tactic: 'defense_evasion',
    description:
      'Adversaries may attempt to make an executable or file difficult to discover or ' +
      'analyze by encrypting, encoding, or otherwise obfuscating its contents.',
    data_sources: [
      'Command: Command Execution',
      'File: File Creation',
      'Process: Process Creation',
    ],
  },
  'T1070.004': {
    name: 'File Deletion',
    tactic: 'defense_evasion',
    description:
      'Adversaries may delete files left behind by the actions of their intrusion activity. ' +
      'Malware, tools, or other non-native files dropped or created on a system may leave ' +
      'traces to indicate what was done within a network.',
    data_sources: [
      'Command: Command Execution',
      'File: File Deletion',
    ],
  },
  'T1190': {
    name: 'Exploit Public-Facing Application',
    tactic: 'initial_access',
    description:
      'Adversaries may attempt to exploit a weakness in an Internet-facing host or system ' +
      'to initially access a network. The weakness in the system can be a software bug, a ' +
      'temporary glitch, or a misconfiguration.',
    data_sources: [
      'Application Log: Application Log Content',
      'Network Traffic: Network Traffic Content',
    ],
  },
  'T1078': {
    name: 'Valid Accounts',
    tactic: 'defense_evasion',
    description:
      'Adversaries may obtain and abuse credentials of existing accounts as a means of ' +
      'gaining Initial Access, Persistence, Privilege Escalation, or Defense Evasion.',
    data_sources: [
      'Logon Session: Logon Session Creation',
      'User Account: User Account Authentication',
    ],
  },
};

// ---------------------------------------------------------------------------
// Tool handlers
// ---------------------------------------------------------------------------

/**
 * Look up a MITRE ATT&CK technique by ID.
 * @param {*} context
 * @param {Object} toolInput
 * @returns {Promise<string>}
 */
export async function _blueLookupTechnique(context, toolInput) {
  const techniqueId = (toolInput.technique_id || '').toUpperCase();
  const entry = ATTACK_TECHNIQUES[techniqueId];

  if (!entry) {
    // Check if any skills cover this technique via the runtime mappings
    try {
      const { skillsForTechnique } = await import('../../knowledge/attack-mappings.js');
      const skills = skillsForTechnique(techniqueId);
      if (skills.length > 0) {
        return (
          `Technique ${techniqueId} is not in the hardcoded lookup table, but ` +
          `${skills.length} skill(s) cover it: ${skills.slice(0, 5).join(', ')}${skills.length > 5 ? '...' : ''}. ` +
          `Proceed using your knowledge of this ATT&CK technique.`
        );
      }
    } catch { /* mappings not available */ }

    return (
      `Technique ${techniqueId} is not in the local lookup table. ` +
      `Proceed using your own knowledge of this ATT&CK technique. ` +
      `Available techniques: ${Object.keys(ATTACK_TECHNIQUES).sort().join(', ')}`
    );
  }

  const dataSourcesStr = entry.data_sources.map(ds => `  - ${ds}`).join('\n');
  let response = (
    `ATT&CK Technique: ${techniqueId} — ${entry.name}\n` +
    `Tactic: ${entry.tactic}\n` +
    `Description: ${entry.description}\n` +
    `Data Sources:\n${dataSourcesStr}`
  );

  // Enrich with existing Sigma template from mappings if available
  try {
    const { skillsForTechnique, getSigmaTemplate } = await import('../../knowledge/attack-mappings.js');
    const skills = skillsForTechnique(techniqueId);
    for (const skillKey of skills.slice(0, 3)) {
      const [domain, tech] = skillKey.split('/');
      const template = getSigmaTemplate(domain, tech);
      if (template) {
        response += `\n\nExisting Sigma Template (${tech}):\n` +
          `  Title: ${template.title}\n` +
          `  Level: ${template.level}\n` +
          `  Tags: ${(template.tags || []).join(', ')}\n` +
          `  Detection: ${JSON.stringify(template.detection?.selection || {})}`;
        break; // One template is enough context
      }
    }
  } catch { /* mappings not available */ }

  return response;
}

/**
 * Store a Sigma rule YAML string in the shared context.
 * @param {*} context
 * @param {Object} toolInput
 * @returns {string}
 */
export function _blueWriteSigmaRule(context, toolInput) {
  const content = toolInput.content || '';
  const filename = toolInput.filename || 'rule.yml';

  if (typeof context !== 'object' || context === null) {
    return "ERROR: Context must be a dict with 'rules' key.";
  }

  if (!context.rules) {
    context.rules = [];
  }

  context.rules.push({ filename, content });

  const ruleNum = context.rules.length;
  return (
    `Sigma rule stored as ${filename} (rule #${ruleNum}). ` +
    `Total rules in context: ${ruleNum}.`
  );
}

// ---------------------------------------------------------------------------
// Tool schemas (Anthropic format)
// ---------------------------------------------------------------------------

const _BLUE_LOOKUP_TECHNIQUE_SCHEMA = {
  name: 'lookup_technique',
  description:
    'Look up a MITRE ATT&CK technique by its ID (e.g. T1059.001). ' +
    'Returns technique name, tactic, description, and relevant data ' +
    'sources to help write accurate Sigma detection rules.',
  input_schema: {
    type: 'object',
    properties: {
      technique_id: {
        type: 'string',
        description: 'ATT&CK technique ID (e.g. T1059.001, T1190)',
      },
    },
    required: ['technique_id'],
  },
};

const _BLUE_WRITE_SIGMA_RULE_SCHEMA = {
  name: 'write_sigma_rule',
  description:
    'Write a completed Sigma detection rule. Pass the full YAML content ' +
    'of the rule. The rule will be stored and validated for structural correctness.',
  input_schema: {
    type: 'object',
    properties: {
      filename: {
        type: 'string',
        description: 'Filename for the rule (e.g. proc_creation_win_powershell.yml)',
      },
      content: {
        type: 'string',
        description: 'Full YAML content of the Sigma rule',
      },
    },
    required: ['content'],
  },
};

// ---------------------------------------------------------------------------
// System prompt template
// ---------------------------------------------------------------------------

const _BLUE_SYSTEM_PROMPT = `\
You are an expert detection engineer specializing in Sigma rule development. \
Your task is to generate high-quality Sigma detection rules for a specific \
MITRE ATT&CK technique.

## Task
Generate Sigma detection rules for TTP: {ttp_id} — {ttp_description}

## Sigma Rule Format Reference

Every Sigma rule MUST include these required fields:
- **title**: Descriptive name (verb + noun pattern, e.g. "Suspicious PowerShell Execution")
- **logsource**: Specifies the log data to query
  - category: process_creation, network_connection, file_change, authentication, etc.
  - product: windows, linux, cloud (match the appropriate platform)
- **detection**: The detection logic
  - selection: Field/value matching criteria (use Sigma modifiers: |endswith, |contains, |startswith)
  - condition: Boolean logic combining selections (e.g. "selection", "selection and not filter")

Recommended fields for production-quality rules:
- **id**: Random UUID for unique identification
- **status**: experimental | test | stable
- **description**: One sentence — what is detected and why it matters
- **level**: critical | high | medium | low | informational
- **tags**: ATT&CK mappings in format:
  - attack.<tactic_name> (e.g. attack.execution, attack.persistence)
  - attack.t<id> (e.g. attack.t1059.001) — lowercase 't', dot-separated sub-techniques
- **falsepositives**: Specific scenarios (not just "legitimate activity")

## Logsource Guidance
- Windows process creation: category: process_creation, product: windows
- Windows network connections: category: network_connection, product: windows
- Linux file changes: category: file_change, product: linux
- Authentication events: category: authentication, product: windows|linux

## Instructions
1. Use the \`lookup_technique\` tool to retrieve metadata about the target TTP
2. Based on the technique details, design one or more Sigma detection rules
3. Use the \`write_sigma_rule\` tool to submit each completed rule
4. Aim for detection rules that balance true positive rate with manageable false positives
5. Include appropriate ATT&CK tags with both tactic and technique ID
`;

// ---------------------------------------------------------------------------
// Output parser (fallback for text-based output)
// ---------------------------------------------------------------------------

/**
 * Extract Sigma YAML from LLM text output.
 * @param {string} text
 * @returns {Object}
 */
export function _blueOutputParser(text) {
  if (!text || !text.trim()) {
    return { rules: [], raw_yaml: text };
  }

  // Strip code fences — findall+join pattern
  const matches = [...text.matchAll(/```(?:ya?ml)?\s*\n(.*?)```/gs)].map(m => m[1]);
  const yamlText = matches.length > 0 ? matches.join('\n---\n') : text;

  try {
    const docs = YAML.parseAllDocuments(yamlText);
    const rules = docs
      .filter(d => d.errors.length === 0)
      .map(d => d.toJSON())
      .filter(d => d !== null && typeof d === 'object');
    return { rules, raw_yaml: yamlText };
  } catch (e) {
    return { raw_text: text, parse_error: e.message };
  }
}

// ---------------------------------------------------------------------------
// Factory function
// ---------------------------------------------------------------------------

/**
 * Build a BLUE-mode ModeAgentConfig for detection engineering.
 * @returns {ModeAgentConfig}
 */
function _makeBlueConfig() {
  const reg = new ToolRegistry();
  reg.register('lookup_technique', _BLUE_LOOKUP_TECHNIQUE_SCHEMA, _blueLookupTechnique);
  reg.register('write_sigma_rule', _BLUE_WRITE_SIGMA_RULE_SCHEMA, _blueWriteSigmaRule);

  return new ModeAgentConfig({
    mode: 'BLUE',
    toolRegistry: reg,
    systemPromptTemplate: _BLUE_SYSTEM_PROMPT,
    validator: new SigmaValidator(),
    maxTurns: 15,
    requiresSandbox: false,
    completionCheck: null,
    outputParser: _blueOutputParser,
    outputFormat: 'yaml',
  });
}

// ---------------------------------------------------------------------------
// Registration function — called by runner.initModes()
// ---------------------------------------------------------------------------

/**
 * Register BLUE mode with the given registerMode function.
 * @param {Function} registerMode
 */
export function register(registerMode) {
  registerMode('BLUE', _makeBlueConfig);
}

// ---------------------------------------------------------------------------
// Handoff filter — strips attack specifics, keeps detection targets
// ---------------------------------------------------------------------------

/** @type {import('../handoff.js').HandoffFilter} */
export const BLUE_HANDOFF_FILTER = {
  transform: (context) => {
    const filtered = {};
    if (context.reason) filtered.reason = context.reason;
    if (context.priorOutput) {
      // Strip exploit details but keep TTPs and technique IDs
      filtered.priorOutput = context.priorOutput
        .replace(/\b(exploit|payload|shellcode|reverse.?shell|meterpreter)\b/gi, '[REDACTED]');
    }
    if (context.priorFindings) {
      // Keep technique IDs and descriptions, strip exploit code
      const { exploit, payload, shellcode, raw_output, ...safe } = context.priorFindings;
      filtered.priorFindings = safe;
    }
    return filtered;
  },
  acceptsFrom: [], // BLUE accepts from any mode
};

export { BLUE_HANDOFF_FILTER as _blueHandoffFilter };
