import * as dns from 'dns';
import { isIP } from 'node:net';
import { isPrivateOrForbidden, pinnedFetch } from '../../../net/ssrf-strict.js';

// These delegate to the hardened range check in src/net/ssrf-strict.ts so that
// WebFetch / DownloadFile / BrowseWeb get the same coverage as MCP and SSH:
// loopback, RFC1918, link-local + cloud metadata (169.254/16, fd00:ec2::),
// CGNAT (100.64/10), 0.0.0.0/8, IPv4-mapped IPv6, NAT64, multicast, reserved.
export function isPrivateIPv4(ip: string): boolean {
  if (isIP(ip) !== 4) return false;
  return isPrivateOrForbidden(ip, 4);
}

export function isPrivateIPv6(ip: string): boolean {
  const normalized = ip.toLowerCase().replace(/^\[|\]$/g, '');
  if (isIP(normalized) !== 6) return false;
  return isPrivateOrForbidden(normalized, 6);
}

export function isHostAllowed(hostname: string, allowedHosts: string[]): boolean {
  return allowedHosts.includes(hostname);
}

/**
 * Block requests to private/forbidden destinations.
 *
 * Resolves ALL addresses for the hostname (not just the first) and rejects if
 * any resolves to a private/forbidden range. An explicit allowlist entry
 * bypasses the check (used for trusted internal hosts).
 *
 * Same policy as `resolvePinnedTarget` (which it delegates to), for callers
 * that cannot pin the connection (e.g. Playwright-driven browsing).
 */
export async function checkSSRF(hostname: string, allowedHosts: string[]): Promise<void> {
  await resolvePinnedTarget(hostname, allowedHosts);
}

/**
 * Resolve a hostname, enforce SSRF policy on ALL addresses, and return the IP
 * the connection must be pinned to.
 *
 * Same policy as `checkSSRF` (localhost block, allowlist bypass, every
 * resolved address must pass `isPrivateOrForbidden`) but it also hands back
 * the address to pin so the caller can connect to the exact IP it validated.
 *
 * Returns `null` when the host is on the allowlist — those go through a normal
 * (non-pinned) fetch so trusted internal hosts keep working. For a literal IP
 * the literal itself is the pin (no DNS round-trip).
 */
async function resolvePinnedTarget(
  hostname: string,
  allowedHosts: string[],
): Promise<{ pinnedIp: string; family: 4 | 6 } | null> {
  if (hostname === 'localhost' && !isHostAllowed(hostname, allowedHosts)) {
    throw new Error(`SSRF blocked: hostname "localhost" is not allowed`);
  }
  if (isHostAllowed(hostname, allowedHosts)) {
    // Trusted host: skip pinning, let fetch resolve it normally.
    return null;
  }

  // Literal IP: no DNS query, pin the literal directly.
  const literalFamily = isIP(hostname);
  if (literalFamily === 4 || literalFamily === 6) {
    const fam = literalFamily as 4 | 6;
    if (isPrivateOrForbidden(hostname, fam)) {
      throw new Error(`SSRF blocked: "${hostname}" is a forbidden IP`);
    }
    return { pinnedIp: hostname, family: fam };
  }

  let addrs: Array<{ address: string; family: number }>;
  try {
    addrs = await dns.promises.lookup(hostname, { all: true });
  } catch (e) {
    throw new Error(`DNS resolution failed for "${hostname}": ${(e as Error).message}`);
  }
  if (addrs.length === 0) {
    throw new Error(`SSRF blocked: "${hostname}" resolved to no addresses`);
  }
  for (const a of addrs) {
    if (isPrivateOrForbidden(a.address, a.family as 4 | 6)) {
      throw new Error(`SSRF blocked: "${hostname}" resolves to forbidden IP "${a.address}"`);
    }
  }
  return { pinnedIp: addrs[0].address, family: addrs[0].family as 4 | 6 };
}

/**
 * SSRF-safe fetch that re-validates every redirect hop AND pins the connection
 * to the validated IP (DNS-rebinding defense).
 *
 * `fetch`'s default redirect following re-resolves DNS and would happily
 * follow a 30x to http://169.254.169.254/ (cloud metadata) or an internal
 * host. This follows redirects manually and runs the SSRF policy against each
 * Location before requesting it, so a public URL cannot bounce the request
 * into a private destination.
 *
 * It also closes the TOCTOU rebinding gap: instead of validating the host with
 * one DNS lookup and then letting `fetch` re-resolve (a sub-second rebind can
 * return a public IP to the check and a private/metadata IP to the connection),
 * each hop resolves once, validates every returned address, and connects to the
 * exact validated IP via `pinnedFetch` (custom undici `connect.lookup`, real
 * Host header preserved from the URL). Allowlisted hosts bypass pinning and use
 * a normal fetch.
 */
export async function ssrfSafeFetch(
  url: string,
  allowedHosts: string[],
  init: RequestInit = {},
  maxRedirects = 5,
): Promise<Response> {
  let current = url;
  for (let hop = 0; hop <= maxRedirects; hop++) {
    const parsed = new URL(current);
    const pin = await resolvePinnedTarget(parsed.hostname, allowedHosts);
    const res = pin
      ? await pinnedFetch(current, {
          ...init,
          redirect: 'manual',
          pinnedIp: pin.pinnedIp,
          family: pin.family,
        })
      : await fetch(current, { ...init, redirect: 'manual' });
    const location = res.status >= 300 && res.status < 400 ? res.headers.get('location') : null;
    if (!location) {
      return res;
    }
    // Discard the abandoned redirect hop's body. Without this the hop's
    // connection stays open, and pinnedFetch's per-call Agent (which closes
    // only once its body is consumed or cancelled) would leak per hop.
    if (res.body && !res.bodyUsed) {
      res.body.cancel().catch(() => {});
    }
    // Resolve relative redirects against the current URL.
    current = new URL(location, current).toString();
  }
  throw new Error('SSRF blocked: too many redirects');
}
