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

/**
 * SSRF guard for user-supplied AI provider base URLs.
 *
 * On a self-hosted box, pointing a provider at a private-LAN service — Ollama
 * at ollama.startos:11434, a LiteLLM/vLLM gateway on 192.168.x, localhost in
 * dev — is a *feature*, so we deliberately ALLOW private and loopback ranges.
 * We block only targets that are never a legitimate provider and are valuable
 * to an attacker: cloud-metadata / link-local (169.254.0.0/16, fe80::/10) and
 * unspecified addresses (0.0.0.0, ::), plus any non-http(s) scheme.
 *
 * We resolve the hostname and check every address it maps to, so a public
 * name that resolves to a blocked range is caught too. (Note: this is a
 * pre-flight check; it does not pin the resolved IP, so a DNS-rebinding race
 * is out of scope — acceptable here since private ranges are allowed anyway.)
 */
export async function assertSafeProviderUrl(rawUrl: string): Promise<void> {
  let parsed: URL;
  try {
    parsed = new URL(rawUrl);
  } catch {
    throw new Error("Invalid URL.");
  }
  if (parsed.protocol !== "http:" && parsed.protocol !== "https:") {
    throw new Error("Only http and https URLs are allowed.");
  }

  // URL.hostname keeps the brackets on IPv6 literals ([::1]); strip BOTH
  // (the /g matters — without it only the leading bracket goes, leaving a
  // trailing ] that fails net.isIP and falls through to a bogus DNS lookup).
  const hostname = parsed.hostname.replace(/^\[|\]$/g, "");

  let addresses: string[];
  if (net.isIP(hostname)) {
    addresses = [hostname];
  } else {
    try {
      const resolved = await lookup(hostname, { all: true });
      addresses = resolved.map((r) => r.address);
    } catch {
      throw new Error(`Could not resolve host: ${hostname}`);
    }
  }

  for (const address of addresses) {
    if (isBlockedAddress(address)) {
      throw new Error(
        `Refusing to connect to ${hostname} (${address}): ` +
          "link-local / cloud-metadata addresses are not allowed.",
      );
    }
  }
}

/**
 * True for addresses that are never a legitimate provider target: the
 * unspecified address and the link-local / cloud-metadata range. Private
 * (RFC1918) and loopback addresses return false on purpose — see the module
 * comment. Exported for unit testing.
 */
export function isBlockedAddress(ip: string): boolean {
  const family = net.isIP(ip);
  if (family === 4) {
    const o = ip.split(".").map(Number);
    if (o[0] === 0) return true; // 0.0.0.0/8 ("this network" — never a valid target)
    if (o[0] === 169 && o[1] === 254) return true; // 169.254.0.0/16 link-local + metadata
    return false;
  }
  if (family === 6) {
    const lower = ip.toLowerCase();
    if (lower === "::") return true; // unspecified
    // fe80::/10 link-local spans fe80..febf
    if (/^fe[89ab]/.test(lower)) return true;
    // IPv4-mapped (::ffff:a.b.c.d) — re-check the embedded v4 address
    const mapped = lower.match(/^::ffff:(\d+\.\d+\.\d+\.\d+)$/);
    if (mapped) return isBlockedAddress(mapped[1]);
    return false;
  }
  return false;
}
