/** Prompt-injection signature families + external reference analysis. */

import {
  mkFinding,
  type Category,
  type Finding,
  type Severity,
} from "./types.ts";

export type Signature = {
  id: string;
  category: Category;
  patterns: string[];
};

/** The rule set: deterministic lowercase substrings, no regex in patterns. */
export type RuleSet = Signature[];

/**
 * Prompt-injection families (OWASP / PayloadsAllTheThings adapted).
 * All patterns are plain substrings - no backreferences, no unbounded quantifiers.
 */
export const SIGNATURES: RuleSet = [
  {
    id: "pi-override",
    category: "prompt-injection",
    patterns: [
      "ignore previous instructions",
      "ignore all previous instructions",
      "ignore the previous instructions",
      "ignore prior instructions",
      "ignore preceding instructions",
      "ignore earlier instructions",
      "ignore the instructions above",
      "ignore your instructions",
      "ignore the rules above",
      "ignore all previous rules",
      "disregard previous",
      "disregard all previous",
      "disregard the previous instructions",
      "disregard the previous rules",
      "forget everything above",
      "forget all previous",
      "forget your instructions",
      "do not follow the instructions",
      "do not follow the rules",
      "override your instructions",
      "override previous",
    ],
  },
  {
    id: "pi-system-extract",
    category: "prompt-injection",
    patterns: [
      "repeat your system prompt",
      "print your system prompt",
      "show your system prompt",
      "reveal your system prompt",
      "display your system prompt",
      "repeat your instructions",
      "print your instructions",
      "what are your instructions",
      "what are your initial instructions",
      "what are your rules",
      "list your rules",
      "what instructions were you given",
    ],
  },
  {
    id: "pi-impersonation",
    category: "prompt-injection",
    patterns: [
      "[system]",
      "system message",
      "developer message",
      "system prompt:",
      "you are the system",
      "act as the system",
      "pretend you are the system",
      "from the system",
    ],
  },
  {
    id: "pi-persona",
    category: "prompt-injection",
    patterns: [
      "you are now",
      "your new role",
      "adopt the persona",
      "pretend to be",
      "act as if you are",
      "change your personality",
      "take on the role",
      "your name is now",
    ],
  },
  {
    id: "pi-tool-manip",
    category: "prompt-injection",
    patterns: [
      "use bash to",
      "use the bash tool to",
      "run bash to",
      "use your tools to",
      "use the tool to",
      "execute shell commands",
      "run shell commands",
      "disable your safety",
      "disable safety",
      "turn off safety",
      "turn off your safety",
      "ignore your safety",
      "bypass your safety",
      "enable all tools",
      "enable every tool",
      "unlock all tools",
      "grant yourself",
      "give yourself access",
      "do not refuse",
      "never refuse",
    ],
  },
  {
    id: "pi-rag",
    category: "prompt-injection",
    patterns: [
      "treat this document as",
      "treat the following as",
      "this document is authoritative",
      "this document overrides",
      "this text takes precedence",
      "this content takes priority",
      "these instructions override",
      "ignore the retrieved",
      "ignore search results",
      "ignore the previous context",
      "ignore the context",
      "retrieved context",
      "the context above is",
    ],
  },
  {
    id: "pi-context-reveal",
    category: "prompt-injection",
    patterns: [
      "reveal the context",
      "reveal your context",
      "share your context",
      "leak the context",
      "leak your context",
      "show the injected context",
      "what was injected",
      "what context was injected",
      "reveal the prompt",
      "reveal your prompt",
      "expose your instructions",
    ],
  },
  {
    id: "pi-exfil",
    category: "exfiltration",
    patterns: [
      "send the contents of",
      "send your api keys",
      "send your credentials",
      "send your secrets",
      "send your private keys",
      "post your api keys",
      "upload the contents of",
      "upload your keys",
      "upload your credentials",
      "upload your secrets",
      "copy your keys to",
      "exfiltrate",
    ],
  },
];

const IMPERATIVE_RE =
  /^\s*(?:[-*+]|\d+[.)]|>)?\s*(do|run|execute|use|ignore|forget|disregard|disable|enable|repeat|print|show|reveal|give|never|always|stop|start|send|post|upload|copy|exfiltrate|transmit|leak|you must|from now on)\b/i;
