// src/lib/spec/audit/ssrf-guard.ts — SSRF defense at the fetch boundary.
//
// This is defense-in-depth, independent of the site-URL input validation
// (`siteInputSchema` only enforces `z.url().max(255)`, which happily accepts
// http://127.0.0.1, http://169.254.169.254, etc.). Every URL the auditor
// fetches (home + robots/sitemap/security.txt/llms.txt, and every redirect
// hop) MUST pass `assertPublicUrl` first.
//
// Pure and testable: `isBlockedHostname` is a synchronous predicate; the DNS
// resolver is injected so tests stay hermetic (no real network).

import { lookup } from "node:dns/promises";
import { isIP } from "node:net";

/** Thrown when a URL is disallowed (bad scheme, internal host/IP, no DNS). */
export class BlockedUrlError extends Error {
  constructor(message: string) {
    super(message);
    this.name = "BlockedUrlError";
  }
}

/** Resolve a hostname to the list of IP addresses it points at. */
export type DnsResolver = (hostname: string) => Promise<string[]>;

/** Default resolver backed by the OS resolver (all A/AAAA records). */
export const defaultResolver: DnsResolver = async (hostname) => {
  const results = await lookup(hostname, { all: true });
  return results.map((r) => r.address);
};

/** Lowercase, strip IPv6 brackets and a trailing FQDN dot. */
function normalizeHostname(hostname: string): string {
  let h = hostname.trim().toLowerCase();
  if (h.startsWith("[") && h.endsWith("]")) h = h.slice(1, -1);
  if (h.endsWith(".")) h = h.slice(0, -1);
  return h;
}

/** True for loopback / private / link-local / CGNAT / unspecified IPv4. */
function isBlockedIPv4(ip: string): boolean {
  const parts = ip.split(".");
  if (parts.length !== 4) return false;
  const octets: number[] = [];
  for (const part of parts) {
    // Reject leading-zero / non-numeric tricks; require a plain 0-255 integer.
    if (!/^\d{1,3}$/.test(part)) return false;
    const n = Number(part);
    if (n > 255) return false;
    octets.push(n);
  }
  const [a, b] = octets;
  if (a === 0) return true; // 0.0.0.0/8 (incl. 0.0.0.0)
  if (a === 127) return true; // loopback 127.0.0.0/8
  if (a === 10) return true; // private 10.0.0.0/8
  if (a === 172 && b >= 16 && b <= 31) return true; // private 172.16.0.0/12
  if (a === 192 && b === 168) return true; // private 192.168.0.0/16
  if (a === 169 && b === 254) return true; // link-local 169.254.0.0/16
  if (a === 100 && b >= 64 && b <= 127) return true; // CGNAT 100.64.0.0/10
  return false;
}

/** Parse an IPv6 literal (incl. `::` compression and embedded IPv4) to 8 groups. */
function parseIPv6(ip: string): number[] | null {
  let s = ip;

  // Embedded IPv4 tail (e.g. ::ffff:127.0.0.1) -> convert to two hextets.
  if (s.includes(".")) {
    const idx = s.lastIndexOf(":");
    if (idx === -1) return null;
    const v4 = s.slice(idx + 1).split(".");
    if (v4.length !== 4) return null;
    const octets: number[] = [];
    for (const part of v4) {
      if (!/^\d{1,3}$/.test(part)) return null;
      const n = Number(part);
      if (n > 255) return null;
      octets.push(n);
    }
    const hi = ((octets[0] << 8) | octets[1]).toString(16);
    const lo = ((octets[2] << 8) | octets[3]).toString(16);
    s = `${s.slice(0, idx + 1)}${hi}:${lo}`;
  }

  const halves = s.split("::");
  if (halves.length > 2) return null;
  const head = halves[0] ? halves[0].split(":") : [];
  const tail = halves.length === 2 && halves[1] ? halves[1].split(":") : [];

  let groups: string[];
  if (halves.length === 1) {
    if (head.length !== 8) return null;
    groups = head;
  } else {
    const missing = 8 - head.length - tail.length;
    if (missing < 0) return null;
    groups = [...head, ...Array(missing).fill("0"), ...tail];
  }
  if (groups.length !== 8) return null;

  const out: number[] = [];
  for (const g of groups) {
    if (!/^[0-9a-f]{1,4}$/.test(g)) return null;
    out.push(parseInt(g, 16));
  }
  return out;
}

