import dns from "node:dns/promises";
import net from "node:net";

/**
 * SSRF protection. Every URL the engine fetches — seed, discovered or redirect hop —
 * must pass `assertUrlAllowed()` before any network activity. We validate the scheme,
 * the hostname, and every resolved address. `safeLookup` is used by the HTTP dispatcher
 * so the socket connects only to a validated address (defeats DNS rebinding).
 */

const BLOCKED_HOSTNAMES = new Set([
  "localhost",
  "localhost.localdomain",
  "ip6-localhost",
  "ip6-loopback",
  "metadata.google.internal",
  "metadata",
  "instance-data",
  "kubernetes.default",
  "kubernetes.default.svc",
]);

const BLOCKED_SUFFIXES = [".localhost", ".local", ".internal", ".localdomain", ".home.arpa", ".in-addr.arpa", ".ip6.arpa", ".maclustr.io", ".ts.net"];

function ipv4ToInt(ip: string): number {
  const p = ip.split(".").map((x) => Number(x));
  return ((p[0]! << 24) >>> 0) + (p[1]! << 16) + (p[2]! << 8) + p[3]!;
}

function inCidr4(ip: string, cidr: string): boolean {
  const [base, bitsStr] = cidr.split("/");
  const bits = Number(bitsStr);
  const mask = bits === 0 ? 0 : (~0 << (32 - bits)) >>> 0;
  return ((ipv4ToInt(ip) & mask) >>> 0) === ((ipv4ToInt(base!) & mask) >>> 0);
}

const BLOCKED_V4 = [
  "0.0.0.0/8",
  "10.0.0.0/8",
  "100.64.0.0/10",
  "127.0.0.0/8",
  "169.254.0.0/16",
  "172.16.0.0/12",
  "192.0.0.0/24",
  "192.0.2.0/24",
  "192.168.0.0/16",
  "198.18.0.0/15",
  "198.51.100.0/24",
  "203.0.113.0/24",
  "224.0.0.0/4",
  "240.0.0.0/4",
  "255.255.255.255/32",
];

export function isBlockedIPv4(ip: string): boolean {
  return BLOCKED_V4.some((c) => inCidr4(ip, c));
}

export function isBlockedIPv6(ip: string): boolean {
  const lower = ip.toLowerCase();
  if (lower === "::" || lower === "::1") return true;
  const mapped = lower.match(/^(?:0*:)*ffff:(\d+\.\d+\.\d+\.\d+)$/);
  if (mapped) return isBlockedIPv4(mapped[1]!);
  const mappedHex = lower.match(/^(?:0*:)*ffff:([0-9a-f]{1,4}):([0-9a-f]{1,4})$/);
  if (mappedHex) {
    const a = parseInt(mappedHex[1]!, 16);
    const b = parseInt(mappedHex[2]!, 16);
    return isBlockedIPv4(`${a >> 8}.${a & 255}.${b >> 8}.${b & 255}`);
  }
  if (/^fe[89ab]/.test(lower)) return true; // link-local
  if (lower.startsWith("fc") || lower.startsWith("fd")) return true; // unique local
  if (lower.startsWith("ff")) return true; // multicast
  // NAT64 (RFC 6052 well-known prefix 64:ff9b::/96, and the local-use 64:ff9b:1::/48): the last 32 bits
  // embed an IPv4 address — apply the IPv4 policy to it instead of blocking the whole prefix, so that
  // IPv6-only networks (464XLAT hotspots) can still reach public IPv4 hosts while private IPv4 stays blocked.
  const nat64 = lower.match(/^64:ff9b:(?:1:[0-9a-f]{1,4}:[0-9a-f]{1,4}:)?(?:0*:)*([0-9a-f]{1,4}):([0-9a-f]{1,4})$/);
  if (nat64) {
    const a = parseInt(nat64[1]!, 16);
    const b = parseInt(nat64[2]!, 16);
    return isBlockedIPv4(`${a >> 8}.${a & 255}.${b >> 8}.${b & 255}`);
  }
  if (lower.startsWith("64:ff9b:")) return true; // malformed NAT64 form
  if (lower.startsWith("2001:db8:")) return true; // documentation
  return false;
}

export function isBlockedIP(ip: string): boolean {
  const v = net.isIP(ip);
  if (v === 4) return isBlockedIPv4(ip);
  if (v === 6) return isBlockedIPv6(ip);
  return true;
}