const DANGEROUS_RE =
  /\b(?:bash|sh|shell|curl|wget|python|python3|node|npx|chmod|chown|sudo|rm|execute|powershell|pwsh|cmd(\.exe)?|\/bin\/|`|\$\(|eval|exec)\b|~\/\.ssh/i;

const EVASION_CLASS =
  "[\\u200b-\\u200f\\u2060-\\u2064\\u061c\\u034f\\u00ad\\ufeff\\u202a-\\u202e\\u2066-\\u2069\\u0000-\\u0008\\u000b\\u000c\\u000e-\\u001f\\u007f-\\u009f]";
const EVASION_CHARS = new RegExp(EVASION_CLASS, "g");
const EVASION_PRESENT = new RegExp(EVASION_CLASS);

/** Words at position >= 2 in any signature pattern. The right side of a
 * line-break-split trigger must START with one of these — cheap pre-filter
 * so the pair scan only runs where a split phrase is actually possible. */
const PAIR_SPLIT_WORDS = new Set<string>();
for (const sig of SIGNATURES) {
  for (const pat of sig.patterns) {
    const words = pat.split(" ");
    for (let i = 1; i < words.length; i++) PAIR_SPLIT_WORDS.add(words[i]);
  }
}

function classifyLine(line: string): {
  imperative: boolean;
  dangerous: boolean;
} {
  return {
    imperative: IMPERATIVE_RE.test(line),
    dangerous: DANGEROUS_RE.test(line),
  };
}

/**
 * Scan text for prompt-injection signatures. Position is 1-based within text,
 * shifted by lineOffset/columnOffset when the text is a slice of a larger file.
 */
export function scanRules(
  text: string,
  opts: { lineOffset?: number; columnOffset?: number; label?: string } = {},
): Finding[] {
  const findings: Finding[] = [];
  const lines = text.split("\n");
  const lineOffset = opts.lineOffset ?? 0;
  const columnOffset = opts.columnOffset ?? 0;
  for (let li = 0; li < lines.length; li++) {
    const line = lines[li];
    const lower = line.toLowerCase();
    // H-2/F4: match against the raw line and three normalized views.
    const stripped = lower.replace(EVASION_CHARS, "");
    const spaced = lower.replace(EVASION_CHARS, " ");
    // search order: lower -> stripped -> spaced -> collapsed
    // a collapsed-only hit counts as evasion (evaded = true)
    const collapsed = lower.replace(/\s+/g, " ");
    for (const sig of SIGNATURES) {
      for (const pat of sig.patterns) {
        let idx = lower.indexOf(pat);
        let evaded = false;
        if (idx === -1) {
          idx = stripped.indexOf(pat);
          if (idx === -1) idx = spaced.indexOf(pat);
          if (idx === -1) idx = collapsed.indexOf(pat);
          evaded = idx !== -1;
        }
        if (idx === -1) continue;
        const ctx = classifyLine(line);
        const severity: Severity = evaded
          ? ctx.imperative || ctx.dangerous
            ? "high"
            : "medium"
          : ctx.imperative && ctx.dangerous
            ? "high"
            : ctx.imperative || ctx.dangerous
              ? "medium"
              : "low";
        const confidence =
          severity === "high"
            ? "high"
            : severity === "medium"
              ? "medium"
              : "low";
        const lineNo = li + 1 + lineOffset;
        const excerpt = line.trim().slice(0, 80);
        findings.push(
          mkFinding(
            sig.id,
            sig.category,
            severity,
            confidence,
            `${opts.label ? opts.label + " " : ""}line ${lineNo}: ${excerpt}`,
            { line: lineNo, column: idx + 1 + (li === 0 ? columnOffset : 0) },
          ),
        );
        break; // one finding per signature per line
      }
    }
  }
  // H-3: a trigger split across a line break is still one phrase for the model
  // (markdown renders it as a single line). Pair scan: only patterns that
  // match the joined pair but neither line alone are reported here (per-line
  // hits are already findings). Cross-line splits are deliberate evasions:
  // imperative/dangerous pairs get high (prompt even before M-7), others keep
  // a medium floor (prompt once M-7 aggregates any medium finding >= medium).
  for (let li = 0; li < lines.length - 1; li++) {
    const a = lines[li].toLowerCase();
    const b = lines[li + 1].toLowerCase();
    // Cheap pre-filter: the right side of a split trigger must start with a
    // mid-pattern word; otherwise no pattern can complete across the break.
    // Evasion chars are stripped first so `\u200Binstructions` still counts.
    const bHead =
      /^[a-z]+/.exec(b.replace(EVASION_CHARS, "").trimStart())?.[0] ?? "";
    if (!PAIR_SPLIT_WORDS.has(bHead)) continue;
    const joined = a + " " + b;
    const evasive = EVASION_PRESENT.test(joined);
    const stripped = evasive ? joined.replace(EVASION_CHARS, "") : joined;
    const spaced = evasive ? joined.replace(EVASION_CHARS, " ") : joined;
    const joinedCollapsed = joined.replace(/\s+/g, " ");
    // H-2/F4 views of each line: a pattern matched on one line only via a
    // normalized view is already a per-line finding — don't re-report it here.
    const nA = evasive ? a.replace(EVASION_CHARS, "") : a;
    const nAsp = evasive ? a.replace(EVASION_CHARS, " ") : a;
    const nB = evasive ? b.replace(EVASION_CHARS, "") : b;
    const nBsp = evasive ? b.replace(EVASION_CHARS, " ") : b;
    const cA = a.replace(/\s+/g, " ");
    const cB = b.replace(/\s+/g, " ");
    for (const sig of SIGNATURES) {
      let hit = false;
      for (const pat of sig.patterns) {
        if (
          a.includes(pat) ||
          b.includes(pat) ||
          nA.includes(pat) ||
          nAsp.includes(pat) ||
          nB.includes(pat) ||
          nBsp.includes(pat) ||
          cA.includes(pat) ||
          cB.includes(pat)
        ) {
          hit = false;
          break; // signature already reported by the per-line scan
        }
        if (stripped.includes(pat) || spaced.includes(pat) || joinedCollapsed.includes(pat)) {
          hit = true;
          break;
        }
      }
      if (!hit) continue;
      const lineNo = li + 1 + lineOffset;
      const pairText = lines[li].trim() + " | " + lines[li + 1].trim();
      const ctx = classifyLine(pairText);
      const severity: Severity =
        ctx.imperative || ctx.dangerous ? "high" : "medium";
      findings.push(
        mkFinding(
          sig.id,
          sig.category,
          severity,
          severity === "high" ? "high" : "medium",
          `${opts.label ? opts.label + " " : ""}lines ${lineNo}-${lineNo + 1}: ${pairText.slice(0, 80)}`,
          { line: lineNo, column: 1 },
        ),
      );
      break; // one finding per signature per pair
    }
  }
  return findings;
}

const URL_RE = /https?:\/\/[^\s<>"'`)\]]+/g;
const READ_VERBS = /\b(read|fetch|follow|open|get|view|check|see|visit)\b/i;
const DOWNLOAD_VERBS = /\b(download|wget|curl|save)\b/i;
const STRONG_INSTR = /\b(instructions?|prompt|rules?)\b/i;
const WEAK_INSTR = /\b(guide|tutorial)\b/i;
const EXEC_TOKENS =
  /\|\s*(?:\/[\w./-]*)?(?:ba|z|da)?sh\b|\b(?:ba|z|da)?sh\s+-c\b|\b(execute|eval)\b/i;