/** True for loopback / unspecified / ULA / link-local / IPv4-mapped-to-blocked. */
function isBlockedIPv6(ip: string): boolean {
  const g = parseIPv6(ip);
  if (!g) return false;

  // Unspecified :: and loopback ::1
  if (g.every((x) => x === 0)) return true;
  if (g.slice(0, 7).every((x) => x === 0) && g[7] === 1) return true;

  // IPv4-mapped ::ffff:a.b.c.d -> validate the embedded IPv4.
  if (
    g[0] === 0 &&
    g[1] === 0 &&
    g[2] === 0 &&
    g[3] === 0 &&
    g[4] === 0 &&
    g[5] === 0xffff
  ) {
    const a = g[6] >> 8;
    const b = g[6] & 0xff;
    const c = g[7] >> 8;
    const d = g[7] & 0xff;
    return isBlockedIPv4(`${a}.${b}.${c}.${d}`);
  }

  // ULA fc00::/7 -> top 7 bits are 1111110.
  if (((g[0] >> 8) & 0xfe) === 0xfc) return true;
  // Link-local fe80::/10 -> top 10 bits are 1111111010.
  if ((g[0] & 0xffc0) === 0xfe80) return true;

  return false;
}

/**
 * Synchronous predicate: is this hostname (literal IP or name) disallowed?
 * Blocks localhost/*.localhost, all private/loopback/link-local/CGNAT/
 * unspecified IPv4 & IPv6 (incl. IPv4-mapped), and bare single-label
 * (non-FQDN) hostnames — a public site is always fully qualified, so a bare
 * label can only point at something internal.
 */
export function isBlockedHostname(hostname: string): boolean {
  const h = normalizeHostname(hostname);
  if (h === "") return true;
  if (h === "localhost" || h.endsWith(".localhost")) return true;

  const kind = isIP(h);
  if (kind === 4) return isBlockedIPv4(h);
  if (kind === 6) return isBlockedIPv6(h);

  // Not an IP literal: a bare label (no dot) is a suspicious non-FQDN.
  if (!h.includes(".")) return true;
  return false;
}

/**
 * Assert that `url` is safe to fetch. Throws `BlockedUrlError` when:
 *  - the URL is unparseable,
 *  - the scheme is not http/https,
 *  - the host (literal IP or name) is internal per `isBlockedHostname`,
 *  - DNS resolution fails or yields no address,
 *  - ANY resolved address is internal (basic anti-rebinding — the caller must
 *    re-assert on every redirect hop, since resolution can change between
 *    checks).
 */
export async function assertPublicUrl(
  url: string,
  resolver: DnsResolver = defaultResolver,
): Promise<void> {
  let parsed: URL;
  try {
    parsed = new URL(url);
  } catch {
    throw new BlockedUrlError(`Unparseable URL: ${url}`);
  }

  if (parsed.protocol !== "http:" && parsed.protocol !== "https:") {
    throw new BlockedUrlError(`Disallowed scheme: ${parsed.protocol}`);
  }

  const hostname = normalizeHostname(parsed.hostname);
  if (isBlockedHostname(hostname)) {
    throw new BlockedUrlError(`Blocked host: ${hostname}`);
  }

  // A literal IP already passed isBlockedHostname; no DNS needed.
  if (isIP(hostname) !== 0) return;

  let addresses: string[];
  try {
    addresses = await resolver(hostname);
  } catch {
    throw new BlockedUrlError(`DNS resolution failed for host: ${hostname}`);
  }
  if (addresses.length === 0) {
    throw new BlockedUrlError(`No addresses resolved for host: ${hostname}`);
  }
  for (const addr of addresses) {
    if (isBlockedHostname(addr)) {
      throw new BlockedUrlError(
        `Host ${hostname} resolves to a blocked address: ${addr}`,
      );
    }
  }
}