export function isBlockedHostname(hostname: string): boolean {
  const h = hostname.toLowerCase().replace(/\.$/, "");
  if (BLOCKED_HOSTNAMES.has(h)) return true;
  if (BLOCKED_SUFFIXES.some((s) => h.endsWith(s))) return true;
  if (!h.includes(".") && net.isIP(h) === 0) return true; // bare single-label hosts
  return false;
}

export class UrlPolicyError extends Error {
  constructor(
    message: string,
    public readonly reason: string,
  ) {
    super(message);
    this.name = "UrlPolicyError";
  }
}

export interface AllowedUrl {
  url: URL;
  hostname: string;
  addresses: string[];
}

const DNS_FALLBACK_SERVERS = ["1.1.1.1", "8.8.8.8", "9.9.9.9"];

/** Resolve all A/AAAA records, falling back to public resolvers when the system resolver fails. */
export async function resolveAll(hostname: string): Promise<string[]> {
  try {
    const res = await dns.lookup(hostname, { all: true, verbatim: true });
    const addrs = res.map((r) => r.address);
    if (addrs.length) return addrs;
  } catch {
    // fall through
  }
  const r = new dns.Resolver({ timeout: 4000, tries: 2 });
  r.setServers(DNS_FALLBACK_SERVERS);
  const out: string[] = [];
  const [a, aaaa] = await Promise.allSettled([r.resolve4(hostname), r.resolve6(hostname)]);
  if (a.status === "fulfilled") out.push(...a.value);
  if (aaaa.status === "fulfilled") out.push(...aaaa.value);
  return out;
}

export interface UrlPolicyOptions {
  /** Resolve DNS and validate each address. Default true. */
  resolve?: boolean;
  /** Allow http:// (default true). */
  allowHttp?: boolean;
}

/** Validate a URL (scheme, host, resolved addresses). Throws UrlPolicyError. */
export async function assertUrlAllowed(input: string | URL, opts: UrlPolicyOptions = {}): Promise<AllowedUrl> {
  let url: URL;
  try {
    url = typeof input === "string" ? new URL(input) : input;
  } catch {
    throw new UrlPolicyError("Malformed URL.", "malformed");
  }
  if (url.protocol !== "https:" && !(url.protocol === "http:" && (opts.allowHttp ?? true))) {
    throw new UrlPolicyError(`Scheme ${url.protocol} not allowed.`, "scheme");
  }
  if (url.username || url.password) throw new UrlPolicyError("Credentials in URL are not allowed.", "credentials");
  const hostname = url.hostname.replace(/^\[|\]$/g, "");
  if (!hostname) throw new UrlPolicyError("Missing host.", "host");
  if (isBlockedHostname(hostname)) throw new UrlPolicyError(`Host ${hostname} is not allowed.`, "blocked_host");
  if (net.isIP(hostname)) {
    if (isBlockedIP(hostname)) throw new UrlPolicyError(`Address ${hostname} is not allowed.`, "blocked_ip");
    return { url, hostname, addresses: [hostname] };
  }
  if (opts.resolve === false) return { url, hostname, addresses: [] };
  const addresses = await resolveAll(hostname);
  if (!addresses.length) throw new UrlPolicyError(`DNS resolution failed for ${hostname}.`, "dns");
  const bad = addresses.find((a) => isBlockedIP(a));
  if (bad) throw new UrlPolicyError(`Host ${hostname} resolves to a blocked address (${bad}).`, "blocked_ip");
  return { url, hostname, addresses };
}

/**
 * `lookup` implementation for net.connect / undici connect options: only returns
 * addresses that pass the policy, so the TCP connection can never reach a private range
 * even if DNS answers differently than during validation (rebinding).
 */
export function safeLookup(
  hostname: string,
  options: unknown,
  callback: (err: NodeJS.ErrnoException | null, address: string | { address: string; family: number }[], family?: number) => void,
): void {
  const all = typeof options === "object" && options !== null && (options as { all?: boolean }).all === true;
  if (isBlockedHostname(hostname)) {
    callback(Object.assign(new Error(`blocked host ${hostname}`), { code: "EBLOCKED" }), all ? [] : "", 4);
    return;
  }
  resolveAll(hostname)
    .then((addrs) => {
      const ok = addrs.filter((a) => !isBlockedIP(a));
      if (!ok.length) {
        callback(Object.assign(new Error(`no allowed address for ${hostname}`), { code: "EBLOCKED" }), all ? [] : "", 4);
        return;
      }
      if (all) callback(null, ok.map((a) => ({ address: a, family: net.isIP(a) })));
      else callback(null, ok[0]!, net.isIP(ok[0]!));
    })
    .catch((err) => callback(err, all ? [] : "", 4));
}