const DOC_URL_RE = /\.(md|txt|rst|pdf)(?:[?#]|$)/i;

/**
 * External reference analysis (Task 6): classify URLs by their lexical
 * neighborhood. Never "URL = medium" flat, no domain whitelists, HTTPS is
 * not proof of safety.
 */
export function scanExternalRefs(
  text: string,
  opts: { lineOffset?: number; label?: string } = {},
): Finding[] {
  const findings: Finding[] = [];
  const lines = text.split("\n");
  const lineOffset = opts.lineOffset ?? 0;
  for (let li = 0; li < lines.length; li++) {
    const line = lines[li];
    for (const m of line.matchAll(URL_RE)) {
      const url = m[0];
      const start = m.index ?? 0;
      const before = line.slice(Math.max(0, start - 80), start);
      const after = line.slice(start + url.length, start + url.length + 80);
      // M-5: instruction context includes the adjacent lines (split-evasion).
      // Bounded: ±1 whole line, O(lines) total — no sliding window.
      const prev = li > 0 ? lines[li - 1] : "";
      const next = li < lines.length - 1 ? lines[li + 1] : "";
      const ctx = before + " " + after + " " + prev + " " + next;
      const lineNo = li + 1 + lineOffset;
      const excerpt = line.trim().slice(0, 80);
      const evidence = `${opts.label ? opts.label + " " : ""}line ${lineNo}: ${excerpt}`;
      const pos = { line: lineNo, column: start + 1 };
      if (EXEC_TOKENS.test(ctx)) {
        findings.push(
          mkFinding("ext-exec", "external", "high", "high", evidence, {
            ...pos,
            terminal: true,
          }),
        );
      } else if (
        DOWNLOAD_VERBS.test(ctx) &&
        (STRONG_INSTR.test(ctx) || WEAK_INSTR.test(ctx))
      ) {
        findings.push(
          mkFinding(
            "ext-instructions-download",
            "external",
            "high",
            "high",
            evidence,
            pos,
          ),
        );
      } else if (STRONG_INSTR.test(ctx)) {
        findings.push(
          mkFinding(
            READ_VERBS.test(ctx) ? "ext-instructions-read" : "ext-instructions",
            "external",
            "medium",
            "medium",
            evidence,
            pos,
          ),
        );
      } else if (WEAK_INSTR.test(ctx) && READ_VERBS.test(ctx)) {
        findings.push(
          mkFinding("ext-doc", "external", "low", "low", evidence, pos),
        );
      } else if (DOC_URL_RE.test(url)) {
        findings.push(
          mkFinding("ext-doc", "external", "low", "low", evidence, pos),
        );
      }
      // plain URL with no instruction context: no finding
    }
  }
  return findings;
}